一个使用detectron2的简单实例分割项目实战

要求：实现对汽车的分割。
配置：ubuntu18.1、python3.6.13、pytorch1.7.1、detectron2 0.4、gtx2080ti。
数据集：从kaggle下载
效果如下（迭代1000次）：

数据集

1、下载后的文件包括下面这些，这里用到train.zip、train_mask.zip，各5088张，其中train_mask对应的是以二值化掩膜的形式标记的汽车区域；

2、将mask标记的文件转为labelme格式，使用mask_to_labelme.py；因为labelme的json文件中包含图片信息，转换后的json文件可以被labelme.exe打开；

3、将labelme转为coco格式数据集，使用labelme_to_json.py；不想自己转换的可以直接下载；
链接：https://pan.baidu.com/s/15895GbAMBJZ58nqQJe8Osg
提取码：1l92

训练

配置文件跟预训练模型都是detectron2自带的，用默认的就行。

import os
from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer, default_argument_parser
from detectron2.data.datasets import register_coco_instances

# 数据集下载：https://www.kaggle.com/c/carvana-image-masking-challenge/data
img_path = r"/home/wang/home/wang/dataset/CarvanaImageMaskingChallenge/train"
json_path = r"/home/wang/home/wang/dataset/CarvanaImageMaskingChallenge/Carvana.json"
config_file = r"/home/wang/home/wang/config/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
model_weights = r"/home/wang/home/wang/ModelWeights/model_final_f10217.pkl"
output_dir = r"/home/wang/home/wang/output/Carvana"

register_coco_instances("Carvana", {}, json_path, img_path)

if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    cfg = get_cfg()
    if len(args.config_file) == 0:
        args.config_file = config_file
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.DATASETS.TRAIN = ("Carvana",)
    cfg.DATASETS.TEST = ()  # no metrics implemented for this dataset
    cfg.DATALOADER.NUM_WORKERS = 8
    cfg.MODEL.WEIGHTS = model_weights
    cfg.SOLVER.IMS_PER_BATCH = 4
    cfg.SOLVER.BASE_LR = 0.0001
    cfg.SOLVER.MAX_ITER = (10000)
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (128)
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
    cfg.OUTPUT_DIR = output_dir
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
    trainer = DefaultTrainer(cfg)
    trainer.resume_or_load(resume=False)
    trainer.train()

测试

from detectron2.engine import DefaultPredictor, default_argument_parser
from detectron2.config import get_cfg
from detectron2.data.detection_utils import read_image
from detectron2.utils.visualizer import ColorMode, Visualizer
import os
import matplotlib.pyplot as plt


config_file = r"/home/wang/home/wang/config/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
model_weights = r"/home/wang/home/wang/output/Carvana/model_final.pth"
image_path = r"/home/wang/home/wang/dataset/CarvanaImageMaskingChallenge/predict/0a0e3fb8f782_01.jpg"


out_filename = os.path.join(os.path.dirname(image_path),
                            os.path.splitext(os.path.basename(image_path))[0] + "_result.jpg")

if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    cfg = get_cfg()
    if len(args.config_file) == 0:
        args.config_file = config_file
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.DATASETS.TEST = ("Carvana",)
    cfg.MODEL.WEIGHTS = model_weights
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
    predictor = DefaultPredictor(cfg)
    image_input = read_image(image_path, format="BGR")
    image_output = predictor(image_input)
    visualizer = Visualizer(image_input, scale=1, instance_mode=ColorMode.IMAGE_BW)
    vis_output = visualizer.draw_instance_predictions(image_output["instances"].to("cpu"))
    np_output = vis_output.get_image()
    plt.subplot(1, 1, 1)
    plt.title('PIL')
    plt.imshow(np_output)
    plt.show()
    vis_output.save(out_filename)  # 保存

参考：UNet实现