参考文档:
https://blog.csdn.net/qq_43631789/article/details/102700231
https://blog.csdn.net/qq_39056987/article/details/106455828
https://blog.csdn.net/qq_36766560/article/details/110009622
https://blog.csdn.net/gsgs1234/article/details/115267777
1. 制作自己数据集(labelme)
使用labelme对自己采集的图像进行标注。
安装和使用过程略。
然后标注的json文件默认保存在图像所在的目录内。标注完成之后一般是这样的:
2. 下载deeplab3+源码包(pytorch)
github地址:
https://github.com/jfzhang95/pytorch-deeplab-xception
下载代码包至自己的指定位置:
3. 将数据集转换为VOC格式
3.1 数据结构介绍
我们首先新建一些列文件夹,文件结构如下:
- ImageSets- Segmentation- train.txt-train.txt- val.txt- JPEGImages- SegmentationClass
ImageSets
目录内单放一个Segmentation
文件夹,然后Segmentation
目录下需要制作3个txt文件:train.txt,train.txt,val.txt 用来表示训练集,验证集,测试集的划分信息。制作方式后面介绍。
JPEGImages
目录内用于存放图像数据集的原图。
SegmentationClass
目录放置的mask图像,mask是原图根据标注信息json文件生成的,生成方式后面介绍,注意mask的图像与原图的名称一一对应。
3.2 生成3个txt文件
txt的格式是每一行一个图像文件名,无后缀,不需要地址。
train,trainval,val自己按照一定比例划分
代码如下:
import osimport numpy as nproot = r"D:\dataset\belt\JPEGImages"output = r"D:\dataset\belt\ImageSets\Segmentation"filename = []#从存放原图的目录中遍历所有图像文件# dirs = os.listdir(root)for root, dir, files in os.walk(root): for file in files: print(file) filename.append(file[:-4]) # 去除后缀,存储#打乱文件名列表np.random.shuffle(filename)#划分训练集、测试集,默认比例6:2:2train = filename[:int(len(filename)*0.6)]trainval = filename[int(len(filename)*0.6):int(len(filename)*0.8)]val = filename[int(len(filename)*0.8):]#分别写入train.txt, test.txtwith open(os.path.join(output,'train.txt'), 'w') as f1, open(os.path.join(output,'trainval.txt'), 'w') as f2,open(os.path.join(output,'val.txt'), 'w') as f3: for i in train: f1.write(i + '\n') for i in trainval: f2.write(i + '\n') for i in val: f3.write(i + '\n')print('成功!')
3.3 根据json,制作对应的mask图像
我们首先将所有的json文件存放到单独的文件夹,实例中表示为: "D:\\dataset\\json"
制作mask的图像需要用到labelme的源码。
我们首先找到labelme源码的安装位置:
用anaconda安装的话,windows一般是Users\用户名\.conda\envs\环境名\Lib\site-packages\labelme
。
然后找到labelme\cli
的位置,先备份一下原来的json_to_dataset.py文件,然后用下面的代码覆盖掉原来的json_to_dataset.py。
import argparseimport base64import jsonimport osimport os.path as ospimport PIL.Imageimport yamlfrom labelme.logger import loggerfrom labelme import utilspath = "D:\\dataset\\json"dirs = os.listdir(path)def label(json_file, out_dir, label_name_to_value): # print("json.load(open(json_file))=", json_file) # json_file = os.path.join(path, json_file) # print("json.load(open(json_file))=", json_file) data = json.load(open(json_file)) if data['imageData']: imageData = data['imageData'] else: imagePath = os.path.join(os.path.dirname(json_file), 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) for shape in sorted(data['shapes'], key=lambda x: x['label']): 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 lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value) label_names = [None] * (max(label_name_to_value.values()) + 1) for name, value in label_name_to_value.items(): label_names[value] = name lbl_viz = utils.draw_label(lbl, img, label_names) PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png')) utils.lblsave(osp.join(out_dir, 'label.png'), lbl) PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png')) with open(osp.join(out_dir, 'label_names.txt'), 'w') as f: for lbl_name in label_names: f.write(lbl_name + '\n') logger.warning('info.yaml is being replaced by label_names.txt') info = dict(label_names=label_names) with open(osp.join(out_dir, 'info.yaml'), 'w') as f: yaml.safe_dump(info, f, default_flow_style=False) logger.info('Saved to: {}'.format(out_dir))def main(): logger.warning('This script is aimed to demonstrate how to convert the' 'JSON file to a single image dataset, and not to handle' 'multiple JSON files to generate a real-use dataset.') parser = argparse.ArgumentParser() parser.add_argument('json_file_dir') parser.add_argument('-o', '--out', default=None) args = parser.parse_args() label_name_to_value = {'_background_': 0} for json_file in dirs: # print("json_file=", json_file) if args.out is None: json_file = os.path.join(path, json_file) out_dir = osp.basename(json_file).replace('.', '_') out_dir = osp.join(osp.dirname(json_file), out_dir) else: # out_dir = args.out json_file = os.path.join(path, json_file) out_dir = osp.basename(json_file).replace('.', '_') out_dir = osp.join(osp.dirname(args.out), out_dir) # print('out_dir=',out_dir) if not osp.exists(out_dir): os.mkdir(out_dir) label(json_file, out_dir, label_name_to_value)if __name__ == '__main__': main()
然后在cli
目录输入命令行:
python .\json_to_dataset.py json_file_dir -o D:\dataset\output_mask\
json_file_dir
:必须带的参数-o
:输出目录
然后输出目录就会增加非常多的文件夹:
每个文件夹下有5个文件,而我们需要将所有的label.png文件放入到VOC格式数据集中的SegmentationClass目录,并且需要改成与原图对应的名称。
类别放在label_names.txt里,默认会有一个_background_
。所以我们正常的语义分割至少要有两类。
抽离label.png
的脚本:
import osimport shutilinputdir = 'D:\\dataset\\output_mask'outputdir = 'D:\\dataset\\belt\\SegmentationClass'for dir in os.listdir(inputdir): # 设置旧文件名(就是路径+文件名) oldname = inputdir + os.sep + dir + os.sep + 'label.png' # os.sep添加系统分隔符 print("oldname=",oldname) png_id = '' print("dir=",oldname) # 之前的mask命名都是以xxx_json的目录明明,现在需要把最后的_json部分去除,还原原图的名称 for i in range(len(dir.split('_'))) : if i == len(dir.split('_'))-1: continue else: # 原图的名称中可有可能包含一个或多个_,除了最后一个_,其余都保留。 if png_id != '': png_id += '_' png_id += dir.split('_')[i] # 设置新文件名 newname = outputdir + os.sep + png_id + '.png' shutil.copyfile(oldname, newname) # 用os模块中的rename方法对文件改名 print(oldname, '======>', newname)
执行结束后就可以发现SegmentationClass
目录内放入了mask图像。
4. 修改deeplab+源码,增加自己的数据集
4.1 mypath.py 中加入自己数据集的路径
实例中增加的数据集名称为belt
belt下的文件结构就是之前提到的VOC结构:
-belt- ImageSets- Segmentation- train.txt-train.txt- val.txt- JPEGImages- SegmentationClass
4.2 在dataloaders/datasets目录下添加文件
复制一份pascal.py
文件,并重命名为自己的数据集名称
然后打开自己的数据集py文件,修改文件内的类别数和数据集名称:
4.3 修改dateloaders目录下utils.py
搜素def get_cityscapes_labels()
函数,然后在上方添加自己数据集的函数,例如get_belt_labels()
.
这个函数的主要意思就是给自己每个类设置一个掩膜颜色,有多少个类,就设置多少种颜色。
然后在decode_segmap
函数内添加代码,其中n_classes
是你要分割的类别数
4.4 在dataloaders目录下修改__init__.py
在第一行添加数据集名称,复制’pascal’数据集描述,把名称修改为自己数据集的名字
if args.dataset == 'belt': train_set = belt.VOCSegmentation(args, split='train') val_set = belt.VOCSegmentation(args, split='val') num_class = train_set.NUM_CLASSES train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs) val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, **kwargs) test_loader = None return train_loader, val_loader, test_loader, num_class
4.5 在同级目录中修改train.py约185行添加自己数据集的名称(可以设置为默认)
5. 开始训练数据
输入指令:
python train.py --backbone mobilenet --lr 0.007 --workers 1 --epochs 50 --batch-size 8 --gpu-ids 0 --checkname deeplab-mobilenet
模型保存的路径是在代码内是设置的,在saver.py可以看到保存的路径:run/[datasetname]/[checkname]
,在示例中的路径就是:run/belt/deeplab-mobilenet
保存的目录中可能会存在很多experiment_*
的目录,这是每一此训练都会保存在一个experiment_{}的目录内,最新的训练结果保存在最后id的目录上。此外,最优的模型还会保存到run/[datasetname]/[checkname]
中的model_best.pth.tar中。
若是出现报错:
AttributeError: ‘DeepLab’ object has no attribute ‘module’
解决方式:
打开train.py,找到报错的地方 'state_dict': self.model.module.state_dict()
,修改为
`'state_dict': self.model.state_dict()
6. 测试
源码中没有测试代码,需要自己放入一个测试py文件。
修改–in-path为数据集的测试图片,最后的结果保存在–in-path中
## demo.py#import argparseimport osimport numpy as npimport time from modeling.deeplab import *from dataloaders import custom_transforms as trfrom PIL import Imagefrom torchvision import transformsfrom dataloaders.utils import *from torchvision.utils import make_grid, save_image def main(): parser = argparse.ArgumentParser(description="PyTorch DeeplabV3Plus Training") parser.add_argument('--in-path', type=str, default='/root/home/zyx/Seg552_VOC/test', help='image to test') # parser.add_argument('--out-path', type=str, required=True, help='mask image to save') parser.add_argument('--backbone', type=str, default='resnet', choices=['resnet', 'xception', 'drn', 'mobilenet'], help='backbone name (default: resnet)') parser.add_argument('--ckpt', type=str, default='deeplab-resnet.pth', help='saved model') parser.add_argument('--out-stride', type=int, default=16, help='network output stride (default: 8)') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--gpu-ids', type=str, default='0', help='use which gpu to train, must be a \ comma-separated list of integers only (default=0)') parser.add_argument('--dataset', type=str, default='belt', choices=['pascal', 'coco', 'cityscapes','belt'], help='dataset name (default: pascal)') parser.add_argument('--crop-size', type=int, default=513, help='crop image size') parser.add_argument('--num_classes', type=int, default=2, help='crop image size') parser.add_argument('--sync-bn', type=bool, default=None, help='whether to use sync bn (default: auto)') parser.add_argument('--freeze-bn', type=bool, default=False, help='whether to freeze bn parameters (default: False)') args = parser.parse_args() args.cuda = not args.no_cuda and torch.cuda.is_available() if args.cuda: try: args.gpu_ids = [int(s) for s in args.gpu_ids.split(',')] except ValueError: raise ValueError('Argument --gpu_ids must be a comma-separated list of integers only') if args.sync_bn is None: if args.cuda and len(args.gpu_ids) > 1: args.sync_bn = True else: args.sync_bn = False model_s_time = time.time() model = DeepLab(num_classes=args.num_classes, backbone=args.backbone, output_stride=args.out_stride, sync_bn=args.sync_bn, freeze_bn=args.freeze_bn)model = nn.DataParallel(model) ckpt = torch.load(args.ckpt, map_location='cpu') model.load_state_dict(ckpt['state_dict']) model = model.cuda() model_u_time = time.time() model_load_time = model_u_time-model_s_time print("model load time is {}".format(model_load_time)) composed_transforms = transforms.Compose([ tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), tr.ToTensor()]) for name in os.listdir(args.in_path): s_time = time.time() image = Image.open(args.in_path+"/"+name).convert('RGB') # image = Image.open(args.in_path).convert('RGB') target = Image.open(args.in_path+"/"+name).convert('L') sample = {'image': image, 'label': target} tensor_in = composed_transforms(sample)['image'].unsqueeze(0) model.eval() if args.cuda: tensor_in = tensor_in.cuda() with torch.no_grad(): output = model(tensor_in) grid_image = make_grid(decode_seg_map_sequence(torch.max(output[:3], 1)[1].detach().cpu().numpy()), 3, normalize=False, range=(0, 255)) save_image(grid_image,args.in_path+"/"+"{}_mask.png".format(name[0:-4])) u_time = time.time() img_time = u_time-s_time print("image:{} time: {} ".format(name,img_time)) # save_image(grid_image, args.out_path) # print("type(grid) is: ", type(grid_image)) # print("grid_image.shape is: ", grid_image.shape) print("image save in in_path.")if __name__ == "__main__": main() # python demo.py --in-path your_file --out-path your_dst_file
注意点:
参数--dataset
代码中加入自己的类:
parser.add_argument('--dataset', type=str, default='belt', choices=['pascal', 'coco', 'cityscapes','belt'], help='dataset name (default: pascal)')
输入测试指令:
python testdemo.py --dataset belt --num_classes 2 --ckpt run/Seg552/deeplab-mobilenet/checkpoint.pth.tar --backbone mobilenet
常见报错:
1. state_dict错误
RuntimeError: Error(s) in loading state_dict for DeepLab: Missing key(s) in state_dict: "
在原版中会出现,在ckpt = torch.load(args.ckpt, map_location='cpu')
之前加入 model = nn.DataParallel(model)
即可。
2. CUDA报错
报错如下:raise AssertionError("Torch not compiled with CUDA enabled") AssertionError: Torch not compiled with CUDA enabled
解决方式:
model = model.cuda()