语义分割|【语义分割项目实战】制作语义分割数据集,并使用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

语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

我们最终的目的是要有两个文件夹
【语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测】第一个文件夹imgs:存放原图
第二个文件夹masks:存放生成的掩膜图像
语义分割就是使用这两个图像进行训练
open dir选择存放原图的imgs,我这里是一些五角星的图片
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

然后使用左下角Create Pologons进行标框即可,每标完一张图片,都要保存在另一个新建文件夹jsons下,最终结果如下
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

但是这不是我们最终需要的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 = https://www.it610.com/article/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 = https://www.it610.com/article/data['imageData'] else: imagePath = os.path.join(os.path.dirname(path), data['imagePath']) with open(imagePath, 'rb') as f: imageData = https://www.it610.com/article/f.read() imageData = base64.b64encode(imageData).decode('utf-8') img = utils.img_b64_to_arr(imageData) label_name_to_value = https://www.it610.com/article/{'_background_': 0} for shape in data['shapes']: label_name = shape['label'] if label_name in label_name_to_value: label_value = https://www.it610.com/article/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文件夹所在的绝对路径,要改成自己的
运行完成后结果如下:
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

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文件,将该代码运行即可,运行完结果如下:
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

然后将这些掩膜图像复制粘贴到GitHub上下载项目的data/masks下,原图复制粘贴到imgs文件夹下
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

这样我们的数据集就准备好了
训练模型 打开我们下载的项目文件,有一些参数需要进行修改
在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 masksoptional arguments: -h, --helpshow this help message and exit --epochs E, -e ENumber of epochs --batch-size B, -b BBatch size --learning-rate LR, -l LR Learning rate --load LOAD, -f LOADLoad 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) --ampUse mixed precision

这样代码就开始运行了
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

代码运行结束后,我们每有一个epoch,就会生成一个.pth文件
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

到此,训练部分就结束了
使用训练好的权重进行预测 先找好需要进行测试的图片
然后在predict.py中:
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

将model的路径改成我们想使用的权重文件的绝对路径
然后在终端中输入以下指令:
python predict.py -i text.jpg -o output.jpg #text.jpg是进行测试的原图,output是生成的掩膜结果

运行结束,output.jpg就会生成,可以查看
我的text.jpg原图:
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

生成的output.jpg图像
语义分割|【语义分割项目实战】制作语义分割数据集,并使用U-Net进行实战检测
文章图片

至此复现结束

    推荐阅读