matlab训练unet,轻松学Pytorch – 构建UNet实现道路裂纹检测

原标题：轻松学Pytorch – 构建UNet实现道路裂纹检测

微信公众号：OpenCV学堂

关注获取更多计算机视觉与深度学习知识

大家好，我又好久没有给大家更新这个系列了，但是我内心一直没有忘记要更新pytorch初学者系列文章，今天给大家分享一下Pytorch如何构建UNet网络并实现模型训练与测试，实现一个道路裂纹检测!

数据集

CrackForest数据集，包括118张标注数据，37张验证与测试数据。数据集的目录有groundtruth、image、seg三个子目录，分别是标注数据、原始图像、分割信息。其中标注信息是matlab格式的文件，通过字典方式实现数据存储与读写，seg文件本质是text文件，按行来组织信息，前面几行是图像属性与格式化信息，data部分的格式如下：

Seg_num+空格+row_index+空格+column1+column2

空格表示space，

seg_num值为0或者1

row_index表示当前行

column1表示开始列位置

column2 表示结束列位置

假设seg中描述的图像宽度为480，高度为320，表示第一行的分割信息表示如下：

00 0 479 表示图像第一行从列0到列479为0，黑色

1200 141 151 表示图像中第200行中列141到151为1，白色

最终解释上述数据集生成的mask数据显示如下：大小均为(480x320)

Pytorch中定义对应数据集类的代码实现如下：

classSegmentationDataset( Dataset):

def__init__( self, image_dir, mask_dir):

self.images = []

self.masks = []

files = os.listdir(image_dir)

sfiles = os.listdir(mask_dir)

fori inrange(len(sfiles)):

img_file = os.path.join(image_dir, files[i])

mask_file = os.path.join(mask_dir, sfiles[i])

# print(img_file, mask_file)

self.images.append(img_file)

self.masks.append(mask_file)

def__len__( self):

returnlen( self.images)

defnum_of_samples( self):

returnlen( self.images)

def__getitem__( self, idx):

iftorch.is_tensor(idx):

idx = idx.tolist

image_path = self.images[idx]

mask_path = self.masks[idx]

else:

image_path = self.images[idx]

mask_path = self.masks[idx]

img = cv.imread(image_path, cv.IMREAD_GRAYSCALE) # BGR order

mask = cv.imread(mask_path, cv.IMREAD_GRAYSCALE)

# 输入图像

img = np.float32(img) / 255.0

img = np.expand_dims(img, 0)

# 目标标签0 ~ 1， 对于

mask[mask <= 128] = 0

mask[mask > 128] = 1

mask = np.expand_dims(mask, 0)

sample = { 'image': torch.from_numpy(img), 'mask': torch.from_numpy(mask),}

returnsample

模型构建

UNet网络是图像语义分割网络，整个网络可以分为两个部分来解释。第一部分是编码网络，不断的