一、yolact 介绍
yolact:
- 特点:实时实例分割,全卷积,在 Titan Xp 上以 33.5 fps在 MS COCO 上实现了 29.8 mAP,提出了 Fast NMS 比标准NMS快12ms
- 论文:https://arxiv.org/abs/1904.02689(ICCV 2019)
- 代码:https://github.com/dbolya/yolact
yolact++:
- 特点:是yolact的期刊扩展,加入可变形卷积(DCN_v2),在 Titan Xp 上以 33.5 fps在 MS COCO 上实现了 34.1mAP
- 论文:https://arxiv.org/abs/1912.06218(TPAMI 2020)
- 代码:同上
yolact_edge
- 特点:侧重边缘设备部署,在 Jetson AGX Xavier 上以 30.8 FPS(在 RTX 2080 Ti 上为 172.7 FPS)运行 ResNet-101 主干网络和 550x550 分辨率图像,应用 TensorRT 优化,利用了视频的时间冗余
- 论文:https://arxiv.org/abs/2012.12259
- 代码:https://github.com/haotian-liu/yolact_edge
二、自定义数据集
使用labelme标注工具
1、源码:https://github.com/wkentaro/labelme
安装完成后在Terminal输入labelme直接打开gui界面
2、打开图片所在文件夹地址,create polygons画轮廓,点成封闭图形后会跳出label框,设置label,选择label,ok,save,下一张。不同类别设置不同label,同一类别不同个体用同一label,例如路人甲和路人乙的label都是person,有的教程说要person-1,person-2,没必要。save的时候保存在图片同一文件夹下面就好了
3、全部标完后在文件夹下有图片和同名json文件,使用github源码下的labelme2coco.py完成数据转换,具体步骤如下
- 首先创建label.txt,内容如下,头两个必须的,后面是自己的label
__ignore__
_background_
person
label1
label2
- 然后运行如下命令,input_dir为包含图片和json文件的路径,output_dir为保存的coco格式数据集路径,如果想区分train和val数据集,就分两个文件夹转换
./labelme2coco.py input_dir output_dir --labels labels.txt
- 最后在output_dir下会生成一个文件和俩个文件夹,两个文件夹一个是原图一个是label可视化结果。
三、训练yolact
1、在yolact的data/config.py下面添加如下代码。my_custom_dataset设置自定义数据集路径和class_names,可以把数据集放data路径下。yolact_my_config设置使用的数据集和训练方式和网络样式
注意,只有一类的话类名后要加逗号,例如:‘class_names’: (‘person’, )
否则num_classes会算成单词的字母数量
my_custom_dataset = dataset_base.copy({
'name': 'MY',
'train_images': './data/my/my_coco_train/',
'train_info': './data/my/my_coco_train/annotations.json',
'valid_images': './data/my/my_coco_val/',
'valid_info': './data/my/my_coco_val/annotations.json',
'has_gt': True,
'class_names': ('person', 'label1', 'label2')
})
yolact_my_config = coco_base_config.copy({
'name': 'yolact_my',
# Dataset stuff
'dataset': my_custom_dataset,
'num_classes': len(my_custom_dataset.class_names) + 1,
# Image Size
'max_size': 550,
# Training params
'lr_steps': (1000, 1500, 2000, 2500),
'max_iter': 3000,
# Backbone Settings
'backbone': resnet50_backbone.copy({
'selected_layers': list(range(1, 4)),
'pred_scales': [[24], [48], [96], [192], [384]],
'pred_aspect_ratios': [ [[1, 1/2, 2]] ]*5,
'use_pixel_scales': True,
'preapply_sqrt': False,
'use_square_anchors': True, # This is for backward compatability with a bug
}),
# FPN Settings
'fpn': fpn_base.copy({
'use_conv_downsample': True,
'num_downsample': 2,
}),
# Mask Settings
'mask_type': mask_type.lincomb,
'mask_alpha': 6.125,
'mask_proto_src': 0,
'mask_proto_net': [(256, 3, {'padding': 1})] * 3 + [(None, -2, {}), (256, 3, {'padding': 1})] + [(32, 1, {})],
'mask_proto_normalize_emulate_roi_pooling': True,
# Other stuff
'share_prediction_module': True,
'extra_head_net': [(256, 3, {'padding': 1})],
'positive_iou_threshold': 0.5,
'negative_iou_threshold': 0.4,
'crowd_iou_threshold': 0.7,
'use_semantic_segmentation_loss': True,
})
2、训练,–config设置为上面的自定义训练方式名字,等训练完成。
python train.py --config=yolact_my_config
四、转为onnx
1、由于官方的yolact有很多无法直接转为onnx的设置,需要改动模型,可以使用@Ma-Dan改好的代码直接转onnx
代码:https://github.com/Ma-Dan/yolact/tree/onnx
2、操作如下,一个用来转出onnx,一个用来测试生成的onnx
# to generate onnx file.
python eval.py --trained_model=weights/yolact_darknet53_54_800000.pth --score_threshold=0.3 --top_k=100 --cuda=False --image=dog.jpg
# to evaluate with onnx.
python onnxeval.py --trained_model=weights/yolact_resnet50_54_800000.pth --score_threshold=0.3 --top_k=100 --cuda=False --image=dog.jpg
3、主要改动如下:
eval.py下的evalimage
from layers import Detect
import torch.onnx
def evalimage(net:Yolact, path:str, save_path:str=None):
frame = torch.from_numpy(cv2.imread(path)).float()
batch = FastBaseTransform()(frame.unsqueeze(0))
pred_outs = net(batch)
#priors = np.array(pred_outs[3])
#np.savetxt('priors.txt', priors, fmt="%f", delimiter=",")
detect = Detect(cfg.num_classes, bkg_label=0, top_k=200, conf_thresh=0.05, nms_thresh=0.5)
preds = detect({'loc': pred_outs[1], 'conf': pred_outs[3], 'mask':pred_outs[2], 'priors': pred_outs[0], 'proto': pred_outs[4]})
dummy_input = Variable(torch.randn(1, 3, 550, 550))
torch.onnx.export(net, dummy_input, "yolact.onnx", verbose=True, opset_version=10)
img_numpy = prep_display(preds, frame, None, None, undo_transform=False)
if save_path is None:
img_numpy = img_numpy[:, :, (2, 1, 0)]
if save_path is None:
plt.imshow(img_numpy)
plt.title(path)
plt.show()
else:
cv2.imwrite(save_path, img_numpy)
yolact.py下的Yolact的forward的return
#return self.detect(pred_outs)
return pred_outs['priors'], pred_outs['loc'], pred_outs['mask'], pred_outs['conf'], pred_outs['proto']
以上两个改动目的是把网络和detect分离,把输出分离(输出顺序可以自由调整,和detect处理对应即可),加入转onnx的代码
yolact.py下的FPN的forward,固定了sizes大小,对于输入不是550X550的网络需要研究一下这里需要改成什么值
j = len(convouts)
sizes = [(69, 69), (35, 35)]
for lat_layer in self.lat_layers:
j -= 1
if j < len(convouts) - 1:
#_, _, h, w = convouts[j].size()
#x = F.interpolate(x, size=(h, w), mode=self.interpolation_mode, align_corners=False)
x = F.interpolate(x, size=sizes[j], mode=self.interpolation_mode, align_corners=False)
x = x + lat_layer(convouts[j])
out[j] = x
yolact.py下设置use_jit = False
#use_jit = torch.cuda.device_count() <= 1
use_jit = False
onnxeval.py主要改动evalimage
import onnxruntime as rt
def evalimage(net:Yolact, path:str, save_path:str=None):
frame = torch.from_numpy(cv2.imread(path)).float()
batch = FastBaseTransform()(frame.unsqueeze(0))
sess = rt.InferenceSession("yolact.onnx")
input_name = sess.get_inputs()[0].name
loc_name = sess.get_outputs()[0].name
conf_name = sess.get_outputs()[1].name
mask_name = sess.get_outputs()[2].name
priors_name = sess.get_outputs()[3].name
proto_name = sess.get_outputs()[4].name
pred_onx = sess.run([loc_name, conf_name, mask_name, priors_name, proto_name], {input_name: batch.cpu().detach().numpy()})
#priors = np.loadtxt('priors.txt', delimiter=',', dtype='float32')
detect = Detect(cfg.num_classes, bkg_label=0, top_k=200, conf_thresh=0.05, nms_thresh=0.5)
preds = detect({'loc': torch.from_numpy(pred_onx[0]), 'conf': torch.from_numpy(pred_onx[1]), 'mask': torch.from_numpy(pred_onx[2]), 'priors': torch.from_numpy(pred_onx[3]), 'proto': torch.from_numpy(pred_onx[4])})
img_numpy = prep_display(preds, frame, None, None, undo_transform=False)
if save_path is None:
img_numpy = img_numpy[:, :, (2, 1, 0)]
if save_path is None:
plt.imshow(img_numpy)
plt.title(path)
plt.show()
else:
cv2.imwrite(save_path, img_numpy)