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

如何解决保存经过训练的 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

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

解决方法

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

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

然后执行以下操作：

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()[:,cv2.COLOR_RGBA2RGB)
plt.imshow(img)

这将显示单个图像的预测