使用完毕在此整理一下代码，这里就只对有改动的地方贴一下啊，其他的直接去github上下载一下吧
https://github.com/Guzaiwang/CE-Net
数据输入文件data.py，其实没改动只是不做扩充加载了原始数据，下面会把改动的地方标为斜体,斜体好像显示不清楚，但是我标了地方都会有*号，注意一下吧，都在偏后，直接往下翻就是了。另外数据存放结构很简单如下所示：
标签就在labels里面了，和图像一一对应且同名。
"""
Based on https://github.com/asanakoy/kaggle_carvana_segmentation
"""
import torch
import torch.utils.data as data
from torch.autograd import Variable as V
from PIL import Image
import cv2
import numpy as np
import os
import scipy.misc as misc
def randomHueSaturationValue(image, hue_shift_limit=(-180, 180),
sat_shift_limit=(-255, 255),
val_shift_limit=(-255, 255), u=0.5):
if np.random.random() < u:
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
h, s, v = cv2.split(image)
hue_shift = np.random.randint(hue_shift_limit[0], hue_shift_limit[1]+1)
hue_shift = np.uint8(hue_shift)
h += hue_shift
sat_shift = np.random.uniform(sat_shift_limit[0], sat_shift_limit[1])
s = cv2.add(s, sat_shift)
val_shift = np.random.uniform(val_shift_limit[0], val_shift_limit[1])
v = cv2.add(v, val_shift)
image = cv2.merge((h, s, v))
#image = cv2.merge((s, v))
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
return image
def randomShiftScaleRotate(image, mask,
shift_limit=(-0.0, 0.0),
scale_limit=(-0.0, 0.0),
rotate_limit=(-0.0, 0.0),
aspect_limit=(-0.0, 0.0),
borderMode=cv2.BORDER_CONSTANT, u=0.5):
if np.random.random() < u:
height, width, channel = image.shape
angle = np.random.uniform(rotate_limit[0], rotate_limit[1])
scale = np.random.uniform(1 + scale_limit[0], 1 + scale_limit[1])
aspect = np.random.uniform(1 + aspect_limit[0], 1 + aspect_limit[1])
sx = scale * aspect / (aspect ** 0.5)
sy = scale / (aspect ** 0.5)
dx = round(np.random.uniform(shift_limit[0], shift_limit[1]) * width)
dy = round(np.random.uniform(shift_limit[0], shift_limit[1]) * height)
cc = np.math.cos(angle / 180 * np.math.pi) * sx
ss = np.math.sin(angle / 180 * np.math.pi) * sy
rotate_matrix = np.array([[cc, -ss], [ss, cc]])
box0 = np.array([[0, 0], [width, 0], [width, height], [0, height], ])
box1 = box0 - np.array([width / 2, height / 2])
box1 = np.dot(box1, rotate_matrix.T) + np.array([width / 2 + dx, height / 2 + dy])
box0 = box0.astype(np.float32)
box1 = box1.astype(np.float32)
mat = cv2.getPerspectiveTransform(box0, box1)
image = cv2.warpPerspective(image, mat, (width, height), flags=cv2.INTER_LINEAR, borderMode=borderMode,
borderValue=(
0, 0,
0,))
mask = cv2.warpPerspective(mask, mat, (width, height), flags=cv2.INTER_LINEAR, borderMode=borderMode,
borderValue=(
0, 0,
0,))
return image, mask
def randomHorizontalFlip(image, mask, u=0.5):
if np.random.random() < u:
image = cv2.flip(image, 1)
mask = cv2.flip(mask, 1)
return image, mask
def randomVerticleFlip(image, mask, u=0.5):
if np.random.random() < u:
image = cv2.flip(image, 0)
mask = cv2.flip(mask, 0)
return image, mask
def randomRotate90(image, mask, u=0.5):
if np.random.random() < u:
image=np.rot90(image)
mask=np.rot90(mask)
return image, mask
def default_loader(img_path, mask_path):
img = cv2.imread(img_path)
# print("img:{}".format(np.shape(img)))
img = cv2.resize(img, (448, 448))
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
mask = 255. - cv2.resize(mask, (448, 448))
img = randomHueSaturationValue(img,
hue_shift_limit=(-30, 30),
sat_shift_limit=(-5, 5),
val_shift_limit=(-15, 15))
img, mask = randomShiftScaleRotate(img, mask,
shift_limit=(-0.1, 0.1),
scale_limit=(-0.1, 0.1),
aspect_limit=(-0.1, 0.1),
rotate_limit=(-0, 0))
img, mask = randomHorizontalFlip(img, mask)
img, mask = randomVerticleFlip(img, mask)
img, mask = randomRotate90(img, mask)
mask = np.expand_dims(mask, axis=2)
#
# print(np.shape(img))
# print(np.shape(mask))
img = np.array(img, np.float32).transpose(2,0,1)/255.0 * 3.2 - 1.6
mask = np.array(mask, np.float32).transpose(2,0,1)/255.0
mask[mask >= 0.5] = 1
mask[mask <= 0.5] = 0
#mask = abs(mask-1)
return img, mask
def default_DRIVE_loader(img_path, mask_path):
img = cv2.imread(img_path)
img = cv2.resize(img, (448, 448))
# mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
mask = np.array(Image.open(mask_path))
mask = cv2.resize(mask, (448, 448))
img = randomHueSaturationValue(img,
hue_shift_limit=(-30, 30),
sat_shift_limit=(-5, 5),
val_shift_limit=(-15, 15))
img, mask = randomShiftScaleRotate(img, mask,
shift_limit=(-0.1, 0.1),
scale_limit=(-0.1, 0.1),
aspect_limit=(-0.1, 0.1),
rotate_limit=(-0, 0))
img, mask = randomHorizontalFlip(img, mask)
img, mask = randomVerticleFlip(img, mask)
img, mask = randomRotate90(img, mask)
mask = np.expand_dims(mask, axis=2)
img = np.array(img, np.float32).transpose(2, 0, 1) / 255.0 * 3.2 - 1.6
mask = np.array(mask, np.float32).transpose(2, 0, 1) / 255.0
mask[mask >= 0.5] = 1
mask[mask <= 0.5] = 0
# mask = abs(mask-1)
return img, mask
def read_ORIGA_datasets(root_path, mode='train'):
images = []
masks = []
if mode == 'train':
read_files = os.path.join(root_path, 'Set_A.txt')
else:
read_files = os.path.join(root_path, 'Set_B.txt')
image_root = os.path.join(root_path, 'images')
gt_root = os.path.join(root_path, 'masks')
for image_name in open(read_files):
image_path = os.path.join(image_root, image_name.split('.')[0] + '.jpg')
label_path = os.path.join(gt_root, image_name.split('.')[0] + '.jpg')
print(image_path, label_path)
images.append(image_path)
masks.append(label_path)
return images, masks
def read_Messidor_datasets(root_path, mode='train'):
images = []
masks = []
if mode == 'train':
read_files = os.path.join(root_path, 'train.txt')
else:
read_files = os.path.join(root_path, 'test.txt')
image_root = os.path.join(root_path, 'save_image')
gt_root = os.path.join(root_path, 'save_mask')
for image_name in open(read_files):
image_path = os.path.join(image_root, image_name.split('.')[0] + '.png')
label_path = os.path.join(gt_root, image_name.split('.')[0] + '.png')
images.append(image_path)
masks.append(label_path)
return images, masks
def read_RIM_ONE_datasets(root_path, mode='train'):
images = []
masks = []
if mode == 'train':
read_files = os.path.join(root_path, 'train_files.txt')
else:
read_files = os.path.join(root_path, 'test_files.txt')
image_root = os.path.join(root_path, 'RIM-ONE-images')
gt_root = os.path.join(root_path, 'RIM-ONE-exp1')
for image_name in open(read_files):
image_path = os.path.join(image_root, image_name.split('.')[0] + '.png')
label_path = os.path.join(gt_root, image_name.split('.')[0] + '-exp1.png')
images.append(image_path)
masks.append(label_path)
return images, masks
def read_DRIVE_datasets(root_path, mode='train'):
images = []
masks = []
image_root = os.path.join(root_path, 'training/images')
gt_root = os.path.join(root_path, 'training/1st_manual')
for image_name in os.listdir(image_root):
image_path = os.path.join(image_root, image_name.split('.')[0] + '.tif')
label_path = os.path.join(gt_root, image_name.split('_')[0] + '_manual1.gif')
images.append(image_path)
masks.append(label_path)
print(images, masks)
return images, masks
def read_Cell_datasets(root_path, mode='train'):
images = []
masks = []
image_root = os.path.join(root_path, 'train-images')
gt_root = os.path.join(root_path, 'train-labels')
for image_name in os.listdir(image_root):
image_path = os.path.join(image_root, image_name)
label_path = os.path.join(gt_root, image_name)
images.append(image_path)
masks.append(label_path)
print(images, masks)
return images, masks
def read_datasets_vessel(root_path, mode='train'):
images = []
masks = []
image_root = os.path.join(root_path, 'training/images')
gt_root = os.path.join(root_path, 'training/mask')
for image_name in os.listdir(image_root):
image_path = os.path.join(image_root, image_name)
label_path = os.path.join(gt_root, image_name)
if cv2.imread(image_path) is not None:
if os.path.exists(image_path) and os.path.exists(label_path):
images.append(image_path)
masks.append(label_path)
print(images[:10], masks[:10])
return images, masks
*def read_own_data(root_path, mode = 'train'):
images = []
masks = []
image_root = os.path.join(root_path, 'train/imgs')
gt_root = os.path.join(root_path, 'train/labels')
for image_name in os.listdir(image_root):
image_path = os.path.join(image_root, image_name)
label_path = os.path.join(gt_root, image_name)
images.append(image_path)
masks.append(label_path)
return images, masks*
*def own_data_loader(img_path, mask_path):
img = cv2.imread(img_path)
img = cv2.resize(img, (512, 512))
mask = np.array(Image.open(mask_path))
mask = cv2.resize(mask, (512, 512))
mask = np.expand_dims(mask, axis=2)
img = np.array(img, np.float32).transpose(2, 0, 1)
mask = np.array(mask, np.float32).transpose(2, 0, 1)
return img, mask*
class ImageFolder(data.Dataset):
def __init__(self,root_path, datasets='Messidor', mode='train'):
self.root = root_path
self.mode = mode
self.dataset = datasets
*assert self.dataset in ['RIM-ONE', 'Messidor', 'ORIGA', 'DRIVE', 'Cell', 'Vessel', 'own_data'],* \
"the dataset should be in 'Messidor', 'ORIGA', 'RIM-ONE', 'Vessel', 'own_data'"
if self.dataset == 'RIM-ONE':
self.images, self.labels = read_RIM_ONE_datasets(self.root, self.mode)
elif self.dataset == 'Messidor':
self.images, self.labels = read_Messidor_datasets(self.root, self.mode)
elif self.dataset == 'ORIGA':
self.images, self.labels = read_ORIGA_datasets(self.root, self.mode)
elif self.dataset == 'DRIVE':
self.images, self.labels = read_DRIVE_datasets(self.root, self.mode)
elif self.dataset == 'Cell':
self.images, self.labels = read_Cell_datasets(self.root, self.mode)
elif self.dataset == 'GAN_Vessel':
self.images, self.labels = read_datasets_vessel(self.root, self.mode)
*elif self.dataset == 'own_data':
self.images, self.labels = read_own_data(self.root, self.mode)*
else:
print('Default dataset is Messidor')
self.images, self.labels = read_Messidor_datasets(self.root, self.mode)
def __getitem__(self, index):
# img, mask = default_DRIVE_loader(self.images[index], self.labels[index])
*img, mask = own_data_loader(self.images[index], self.labels[index])*
img = torch.Tensor(img)
mask = torch.Tensor(mask)
return img, mask
def __len__(self):
assert len(self.images) == len(self.labels), 'The number of images must be equal to labels'
return len(self.images)
训练文件main.py，这里面就按照他人要求加了学习率下降策略，可以和源码对照下增加或者保留想要的。
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
import cv2
import os
from tqdm import tqdm
from time import time
import torch
import torch.nn as nn
import torch.utils.data as data
from torch.autograd import Variable as V
from networks.cenet import CE_Net_
from framework import MyFrame
from loss import dice_bce_loss
from data import ImageFolder
from Visualizer import Visualizer
import Constants
import image_utils
# Please specify the ID of graphics cards that you want to use
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
def CE_Net_Train():
NAME = 'road'
solver = MyFrame(CE_Net_, dice_bce_loss, 1e-3)
batchsize = torch.cuda.device_count() * Constants.BATCHSIZE_PER_CARD
scheduler=torch.optim.lr_scheduler.ExponentialLR(solver.optimizer, 0.9)
dataset = ImageFolder(root_path=Constants.ROOT, datasets='own_data')
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=batchsize,
shuffle=True,
num_workers=0)
mylog = open('logs/' + NAME + '.log', 'w')
no_optim = 0
total_epoch = Constants.TOTAL_EPOCH
train_epoch_best_loss = Constants.INITAL_EPOCH_LOSS
for epoch in range(1, total_epoch + 1):
data_loader_iter = iter(data_loader)
train_epoch_loss = 0
index = 0
scheduler.step()
for img, mask in tqdm(data_loader_iter):
solver.set_input(img, mask)
train_loss, pred = solver.optimize()
train_epoch_loss += train_loss
index = index + 1
train_epoch_loss = train_epoch_loss/len(data_loader_iter)
mylog.write('epoch: '+ str(epoch) + ' ' + ' train_loss: ' + str(train_epoch_loss.cpu().numpy()) + '\n')
print('epoch:', epoch, 'train_loss:', train_epoch_loss.cpu().numpy(), 'lr: ' + format(scheduler.get_lr()[0]))
solver.save('./weights/'+str(epoch)+'.th')
mylog.flush()
# print(mylog, 'Finish!')
print('Finish!')
mylog.close()
if __name__ == '__main__':
print(torch.__version__)
CE_Net_Train()
预测文件test.py这个是复制了源码，新建的文件源码也有，可以对照一下，没什么大的变化，其实只是因为没做任何预处理，所以把代码里的相关多余操作去掉了，实际使用的是斜体那段代码。
import torch
import torch.nn as nn
import torch.utils.data as data
from torch.autograd import Variable as V
import sklearn.metrics as metrics
import cv2
import os
import numpy as np
from time import time
from PIL import Image
import warnings
warnings.filterwarnings('ignore')
from networks.cenet import CE_Net_
BATCHSIZE_PER_CARD = 8
class TTAFrame():
def __init__(self, net):
self.net = net().cuda()
self.net = torch.nn.DataParallel(self.net, device_ids=range(torch.cuda.device_count()))
def test_one_img_from_path(self, path, evalmode = True):
if evalmode:
self.net.eval()
batchsize = torch.cuda.device_count() * BATCHSIZE_PER_CARD
if batchsize >= 8:
return self.test_one_img_from_path_1(path)
elif batchsize >= 4:
return self.test_one_img_from_path_2(path)
elif batchsize >= 2:
return self.test_one_img_from_path_4(path)
*def test_one_img_from_path_8(self, path):
img = cv2.imread(path)#.transpose(2,0,1)[None]
img = cv2.resize(img,(512,512))
img90 = np.array(np.rot90(img))
img1 = np.concatenate([img[None],img90[None]])
img2 = np.array(img1)[:,::-1]
img3 = np.array(img1)[:,:,::-1]
img4 = np.array(img2)[:,:,::-1]
img1 = img1.transpose(0,3,1,2)
img2 = img2.transpose(0,3,1,2)
img3 = img3.transpose(0,3,1,2)
img4 = img4.transpose(0,3,1,2)
img1 = V(torch.Tensor(np.array(img1, np.float32)).cuda())
img2 = V(torch.Tensor(np.array(img2, np.float32)).cuda())
img3 = V(torch.Tensor(np.array(img3, np.float32)).cuda())
img4 = V(torch.Tensor(np.array(img4, np.float32)).cuda())
maska = self.net.forward(img1).squeeze().cpu().data.numpy()
maskb = self.net.forward(img2).squeeze().cpu().data.numpy()
maskc = self.net.forward(img3).squeeze().cpu().data.numpy()
maskd = self.net.forward(img4).squeeze().cpu().data.numpy()
mask1 = maska + maskb[:,::-1] + maskc[:,:,::-1] + maskd[:,::-1,::-1]
mask2 = mask1[0] + np.rot90(mask1[1])[::-1,::-1]
return mask2*
def test_one_img_from_path_4(self, path):
img = cv2.imread(path)#.transpose(2,0,1)[None]
img = cv2.resize(img,(512,512))
img90 = np.array(np.rot90(img))
img1 = np.concatenate([img[None],img90[None]])
img2 = np.array(img1)[:,::-1]
img3 = np.array(img1)[:,:,::-1]
img4 = np.array(img2)[:,:,::-1]
img1 = img1.transpose(0,3,1,2)
img2 = img2.transpose(0,3,1,2)
img3 = img3.transpose(0,3,1,2)
img4 = img4.transpose(0,3,1,2)
img1 = V(torch.Tensor(np.array(img1, np.float32)/255.0 * 3.2 -1.6).cuda())
img2 = V(torch.Tensor(np.array(img2, np.float32)/255.0 * 3.2 -1.6).cuda())
img3 = V(torch.Tensor(np.array(img3, np.float32)/255.0 * 3.2 -1.6).cuda())
img4 = V(torch.Tensor(np.array(img4, np.float32)/255.0 * 3.2 -1.6).cuda())
maska = self.net.forward(img1).squeeze().cpu().data.numpy()
maskb = self.net.forward(img2).squeeze().cpu().data.numpy()
maskc = self.net.forward(img3).squeeze().cpu().data.numpy()
maskd = self.net.forward(img4).squeeze().cpu().data.numpy()
mask1 = maska + maskb[:,::-1] + maskc[:,:,::-1] + maskd[:,::-1,::-1]
mask2 = mask1[0] + np.rot90(mask1[1])[::-1,::-1]
return mask2
def test_one_img_from_path_2(self, path):
img = cv2.imread(path)#.transpose(2,0,1)[None]
img = cv2.resize(img,(512,512))
img90 = np.array(np.rot90(img))
img1 = np.concatenate([img[None],img90[None]])
img2 = np.array(img1)[:,::-1]
img3 = np.concatenate([img1,img2])
img4 = np.array(img3)[:,:,::-1]
img5 = img3.transpose(0,3,1,2)
# img5 = np.array(img5, np.float32)/255.0 * 3.2 -1.6
img5 = np.array(img5, np.float32)
img5 = V(torch.Tensor(img5).cuda())
img6 = img4.transpose(0,3,1,2)
# img6 = np.array(img6, np.float32)/255.0 * 3.2 -1.6
img6 = np.array(img6, np.float32)
img6 = V(torch.Tensor(img6).cuda())
maska = self.net.forward(img5).squeeze().cpu().data.numpy()#.squeeze(1)
maskb = self.net.forward(img6).squeeze().cpu().data.numpy()
mask1 = maska + maskb[:,:,::-1]
mask2 = mask1[:2] + mask1[2:,::-1]
mask3 = mask2[0] + np.rot90(mask2[1])[::-1,::-1]
return mask3
def test_one_img_from_path_1(self, path):
img = cv2.imread(path)#.transpose(2,0,1)[None]
img = cv2.resize(img,(512,512))
img90 = np.array(np.rot90(img))
img1 = np.concatenate([img[None],img90[None]])
img2 = np.array(img1)[:,::-1]
img3 = np.concatenate([img1,img2])
img4 = np.array(img3)[:,:,::-1]
img5 = np.concatenate([img3,img4]).transpose(0,3,1,2)
# img5 = np.array(img5, np.float32)/255.0 * 3.2 -1.6
img5 = np.array(img5, np.float32)
img5 = V(torch.Tensor(img5).cuda())
mask = self.net.forward(img5).squeeze().cpu().data.numpy()#.squeeze(1)
mask1 = mask[:4] + mask[4:,:,::-1]
mask2 = mask1[:2] + mask1[2:,::-1]
mask3 = mask2[0] + np.rot90(mask2[1])[::-1,::-1]
return mask3
def load(self, path):
self.net.load_state_dict(torch.load(path))
source = './road/val/imgs/'
val = os.listdir(source)
solver = TTAFrame(CE_Net_)
solver.load('./weights/100.th')
tic = time()
target = './result/'
os.mkdir(target)
for i,name in enumerate(val):
mask = solver.test_one_img_from_path(source+name)
mask[mask>0.5] = 255
mask[mask<=0.5] = 0
mask=cv2.resize(mask,(500,500),interpolation = cv2.INTER_NEAREST)
mask = np.concatenate([mask[:,:,None],mask[:,:,None],mask[:,:,None]],axis=2)
cv2.imwrite(target+name,mask.astype(np.uint8))
精度评定代码eval.py,这是二分类，代码如下：
# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np
class IOUMetric:
"""
Class to calculate mean-iou using fast_hist method
"""
def __init__(self, num_classes):
self.num_classes = num_classes
self.hist = np.zeros((num_classes, num_classes))
def _fast_hist(self, label_pred, label_true):
# 找出标签中需要计算的类别,去掉了背景
mask = (label_true >= 0) & (label_true < self.num_classes)
# # np.bincount计算了从0到n**2-1这n**2个数中每个数出现的次数，返回值形状(n, n)
hist = np.bincount(
self.num_classes * label_true[mask].astype(int) +
label_pred[mask], minlength=self.num_classes ** 2).reshape(self.num_classes, self.num_classes)
return hist
# 输入：预测值和真实值
# 语义分割的任务是为每个像素点分配一个label
def evaluate(self, predictions, gts):
for lp, lt in zip(predictions, gts):
assert len(lp.flatten()) == len(lt.flatten())
self.hist += self._fast_hist(lp.flatten(), lt.flatten())
# miou
iou = np.diag(self.hist) / (self.hist.sum(axis=1) + self.hist.sum(axis=0) - np.diag(self.hist))
miou = np.nanmean(iou)
# -----------------其他指标------------------------------
# mean acc
acc = np.diag(self.hist).sum() / self.hist.sum()
acc_cls = np.nanmean(np.diag(self.hist) / self.hist.sum(axis=1))
freq = self.hist.sum(axis=1) / self.hist.sum()
fwavacc = (freq[freq > 0] * iou[freq > 0]).sum()
return acc, acc_cls, iou, miou, fwavacc
if __name__ == '__main__':
label_path = './road/val/labels/'
predict_path = './result/'
pres = os.listdir(predict_path)
labels = []
predicts = []
for im in pres:
if im[-4:] == '.png':
label_name = im.split('.')[0] + '.png'
lab_path = os.path.join(label_path, label_name)
pre_path = os.path.join(predict_path, im)
label = cv2.imread(lab_path,0)
pre = cv2.imread(pre_path,0)
pre[pre>0] = 1
labels.append(label)
predicts.append(pre)
el = IOUMetric(2) #注意了二分类写2的额
acc, acc_cls, iou, miou, fwavacc = el.evaluate(predicts, labels)
print('acc: ',acc)
print('acc_cls: ',acc_cls)
print('iou: ',iou)
print('miou: ',miou)
print('fwavacc: ',fwavacc)
最后精度评定如下：
('acc: ', 0.9635741133786848)
('acc_cls: ', 0.9319977560822872)
('iou: ', array([0.96197102, 0.53645078]))
('miou: ', 0.7492109018048626)
('fwavacc: ', 0.9419769529495335)
结果并不是很好，你们自己调一下吧
部分结果如下：
1831
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
