Detectron2 demo article - gist 3

detectron2_demo_3.py

import cv2
import numpy as np
import re

from detectron2.data import MetadataCatalog
from detectron2.structures import Instances
from detectron2.utils.visualizer import Visualizer, VisImage


if __name__ == "__main__":
    args: argparse.Namespace = _get_parsed_args()

    cfg: CfgNode = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file(args.base_model))
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.4
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(args.base_model)
    predictor: DefaultPredictor = DefaultPredictor(cfg)

    image_file: str
    for image_file in args.images:
        img: np.ndarray = cv2.imread(image_file)

        output: Instances = predictor(img)["instances"]
        v = Visualizer(img[:, :, ::-1],
                       MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
                       scale=1.0)
        result: VisImage = v.draw_instance_predictions(output.to("cpu"))
        result_image: np.ndarray = result.get_image()[:, :, ::-1]

        out_file_name: str = re.search(r"(.*)\.", image_file).group(0)[:-1]
        out_file_name += "_processed.png"

        cv2.imwrite(out_file_name, result_image)