From 046601bc76f154b09e09ba6e33d1e59cf6f3e01e Mon Sep 17 00:00:00 2001 From: Chenxiaon Niu Date: Fri, 16 Oct 2020 16:42:50 +0800 Subject: [PATCH] make sure dataset meta loaded successfully so that vis is with label --- demo/predictor.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/demo/predictor.py b/demo/predictor.py index aa75c9d..e8e5927 100644 --- a/demo/predictor.py +++ b/demo/predictor.py @@ -6,7 +6,7 @@ import cv2 import torch -from detectron2.data import MetadataCatalog +from detectron2.data import MetadataCatalog, DatasetCatalog from detectron2.engine.defaults import DefaultPredictor from detectron2.utils.video_visualizer import VideoVisualizer from detectron2.utils.visualizer import ColorMode, Visualizer @@ -21,6 +21,9 @@ def __init__(self, cfg, instance_mode=ColorMode.IMAGE, parallel=False): parallel (bool): whether to run the model in different processes from visualization. Useful since the visualization logic can be slow. """ + DatasetCatalog.get( + cfg.DATASETS.TRAIN[0] if len(cfg.DATASETS.TRAIN) else "__unused" + ) self.metadata = MetadataCatalog.get( cfg.DATASETS.TRAIN[0] if len(cfg.DATASETS.TRAIN) else "__unused" )