Skip to content

Instantly share code, notes, and snippets.

@JosephKJ

JosephKJ/visualize.py

Created December 24, 2019 14:15
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save JosephKJ/e523f77af6c5768538eda0c7e2af375a to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save JosephKJ/e523f77af6c5768538eda0c7e2af375a to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. JosephKJ created this gist Dec 24, 2019.
    26 changes: 26 additions & 0 deletions visualize.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,26 @@
    import cv2
    from detectron2.utils.logger import setup_logger
    setup_logger()

    from detectron2.config import get_cfg
    from detectron2.engine import DefaultPredictor
    from detectron2.utils.visualizer import Visualizer
    from detectron2.data import MetadataCatalog

    # Get image
    im = cv2.imread("000002.jpg")

    # Get the configuration ready
    cfg = get_cfg()
    cfg.merge_from_file("configs/PascalVOC-Detection/faster_rcnn_R_50_C4.yaml")
    cfg.MODEL.WEIGHTS = "model_final.pth"
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5

    predictor = DefaultPredictor(cfg)
    outputs = predictor(im)

    v = Visualizer(im[:,:,::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
    v = v.draw_instance_predictions(outputs['instances'].to('cpu'))
    img = v.get_image()[:, :, ::-1]
    cv2.imwrite('output.jpg', img)