引言
为了从图片分割出我们想要的特征，博主调研了几种可以部署的语义分割方法，在这主要讲解下如何使用Deeplabv3plus来训练自己的数据集。
一、准备数据集
首先要知道Deeplabv3plus网络需要准备什么样的数据集，数据集总体的格式如下图：
其中，PV是我数据集的名称，下面共包含 三个文件夹。第一个是ImageSets，里面的子文件夹Segmentation存放了训练和验证图片的路径txt；JPEGImage存放了原图，SegmentationClass存放的是分割后的图片，下面我们正式开始准备数据集。
1、首先我们使用labelme工具制作我们的数据集，具体制作方法这里不赘述，最终我们可以得到JPG原图和对应的JSON文件。
2、将JSON转化为图片，具体代码如下，根据自己JSON文件以及label.txt地址来修改：
from __future__ import print_function
import argparse
import glob
import os
import os.path as osp
import sys
import imgviz
import numpy as np
import labelme
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument("--input_dir",default="E:\File\Pycharm\hed-master\hed-data\PV_label_deeplab", help="input annotated directory") #json路径
parser.add_argument("--output_dir",default="E:\File\Pycharm\hed-master\hed-data\zw", help="output dataset directory") #输出地址
parser.add_argument("--labels",default="E:\File\Pycharm\hed-master\hed-data\labels\label.txt", help="labels file") #标签txt
parser.add_argument(
"--noviz", help="no visualization", action="store_true"
)
args = parser.parse_args()
if osp.exists(args.output_dir):
print("Output directory already exists:", args.output_dir)
sys.exit(1)
os.makedirs(args.output_dir)
os.makedirs(osp.join(args.output_dir, "JPEGImages"))
os.makedirs(osp.join(args.output_dir, "SegmentationClass"))
os.makedirs(osp.join(args.output_dir, "SegmentationClassPNG"))
if not args.noviz:
os.makedirs(
osp.join(args.output_dir, "SegmentationClassVisualization")
)
print("Creating dataset:", args.output_dir)
class_names = []
class_name_to_id = {}
for i, line in enumerate(open(args.labels).readlines()):
class_id = i - 1 # starts with -1
class_name = line.strip()
class_name_to_id[class_name] = class_id
if class_id == -1:
assert class_name == "__ignore__"
continue
elif class_id == 0:
assert class_name == "_background_"
class_names.append(class_name)
class_names = tuple(class_names)
print("class_names:", class_names)
out_class_names_file = osp.join(args.output_dir, "class_names.txt")
with open(out_class_names_file, "w") as f:
f.writelines("\n".join(class_names))
print("Saved class_names:", out_class_names_file)
for filename in glob.glob(osp.join(args.input_dir, "*.json")):
print("Generating dataset from:", filename)
label_file = labelme.LabelFile(filename=filename)
base = osp.splitext(osp.basename(filename))[0]
out_img_file = osp.join(args.output_dir, "JPEGImages", base + ".jpg")
out_lbl_file = osp.join(
args.output_dir, "SegmentationClass", base + ".npy"
)
out_png_file = osp.join(
args.output_dir, "SegmentationClassPNG", base + ".png"
)
if not args.noviz:
out_viz_file = osp.join(
args.output_dir,
"SegmentationClassVisualization",
base + ".jpg",
)
with open(out_img_file, "wb") as f:
f.write(label_file.imageData)
img = labelme.utils.img_data_to_arr(label_file.imageData)
lbl, _ = labelme.utils.shapes_to_label(
img_shape=img.shape,
shapes=label_file.shapes,
label_name_to_value=class_name_to_id,
)
labelme.utils.lblsave(out_png_file, lbl)
np.save(out_lbl_file, lbl)
if not args.noviz:
viz = imgviz.label2rgb(
label=lbl,
#img改成image，labelme接口的问题不然会报错
#img=imgviz.rgb2gray(img),
image=imgviz.rgb2gray(img),
font_size=15,
label_names=class_names,
loc="rb",
)
imgviz.io.imsave(out_viz_file, viz)
if __name__ == "__main__":
main()
执行完上面代码，会生成下面的文件：
这样，我们就剩Imagesets没有了。
3、生成train.txt、trainval.txt以及val.txt
准备好上面生成的JPEGImages（原图），执行以下代码：
import os
import numpy as np
root = r"E:\File\Pycharm\pytorch-deeplab-xception-master\dataset\PV\JPEGImages"
output = r"E:\File\Pycharm\pytorch-deeplab-xception-master\dataset\PV\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:2
train = 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.txt
with 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('成功！')
执行完上面代码，会生成三个txt文件。至此，我们所有需要准备的文件都弄好了，我们只需要将所有文件按照一开始给的格式进行存放即可。
二、增加数据集代码
1、下载项目代码：GitHub - jfzhang95/pytorch-deeplab-xception: DeepLab v3+ model in PyTorch. Support different backbones. 下载好之后解压至任意目录下。
2、打开mypath.py，增加自己的数据集，我们的数据集名称是PV，PV下面就上面准备的三个文件夹。
3、在dataloaders/datasets/目录下，复制pascal.py一份，重命名为你数据集的名称，我这里就是PV.py，并且修改PV.py下你数据集的类别数以及数据集名称，红线部分。
4、修改dataloaders/utils.py，如下图，在def get_cityscapes_labels()函数上面加入您数据集的颜色设置，有多少类别就加多少泪类，颜色随意设置即可。
接着在同文件下修改decode_segmap()函数，对刚才增加的方法进行调用，如红色部分：
5、在dataloaders文件下的__init__.py进行修改：
增加自己数据集的读取代码：
if args.dataset == 'PV':
train_set = PV.VOCSegmentation(args, split='train')
val_set = PV.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
到此，代码已经可以读取我们自己的数据集了。
三、训练和测试
1、修改train.py中的训练数据参数，如下图
2、输入训练代码：
python train.py --backbone mobilenet --lr 0.007 --workers 1 --epochs 50 --batch-size 8 --gpu-ids 0 --checkname deeplab-mobilenet
一般来说就可以成功训练了，但是：很多小伙伴遇到什么Assert Error，就什么断言错误，出错地方大部分就是下图的位置（PV.py文件下）：
出错原因：其实很简单，有这个错误基本上是你txt中的图片名称与图片路径中的图片名称对应不上，txt中的图片名称如下图,好好的检查你图片的名称是否对的上，并且检查图片的后缀，因为原图是JPG，分割后的图是PNG。
还有一个原因，就是你txt文件最后你手贱多打了个换行，导致程序认为这也是图片（只不过名字是空的： .jpg），把换行删除即可。
如果还有问题，你就自己print(_image)和print(_cat)，看看打印出来的图片名称是否有问题。
2、测试
原代码并没有测试代码，这里先放一个简单的demo，日后再补充。
#
# demo.py
#
import argparse
import os
import numpy as np
import time
from modeling.deeplab import *
from dataloaders import custom_transforms as tr
from PIL import Image
from torchvision import transforms
from 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
未完待续。。。。。。
有问题直接留言，博主一一解答。