PyTorch项目实战---Unet实现道路裂纹缺陷检测

置顶

是海棠花开

已于 2022-06-07 18:06:41 修改

阅读量6.5k

点赞数 18

于 2022-06-07 11:38:05 首次发布

本文链接：https://blog.csdn.net/m0_55750625/article/details/125159942

版权

Unet 图像分割数据预处理 PyTorch 模型训练

关键词由CSDN通过智能技术生成

# -*- coding: utf-8 -*-
from os.path import isdir
from scipy import io
import os, sys
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

if __name__ == '__main__':
    file_path = './CrackForest-dataset-master/groundTruth/'
    png_img_dir = './CrackForest-dataset-master/groundTruthPngImg/'
    if not isdir(png_img_dir):
        os.makedirs(png_img_dir)
    image_path_lists = os.listdir(file_path)
    images_path = []
    for index in range(len(image_path_lists)):
        image_file = os.path.join(file_path, image_path_lists[index])
        #print(image_file)#./CrackForest-dataset-master/groundTruth/001.mat
        images_path.append(image_file)
        image_mat = io.loadmat(image_file)
        segmentation_image = image_mat['groundTruth']['Segmentation'][0]
        segmentation_image_array = np.array(segmentation_image[0])
        image = Image.fromarray((segmentation_image_array -1) * 255)
        png_image_path = os.path.join(png_img_dir, "%s.png" % image_path_lists[index][0:3])
        #print(png_image_path)#./CrackForest-dataset-master/groundTruthPngImg/001.png
        image.save(png_image_path)
        plt.figure()
        plt.imshow(image)
        #plt.pause(2)
        plt.pause(0.001)
        #plt.show()

groundTruth文件下的label是.mat的文件格式：

groundTruthPngImg文件夹下为转换后的label是.png的图像格式：

3.数据处理

# -*- coding: utf-8 -*-
import os, sys
import numpy as np
import cv2 as cv
import torch
from torch.utils.data import Dataset,DataLoader
import matplotlib.pylab as plt

class SegmentationDataset(object):
    def __init__(self, image_dir, mask_dir):
        self.images = []
        self.masks = []
        files = os.listdir(image_dir)
        sfiles