deployment - 保存经过训练的 Detectron2 模型并对单个图像进行预测

Question

我是 detectorron2 的新手，这是我的第一个项目。阅读文档并使用教程作为指导后，我在自定义数据集上训练了我的模型并执行了评估。

我现在想通过加载此保存的模型来对通过 API 收到的图像进行预测。我找不到任何可以帮助我完成这项任务的阅读 Material 。

为了保存我的模型，我使用此链接作为引用 - https://detectron2.readthedocs.io/en/latest/tutorials/models.html

我可以使用以下代码保存我训练过的模型-

from detectron2.modeling import build_model
model = build_model(cfg) # returns a torch.nn.Module

from detectron2.checkpoint import DetectionCheckpointer

checkpointer = DetectionCheckpointer(model, save_dir="output")
checkpointer.save("model_final") # save to output/model_final.pth

但我仍然对如何实现我想要的东西感到困惑。我可以使用一些指导来指导我下一步应该做什么。将非常感谢任何可以提供帮助的人。

Answer 1

对于单个图像，创建数据列表。将图像路径放入file_name中，如下所示:

test_data = [{'file_name': '.../image_1jpg',
              'image_id': 10}]

然后运行以下命令:

from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.data import MetadataCatalog
from detectron2.utils.visualizer import Visualizer, ColorMode
import matplotlib.pyplot as plt
import cv2.cv2 as cv2


test_data = [{'file_name': '.../image_1jpg',
                  'image_id': 10}]

cfg = get_cfg()
cfg.merge_from_file("model config")
cfg.MODEL.WEIGHTS = "model_final.pth" # path for final model
predictor = DefaultPredictor(cfg)
im = cv2.imread(test_data[0]["file_name"])
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
               metadata=MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
               scale=0.5,
               instance_mode=ColorMode.IMAGE_BW)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
img = cv2.cvtColor(out.get_image()[:, :, ::-1], cv2.COLOR_RGBA2RGB)
plt.imshow(img)

这将显示单个图像的预测