时间戳

首页 > 解决方案 > 保存经过训练的 Detectron2 模型并对单个图像进行预测

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

问题描述

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

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

为了保存我的模型，我使用此链接作为参考 - 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

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

标签： deploymentpytorchpredictionfastapidetectron

解决方案

对于单个图像，创建数据列表。将图像路径放入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)

这将显示单个图像的预测

推荐阅读