Detectron2 利用代码进行模型测试

参考博客：detectron2训练自己的数据集

在上一节章节中，我们已经搭建好了Detectron2的运行环境，我在网上还发现了另一种运行方式。

在Detectron2根目录下新建两个py文件：Detector.py和test.py，代码如下：

test.py：

from Detector import *

detector = Detector(model_type="IS")

detector.onImage('demo/images/1.jpg')

Detector.py：

from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2 import model_zoo
import os

import cv2


class Detector:
    def __init__(self, model_type="OD"):
        self.cfg = get_cfg()

        if model_type == "OD":
            self.cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))
            self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml")
        elif model_type == "IS":
            self.cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
            self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
            # 调用自定义pth模型
            # self.cfg.MODEL.WEIGHTS = os.path.join(self.cfg.OUTPUT_DIR, "model_final.pth")
        self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.4  # set the testing threshold for this model
        # 使用自定义模型时，需要设置该参数
        # self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1

        # cuda、cpu
        self.cfg.MODEL.DEVICE = "cpu"

        self.predictor = DefaultPredictor(self.cfg)

    def onImage(self, imagepath):
        image = cv2.imread(imagepath, flags=1)

        predictions = self.predictor(image)

        viz = Visualizer(image[:, :, ::-1], metadata=MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]),
                         instance_mode=ColorMode.SEGMENTATION)

        output = viz.draw_instance_predictions(predictions["instances"].to("cpu"))

        cv2.imshow("Result", output.get_image()[:, :, ::-1])
        cv2.waitKey(0)

注意：

1、这里有个坑爹的地方，在上一节中我们有下载一个model_final_f10217.pkl文件，放在models目录下，如下：

然后，但是在Detector.py代码中，将self.cfg.MODEL.WEIGHTS的值修改成pkl的文件路径，程序是无法执行的，但是按网上的说明，改成yaml文件后，执行test.py时，坑爹是程序会重新下载个一个pkl文件，并且你还找不到该文件下载到哪个目录下了。。。

2、还有一个注意点，第一个跑test.py文件时，会报下面这个错误：

后来，我按上一节的教程，跑了一次demo/demo.py文件，然后再跑test.py文件，好家伙，跑成功了。

3、我这里还是推荐直接跑官方的demo.py文件比较靠谱，网上这个代码有点坑。。。

4、这个Detector.py代码不是每次都能跑成功，基本上隔三差五的都会报上面这几个错。

运行效果如下：