Search code examples
pytorchobject-detectionopencvfaster-rcnndetectronpython

How to test custom Faster RCNN model(using Detectron 2 and pytorch) on video?


I have trained a Faster RCNN model on a custom dataset for object detection and want to test it on Videos. I could test the results on images but am stuck on how to do that for a video.

Here is the code for inference on images:

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.DATASETS.TEST = ("my_dataset_test", )
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model
predictor = DefaultPredictor(cfg)
test_metadata = MetadataCatalog.get("my_dataset_test")
from detectron2.utils.visualizer import ColorMode
import glob

for imageName in glob.glob('/content/test/*jpg'):
  im = cv2.imread(imageName)
  outputs = predictor(im)
  v = Visualizer(im[:, :, ::-1],
                metadata=test_metadata, 
                scale=0.8
                 )
  out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
  cv2_imshow(out.get_image()[:, :, ::-1])

Please somebody let me know how tweak this code to work for detection on videos?

Platform used: Google Colab

Tech Stack:Detectron 2, Pytorch


Solution

  • Check this loop out:

    from detectron2.utils.visualizer import ColorMode
    import glob
    import cv2
    from google.colab.patches import cv2_imshow
    
    cap = cv2.VideoCapture('/path/to/video')
    while cap.isOpened():
        ret, frame = cap.read()
        # if frame is read correctly ret is True
        if not ret:
            break
    
        outputs = predictor(im)
        v = Visualizer(im[:, :, ::-1],
        metadata=test_metadata, 
        scale=0.8
        )
        out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
        cv2_imshow(out.get_image()[:, :, ::-1])
        
        if cv2.waitKey(1) == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()