Saving a trained Detectron2 model and making predictions on a single image

Question

I am new to detectron2 and this is my first project. After reading the docs and using the tutorials as a guide, I trained my model on the custom dataset and performed the evaluation.

I would now like to make predictions on images I receive via an API by loading this saved model. I could not find any reading materials that could help me with this task.

To save my model, I have used this link as a reference - https://detectron2.readthedocs.io/en/latest/tutorials/models.html

I am able to save my trained model using the following code-

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

But I am still confused as to how I can go about implementing what I want. I could use some guidance on what my next steps should be. Would be extremely grateful to anyone who can help.

Answer 1

for a single image, create a list of data. Put image path in the file_name as below:

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

Then do run the following:

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)

This will show the prediction for the single image