yolact 训练及推理

一、数据集制作

1.1 标注

我采用的是Labelme工具:https://github.com/wkentaro/labelme
分割需要得到物体的掩膜信息，这里是利用Labelme标注出物体的一圈轮廓，然后再将其转换为掩膜及对应的目标框。

标注完成后一张图片会对应一个json文件，json文件里面以Point记录每一个物体的轮廓点

标注完成后可以借助json_to_dateset.py检查标注是否正确

1.2 转换为COCO数据集格式

COCO数据集格式如下：

• JPEGImages
  存放转换后的jpg图片
• Visualization
  存放转换后实例可视化图片
• annotations.json
  存放所有图片的标注信息

转换脚本可以参考labelme2coco.py
注意：，转换需要定义一个标签文件，如上图标注了桌子和人，那其标签文件第一行放__ignore__，第二行写_background_，然后是自定义的类别：

__ignore__
_background_
desk
person

二、训练及验证

经过前面步骤得到了训练数据

2.1 修改data/config.py

• 添加自己的数据类
  主要记录类别、数据集存放路径，可以参考以下实例，注意： 如果只有一类，那类别那里的逗号不能省略

CUSTOM_CLASSES = ("desk", "person")  # notice: if you have only one class, the value is ("desk", ),  the comma can't leave out
CUSTOM_LABEL_MAP = {1: 1, 2: 2}
my_dataset = Config(
    {
        "name": "my_dataset",
        "train_images": r"D:\coco_train",
        "train_info": r"D:\annotations.json",
        "valid_images": r"D:\coco_val",
        "valid_info": r"D:\annotations.json",
        "class_names": CUSTOM_CLASSES,
        "label_map": CUSTOM_LABEL_MAP
    }
)

• 添加配置
  主要设置模型的骨架网络、迭代次数等信息

yolact_resnet50_my_config = yolact_resnet50_config.copy({
    "name": "my",
    "dataset": my_dataset,
    "num_classes": len(my_dataset.class_names) + 1,
    "max_iter": 20000,                                   # it can calculate the epoch: iter / (data_size / batch)
    "lr_steps": (8000, 15000, 18000),
    "label_map": CUSTOM_LABEL_MAP
})

2.2 训练

设置好配置文件和批处理大小后直接在主目录运行train.py脚本
如：

python train.py 
    --config=yolact_resnet50_my_config 
    --batch_size 4

2.3 验证

得到训练模型后，对给定测试数据进行验证，查看模型实际效果。设置好配置文件、模型路径以及测试路径和结果保存路径以及置信度阈值等，top_K表示一张图片中最多包含的实例个数。

python .\eval.py 
    --config=yolact_resnet50_my_config 
    --trained_model="weights/my_model.pth" 
    --images "E:/data/img/MeshTest***save_path***./results" 
    --top_k 20 
    --score_threshold 0.8

三、模型部署

利用yolact训练了一波自己的数据，目前效果还可以，准备利用C++调用，这里记录一下模型部署过程及中间遇到的一些问题。

3.1. 模型转换为ONNX格式

1、设置配置文件

依据自己训练的配置，设置配置文件

set_cfg("yolact_resnet50_my_config")

2、精简yolact.py

依据配置参数，删除一些判断分支语句，yolact只保留初始化、加载权重和推理。

3、加载权重

weight_path = "../weights/my_model.pth"
net = Yolact()
net.load_weights(weight_path)
net.eval()
net.to(device)

4、转换

inputs = torch.randn(1, 3, 550, 550).to(device)
onnx_model_path = "./yolact.onnx"
print("convert net to ", onnx_model_path, " ... ")
torch.onnx.export(
    net,
    (inputs,),
    onnx_model_path,
    verbose=True,
    input_names=["img"],
    output_names=["loc", "conf", "mask", "proto"],
    opset_version=12
)
print("converted successed!")

3.2、NCNN模型推理

中间部分细节可参考nihui大神的说明
详细记录YOLACT实例分割ncnn实现
C++推理脚本：https://github.com/Tencent/ncnn/blob/master/examples/yolact.cpp

3.3 OpenCV模型推理

和NCNN推理类似，这里记录一下几个关键点，主要是将这几个关键点转换为C++脚本。

• 获取维度信息
  输出P3---P7特征层大小，以及最终输出的信息。

net = Yolact()
print("net: ", net)
net.train()
x = torch.zeros((1, 3, cfg.max_size, cfg.max_size))
y = net(x)

# FPN后输入预测网络的P3---p7特征层
for p in net.prediction_layers:
    print(p.last_conv_size)

# 输出维度
for k, a in y.items():
    print(k + ': ', a.size(), torch.sum(a))

输出结果如下：

(69, 69)
(35, 35)
(18, 18)
(9, 9)
(5, 5)

loc:  torch.Size([1, 19248, 4]) tensor(-5334.4219, grad_fn=<SumBackward0>)
conf:  torch.Size([1, 19248, 81]) tensor(-1648.6711, grad_fn=<SumBackward0>)
mask:  torch.Size([1, 19248, 32]) tensor(-24752.4023, grad_fn=<SumBackward0>)
priors:  torch.Size([19248, 4]) tensor(21850.8438)
proto:  torch.Size([1, 138, 138, 32]) tensor(59765.8398, grad_fn=<SumBackward0>)
segm:  torch.Size([1, 80, 69, 69]) tensor(15391.8555, grad_fn=<SumBackward0>)

• 生成候选框
  这个可以参考yolact.py中的make_priors函数，大体思路还是上一篇提到的依据P3---P7特征图，对其每一个特征点生成不同大小的候选框。

def make_priors(self, conv_h, conv_w, device):
    """ Note that priors are [x,y,width,height] where (x,y) is the center of the box. """

• 掩膜解析
  mask需要乘上对应的系数，然后再缩放回原图大小，可参考layers/output_utils.py中的postprocess函数。

def postprocess(det_output, w, h, batch_idx=0, interpolation_mode='bilinear',
                visualize_lincomb=False, crop_masks=True, score_threshold=0):
    ...
    proto_data = dets['proto']
    masks = proto_data @ masks.t()
    masks = cfg.mask_proto_mask_activation(masks)
    ...
    
    # Permute into the correct output shape [num_dets, proto_h, proto_w]
    masks = masks.permute(2, 0, 1).contiguous()
    
    # Scale masks up to the full image
    masks = F.interpolate(masks.unsqueeze(0), (h, w), mode=interpolation_mode, align_corners=False).squeeze(0)

• 结果解析
  获取到每张图所有实例以及对应的类别、包围框和mask信息。参考yolact.py文件中Detect函数 以及 eval.py中prep_display函数

# yolact.py
# For use in evaluation
self.detect = Detect(cfg.num_classes, bkg_label=0, top_k=cfg.nms_top_k,
                     conf_thresh=cfg.nms_conf_thresh, nms_thresh=cfg.nms_thresh)

# prep_display.py
def prep_display(dets_out, img, h, w, undo_transform=True, class_color=False, mask_alpha=0.45, fps_str=''):

四、常见问题

4.1、FPN不是子模块

• 错误信息
  RuntimeError: Tried to trace <torch.yolact.FPN object at 0x9c26df60> but it is not part of the active trace. Modules that are called during a trace must be registered as submodules of the thing being traced.

• 解决方法
  找到yolact.py文件中FPN类的位置
  将

class FPN(ScriptModuleWrapper):

修改为

class FPN(nn.Module):

4.2、device报错

• 错误信息
  RuntimeError: legacy constructor for device type: cpu was passed device type: cuda, but device type must be: cpu.

• 解决方法
  在yolact开头添加

device = "cuda" if torch.cuda.is_available() else "cpu"

找到yolact.py文件中

self.priors = torch.Tensor(prior_data, device=device).view(-1, 4).detach()

将其更改为

self.priors = torch.Tensor(prior_data).view(-1, 4).detach().to(device)