viz_singleExpert.py

import random

import cv2
from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
# import some common detectron2 utilities
from detectron2.utils.visualizer import Visualizer

from dataloader.dataset import Dataset

if __name__ == "__main__":
    args = "ImagesQ_hd/"  # change to ImagesQhd/
    # args = "DepthJetQhd/"
    x = Dataset(args)
    cfg = get_cfg()
    model = "COCO-Detection/faster_rcnn_R_50_C4_3x.yaml"
    cfg.merge_from_file(model_zoo.get_config_file(model))
    cfg.DATASETS.TRAIN = ("InOutDoorDepth_train",)
    print("InOutDoorDepth_train")
    cfg.MODEL.RPN.PRE_NMS_TOPK_TRAIN = 3000
    cfg.MODEL.RPN.PRE_NMS_TOPK_TEST = 3000
    cfg.MODEL.RPN.POST_NMS_TOPK_TRAIN = 500
    cfg.MODEL.RPN.POST_NMS_TOPK_TEST = 500
    cfg.DATASETS.TEST = ()
    cfg.OUTPUT_DIR = args
    cfg.DATALOADER.NUM_WORKERS = 2
    cfg.MODEL.WEIGHTS = 'output/model_0024999.pth'
    cfg.SOLVER.IMS_PER_BATCH = 2
    cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
    cfg.SOLVER.MAX_ITER = 10000  # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 50  # faster, and good enough for this toy dataset (default: 512)
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)

    predictor = DefaultPredictor(cfg)

    from detectron2.utils.visualizer import ColorMode
    dataset_dicts = x.load_annotations()
    for d in random.sample(dataset_dicts, 1):
        # im = cv2.imread(d["file_name"])
        im = cv2.imread('docs/imagehd.png')
        im = cv2.resize(im, (1920, 1080))
        # im = cv2.imread('/mnt/AAB281B7B2818911/datasets/InOutDoorPeopleRGBD/DepthJetQhd/seq_1454085791.0435234190.png')
        outputs = predictor(im)

        # format is documented at https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
        v = Visualizer(im[:, :, ::-1],
                       metadata=x.meta,
                       scale=0.5,
                       instance_mode=ColorMode.IMAGE_BW
                       # remove the colors of unsegmented pixels. This option is only available for segmentation models
                       )
        out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
        cv2.imshow('test',out.get_image()[:, :, ::-1])
        cv2.waitKey(0)