SOLO实例分割模型训练过程

一、SOLO介绍

SOLO（Segmenting Objects by Locations）是阿莱德大学和字节跳动联合研究的，算法的核心思想是将分割问题转化为位置分类问题，从而做到不需要 anchor（锚框）及 bounding box，而是根据实例的位置和大小，对每个实例的像素点赋予一个类别从而达到对实例对象进行分割的效果。

代码地址：GitHub - WXinlong/SOLO: SOLO and SOLOv2 for instance segmentation, ECCV 2020 & NeurIPS 2020.

参考博客：

SOLO实战——用自己的数据集训练实例分割模型_from .cocodataset import cocodataset_CV-deeplearning的博客-优快云博客

win10上用实例分割网络SOLO训练自己的数据集（mmdetection版本）_solo实力分割自己的数据集_不打代码也没有头发的博客-优快云博客

SOLO分割模型环境搭建踩坑记录（ubuntu18.04）_solo模型_丹啊丹的博客-优快云博客

二、数据集制作

1）运行labelme2coco.py文件，将所有的json文件读取到instances_train2014.json中。测试集同理。

2）数据集格式如下：annotations中是训练集和测试集的json文件，train2014中是jpg图片和json文件。

三、环境配置

部分包的版本：

torch                  1.10.0+cu113

torchaudio              0.10.0+cu113

torchvision              0.11.1+cu113

mmcv                  0.2.16

mmdet                 1.0.0+unknown /root/anaconda3/envs/torch1.10/lib/python3.8/site-packages/mmdet-1.0.0+unknown-py3.8-linux-x86_64.egg

model-index            0.1.11

networkx               3.1

numpy                 1.23.0

步骤如下：

​
conda create -n solo python=3.7 -y

conda activate solo

conda install -c pytorch pytorch torchvision -y

conda install cython -y

git clone GitHub - WXinlong/SOLO: SOLO and SOLOv2 for instance segmentation, ECCV 2020 & NeurIPS 2020.

cd SOLO

pip install -r requirements/build.txt

pip install "git+GitHub - cocodataset/cocoapi: COCO API - Dataset @ http://cocodataset.org/"

pip install -v -e .

四、注册数据集及修改配置文件

五、训练

训练命令：

python tools/train.py configs/solo/solo_r101_fpn_8gpu_3x.py　＃我用的是101这个

六、模型推理

推理命令：

python demo/inference_demo.py

结果图：

批量图片的测试： 

from mmdet.apis import init_detector, inference_detector, show_result_pyplot, show_result_ins
import mmcv
import os,time

config_file = '/home/ubuntu/wj/SOLO-master/configs/solo/solo_r101_fpn_8gpu_3x.py'
# download the checkpoint from model zoo and put it in `checkpoints/`
checkpoint_file = '/home/ubuntu/wj/SOLO-master/work_dirs/solo_release_r101_fpn_8gpu_3x/epoch_36.pth'

# build the model from a config file and a checkpoint file
model = init_detector(config_file, checkpoint_file, device='cuda:0')

# 批量测试图片
def imgs(imgs_path,results_path):
    t0 = time.time()
    img_list = os.listdir(imgs_path)
    n = len(img_list)
    for img_name in img_list:
        img_path = imgs_path + "/" + img_name
        result = inference_detector(model, img_path)
        # print("********",result)
        new_name = img_name.split(".")[0] + "_1.jpg"
        out_path = results_path + "/" + new_name
        # print("*****", model.CLASSES)
        model.CLASSES = ['big', 'small', 'truck_back', 'other']
        show_result_ins(img_path, result, model.CLASSES, score_thr=0.25, out_file=out_path)
        print(f"%s处理完毕！！！" %  img_name)

    t1 = time.time()
    print("平均用时：", (t1 - t0)/60/n)

imgs('/home/ubuntu/wj/SOLO-master/test_img/imgs','/home/ubuntu/wj/SOLO-master/test_img/results')

七、遇到的问题及解决措施

1.执行训练代码时，'ore is not in the dataset registry'

根据博客在mmdet/datasets/__init__.py修改代码没有修改from .Your_dataset import Your_Dataset，以及在ore.py中的类名。

2.Corrupt JPEG data: 2 extraneous bytes before marker 0xd9

代码读取图片的位置：/home/ubuntu/wj/SOLO-master/mmdet/apis/inference.py，不管它，不影响训练。

3.module 'numpy' has no attribute 'int'

AttributeError: module numpy has no attribute int .报错解决方案_小恶魔饿了的博客-优快云博客

4.ImportError: numpy.core.multiarray failed to import

问题3和问题4是同一个问题，是numpy版本有问题，经过尝试，numpy==1.23.0

5.发现测试结果图上的标签不对应，在inference_demo.py中添加

model.CLASSES = ['big', 'small', 'truck_back', 'other'] show_result_ins(img_path, result, model.CLASSES, score_thr=0.25, out_file=out_path) # 在这一行上面