全流程，从零搞懂基于PaddlePaddle的图像分割

import ospath = PATH_TO_JSON
json_file = os.listdir(path)
#os.system("conda activate labelme")for file in json_file:    os.system("labelme_json_to_dataset %s"%(path+"/"+file))

import osimport ioimport numpy as npimport matplotlib.pyplot as pltfrom PIL import Image as PilImageimport paddlefrom paddle.nn import functional as Fimport warnings
warnings.filterwarnings('ignore') 

paddle.__version__

!unzip -q 'data/data75217/doves.zip' -d work/doves

.
├── label.png
├── img.png
├── label_names.txt
└── label_viz.png

images_path = "work/doves"image_count = len([os.path.join(images_path, image_name) 
          for image_name in os.listdir(images_path)])print("用于训练的图片样本数量:", image_count)

def _sort_images(image_dir):
    """
    对文件夹内的图像进行按照文件名排序
    """
    images = []
    labels = []    for image_name in os.listdir(image_dir):        if os.path.isdir(os.path.join(image_dir, image_name)):
            images.append(os.path.join(os.path.join(image_dir, image_name), 'img.png'))
            labels.append(os.path.join(os.path.join(image_dir, image_name), 'label.png'))    return sorted(images), sorted(labels)"""
这里的分割符是\t，后面使用PaddleSeg的时候要注意修改相关代码，因为PaddleSeg读取文件时默认的分割符是空格。
当然也可以将这里的\t替换为空格。
"""def write_file(mode, images, labels):
    with open('./{}.txt'.format(mode), 'w') as f:        for i in range(len(images)):
            f.write('{}\t{}\n'.format(images[i], labels[i]))    
"""
由于所有文件都是散落在文件夹中，在训练时我们需要使用的是数据集和标签对应的数据关系，
所以我们第一步是对原始的数据集进行整理，得到数据集和标签两个数组，分别一一对应。
这样可以在使用的时候能够很方便的找到原始数据和标签的对应关系，否则对于原有的文件夹图片数据无法直接应用。
"""images, labels = _sort_images(images_path)
eval_num = int(image_count * 0.15)"""
由于图片数量有限，这里的测试集和验证集采用相同的一组图片。
"""write_file('train', images[:-eval_num], labels[:-eval_num])
write_file('test', images[-eval_num:], labels[-eval_num:])
write_file('eval', images[-eval_num:], labels[-eval_num:])

with open('./train.txt', 'r') as f:
    i = 0

    for line in f.readlines():
        image_path, label_path = line.strip().split('\t')
        image = np.array(PilImage.open(image_path))
        label = np.array(PilImage.open(label_path))    
        if i > 2:            break
        # 进行图片的展示
        plt.figure()

        plt.subplot(1,2,1), 
        plt.title('Train Image')
        plt.imshow(image.astype('uint8'))
        plt.axis('off')

        plt.subplot(1,2,2), 
        plt.title('Label')
        plt.imshow(label.astype('uint8'), cmap='gray')
        plt.axis('off')

        plt.show()
        i = i + 1

class MyDataset(Dataset):
    def __init__(self):
        ...        
    # 每次迭代时返回数据和对应的标签
    def __getitem__(self, idx):
        return x, y    # 返回整个数据集的总数
    def __len__(self):
        return count(samples)

import paddlefrom code.unet import DoveNet

num_classes = 2IMAGE_SIZE = (224, 224)
network = DoveNet(num_classes)
model = paddle.Model(network)
model.summary((-1, 3,) + IMAGE_SIZE)

!python code/train.py

import paddlefrom code.unet import DoveNetfrom code.dove_dataset import DoveDataset

num_classes = 2network = DoveNet(num_classes)

state_dict = paddle.load('unet.pdparams')
network.set_state_dict(state_dict)
model = paddle.Model(network)

predict_dataset = DoveDataset(mode='test')
test_loader = paddle.io.DataLoader(predict_dataset, places=paddle.CUDAPlace(0), batch_size= 32)
model.prepare(paddle.nn.CrossEntropyLoss(axis=1))
predict_results = model.predict(test_loader)

import numpy as npimport matplotlib.pyplot as pltfrom paddle.vision.transforms import transforms as Tfrom PIL import Image as PilImage

plt.figure(figsize=(10, 10))

IMAGE_SIZE = (224, 224)
i = 0idx = 0with open('./test.txt', 'r') as f:    for line in f.readlines():
        image_path, label_path = line.strip().split('\t')
        resize_t = T.Compose([
            T.Resize(IMAGE_SIZE)
        ])
        image = resize_t(PilImage.open(image_path))
        label = resize_t(PilImage.open(label_path))

        image = np.array(image).astype('uint8')
        label = np.array(label).astype('uint8')        if i > 8: 
            break
        plt.subplot(3, 3, i + 1)
        plt.imshow(image)
        plt.title('Input Image')
        plt.axis("off")

        plt.subplot(3, 3, i + 2)
        plt.imshow(label, cmap='gray')
        plt.title('Label')
        plt.axis("off")
        
        data = predict_results[0][0][idx].transpose((1, 2, 0))
        mask = np.argmax(data, axis=-1)

        plt.subplot(3, 3, i + 3)
        plt.imshow(mask.astype('uint8'), cmap='gray')
        plt.title('Predict')
        plt.axis("off")
        i += 3
        idx += 1plt.show()

!pip install paddleseg

!git clone https://github.com/PaddlePaddle/PaddleSeg.git

!unzip -q data/data75217/PaddleSeg.zip -d work

## 使用U-net进行训练!python work/PaddleSeg/train.py --config unet.yml --save_interval 2000

## 将测试图片拷贝到新建的文件夹中import osimport shutilif not os.path.exists('test_imgs'):
    os.mkdir('test_imgs')with open('test.txt', 'r') as f:
    paths = f.readlines()
    i = 0
    for path in paths:
        img, label = path.strip().split('\t')
        shutil.copy(img, 'test_imgs/t'+str(i)+'.png')
        i += 1

## 将标签图片拷贝到新建的文件夹中if not os.path.exists('label_imgs'):
    os.mkdir('label_imgs')with open('seg_test.txt', 'r') as f:
    paths = f.readlines()
    i = 0
    for path in paths:
        img, label = path.strip().split('\t')
        shutil.copy(label, 'label_imgs/t'+str(i)+'.png')
        i += 1

## 开始预测并保存预测后的图片!python work/PaddleSeg/predict.py --image_path test_imgs \
--model_path output/iter_10000/model.pdparams \
--save_dir saved_imges \
--crop_size 512 512 \
--config unet.yml

## 创建需要可视化的图片列表test_images = os.listdir('test_imgs')
label_images = os.listdir('label_imgs')
predicted_images = os.listdir('saved_imges/pseudo_color_prediction')

test_images = ['test_imgs/' + path for path in test_images if not path.startswith('.')]
label_images = ['label_imgs/' + path for path in label_images if not path.startswith('.')]
predicted_images = ['saved_imges/pseudo_color_prediction/' + path for path in predicted_images if not path.startswith('.')]

test_images.sort()
label_images.sort()
predicted_images.sort()

## 开始可视化from paddle.vision.transforms import transforms as Timport matplotlib.pyplot as pltfrom PIL import Image as PilImageimport numpy as np

plt.figure(figsize=(10, 10))

IMAGE_SIZE = (224, 224)
i = 0for j in range(len(test_images)):
    test_path = test_images[j]
    label_path = label_images[j]
    predicted_path = predicted_images[j]

    resize_t = T.Compose([
        T.Resize(IMAGE_SIZE)
    ])
    image = resize_t(PilImage.open(test_path))
    label = resize_t(PilImage.open(label_path))
    predicted = resize_t(PilImage.open(predicted_path))

    image = np.array(image).astype('uint8')
    label = np.array(label).astype('uint8')
    predicted = np.array(predicted).astype('uint8')    if i > 8: 
        break
    plt.subplot(3, 3, i + 1)
    plt.imshow(image)
    plt.title('Input Image')
    plt.axis("off")

    plt.subplot(3, 3, i + 2)
    plt.imshow(label, cmap='gray')
    plt.title('Label')
    plt.axis("off")

    plt.subplot(3, 3, i + 3)
    plt.imshow(predicted, cmap='gray')
    plt.title('Predict')
    plt.axis("off")
    i += 3
    j += 1plt.show()