【语义分割项目实战】制作语义分割数据集，并使用U-Net进行实战检测

首先上项目的链接：

GitHub - milesial/Pytorch-UNet: PyTorch implementation of the U-Net for image semantic segmentation with high quality images

这个好像是官方给的项目代码，用起来非常的方便，下载完直接解压压缩包，然后使用pycharm打开即可

Labelme工具安装

首先在anaconda环境下创建一个新的环境

conda create -n labelme python=3.6

安装完成后激活该环境

conda activate labelme

安装

pip install labelme

安装完后直接打开即可

labelme

我们最终的目的是要有两个文件夹

第一个文件夹imgs：存放原图

第二个文件夹masks：存放生成的掩膜图像

语义分割就是使用这两个图像进行训练 

open dir选择存放原图的imgs，我这里是一些五角星的图片

然后使用左下角Create Pologons进行标框即可，每标完一张图片，都要保存在另一个新建文件夹jsons下，最终结果如下

 

 但是这不是我们最终需要的mask图像，还要使用别的脚本进行生成

我这里借鉴了博主 奶盖芒果的脚本，看了好多贴子只有这个姐姐讲的最好QAQ

原帖链接

(22条消息) U-net复现pytorch版本 以及制作自己的数据集并训练_奶盖芒果的博客-CSDN博客_unet训练

labelme批量实现json_to_dataset方法（2021）_奶盖芒果的博客-CSDN博客_json_to_dataset.py

首先按照文中的方法，找到json_to_dataset的文件，将其改写为

import argparse
import json
import os
import os.path as osp
import warnings
 
import PIL.Image
import yaml
 
from labelme import utils
import base64
 
import numpy as np
from skimage import img_as_ubyte
 
def main():
    warnings.warn("This script is aimed to demonstrate how to convert the\n"
                  "JSON file to a single image dataset, and not to handle\n"
                  "multiple JSON files to generate a real-use dataset.")
    parser = argparse.ArgumentParser()
    parser.add_argument('json_file')
    parser.add_argument('-o', '--out', default=None)
    args = parser.parse_args()
 
    json_file = args.json_file
 
    count = os.listdir(json_file) 
    for i in range(0, len(count)):
        path = os.path.join(json_file, count[i])
        if os.path.isfile(path):
            data = json.load(open(path))
 
			##############################
			#save_diretory
            out_dir1 = osp.basename(path).replace('.', '_')
            save_file_name = out_dir1
            out_dir1 = osp.join(osp.dirname(path), out_dir1)
 
            if not osp.exists(json_file + '\\' + 'labelme_json'):
                os.mkdir(json_file + '\\' + 'labelme_json')
            labelme_json = json_file + '\\' + 'labelme_json'
 
            out_dir2 = labelme_json + '\\' + save_file_name
            if not osp.exists(out_dir2):
                os.mkdir(out_dir2)
 
			#########################
 
            if data['imageData']:
                imageData = data['imageData']
            else:
                imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
                with open(imagePath, 'rb') as f:
                    imageData = f.read()
                    imageData = base64.b64encode(imageData).decode('utf-8')
            img = utils.img_b64_to_arr(imageData)
            label_name_to_value = {'_background_': 0}
            for shape in data['shapes']:
                label_name = shape['label']
                if label_name in label_name_to_value:
                    label_value = label_name_to_value[label_name]
                else:
                    label_value = len(label_name_to_value)
                    label_name_to_value[label_name] = label_value
            
            # label_values must be dense
            label_values, label_names = [], []
            for ln, lv in sorted(label_name_to_value.items(), key=lambda x: x[1]):
                label_values.append(lv)
                label_names.append(ln)
            assert label_values == list(range(len(label_values)))
            
            lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
            
			
			
            captions = ['{}: {}'.format(lv, ln)
                for ln, lv in label_name_to_value.items()]
            lbl_viz = utils.draw_label(lbl, img, captions)
             
            PIL.Image.fromarray(img).save(out_dir2+'\\'+save_file_name+'_img.png')
            #PIL.Image.fromarray(lbl).save(osp.join(out_dir2, 'label.png'))
            utils.lblsave(osp.join(out_dir2, save_file_name+'_label.png'), lbl)
            PIL.Image.fromarray(lbl_viz).save(out_dir2+'\\'+save_file_name+
            '_label_viz.png')
 
            with open(osp.join(out_dir2, 'label_names.txt'), 'w') as f:
                for lbl_name in label_names:
                    f.write(lbl_name + '\n')
 
            warnings.warn('info.yaml is being replaced by label_names.txt')
            info = dict(label_names=label_names)
            with open(osp.join(out_dir2, 'info.yaml'), 'w') as f:
                yaml.safe_dump(info, f, default_flow_style=False)
				
				
			#save png to another directory
            if not osp.exists(json_file + '\\' + 'mask_png'):
                os.mkdir(json_file + '\\' + 'mask_png')
            mask_save2png_path = json_file + '\\' + 'mask_png'
 
            utils.lblsave(osp.join(mask_save2png_path, save_file_name+'_label.png'), lbl)
 
            print('Saved to: %s' % out_dir2)
            
if __name__ == '__main__':
    main()

然后将这个python文件复制到和imgs，jsons同一个文件夹下，在终端中输入以下指令

python json_to_dataset.py H:\pictures\jsons

最后那个路径是jsons文件夹所在的绝对路径，要改成自己的

运行完成后结果如下：

mask_png中存放的就是生成的掩膜图像

后面还需要改写一下这些图像的命名，因为后面我们为了便于图像的读取，需要将原图和掩膜图像改成一样且一一对应的，例如，原图的命名为1.jpg，那么它对应的掩膜图像应该是1.jpg或者1.png

使用的改名脚本如下，我在这里对原博主的代码进行了一点点小的修改

 

import os
import sys


# 修改生成的文件名字
def rename():
    path = input("请输入路径(例如D:\\\\picture)：")
    startNumber = input("请输入开始数:")
    fileType = input("请输入后缀名（如 .jpg、.txt等等）:")
    print("正在生成以" + startNumber  + fileType + "迭代的文件名")
    count = 0
    filelist = os.listdir(path)
    for files in filelist:
        Olddir = os.path.join(path, files)
        if os.path.isdir(Olddir):
            continue
        Newdir = os.path.join(path, str(count + int(startNumber))+ fileType)
        os.rename(Olddir, Newdir)
        count += 1
    print("一共修改了" + str(count) + "个文件")


rename()

然后在该项目下新建一个python文件，将该代码运行即可，运行完结果如下：

然后将这些掩膜图像复制粘贴到GitHub上下载项目的data/masks下，原图复制粘贴到imgs文件夹下

这样我们的数据集就准备好了

 训练模型

打开我们下载的项目文件，有一些参数需要进行修改

在utils/data_loading.py的最后一行中，将mask_suffix后面的'_mask'删掉，因为我们已经把mask的命名和imgs的命名改成一样的了，就不需要再加这一串字符

class CarvanaDataset(BasicDataset):
    def __init__(self, images_dir, masks_dir, scale=1):
        super().__init__(images_dir, masks_dir, scale, mask_suffix='')#在此处修改

然后就可以在终端中运行

python train.py

对于epoch，batch-size等变量，可以在train的源代码中修改，也可以在终端中直接修改

具体修改可以参考GitHub原作者的注释

> python train.py -h
usage: train.py [-h] [--epochs E] [--batch-size B] [--learning-rate LR]
                [--load LOAD] [--scale SCALE] [--validation VAL] [--amp]

Train the UNet on images and target masks

optional arguments:
  -h, --help            show this help message and exit
  --epochs E, -e E      Number of epochs
  --batch-size B, -b B  Batch size
  --learning-rate LR, -l LR
                        Learning rate
  --load LOAD, -f LOAD  Load model from a .pth file
  --scale SCALE, -s SCALE
                        Downscaling factor of the images
  --validation VAL, -v VAL
                        Percent of the data that is used as validation (0-100)
  --amp                 Use mixed precision

这样代码就开始运行了

 代码运行结束后，我们每有一个epoch，就会生成一个.pth文件

到此，训练部分就结束了

使用训练好的权重进行预测

 先找好需要进行测试的图片

然后在predict.py中：

将model的路径改成我们想使用的权重文件的绝对路径

然后在终端中输入以下指令：

python predict.py -i text.jpg -o output.jpg
#text.jpg是进行测试的原图，output是生成的掩膜结果

 运行结束，output.jpg就会生成，可以查看

我的text.jpg原图：

生成的output.jpg图像

 

至此复现结束