Home > Software engineering >  Detectron2 prediction problem on my local machine with custom model
Detectron2 prediction problem on my local machine with custom model

Time:03-08

I trained a custom model with detectron2 on google colab, and ok, it's working correctly. The model was trained, the predictions were ok, this on google colab. But when I made predictions on my local machine did'nt work. Here a similar example on google colab: https://colab.research.google.com/drive/1bSlH5Am_zFEWbJ9zTRu2wFEDKDvn0LUv?usp=sharing

I exported de final model and ran with this code:

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

cfg = get_cfg()
cfg.merge_from_file("./detectron2_repo/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = "model_final.pth" # path for final model
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8
predictor = DefaultPredictor(cfg)
im = cv2.imread('0.jpg')
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)
cv2.imwrite('img.jpg',img)

I supose that the cfg.merge_from_file is the problem. Is there other file? Where I find on colab?

I tested the standard models and worked well on my local machine, the problem is with the custom model.

CodePudding user response:

I saved the configs with this comand and then I downloaded.

f = open('config.yml','w')
f.write(cfg.dump())
f.close()

and replaced:

cfg.merge_from_file("./detectron2_repo/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")

by

cfg.merge_from_file("config.yml")

and worked.

  • Related