allow visualizer to work without metadata

Differential Revision: D18021490

fbshipit-source-id: 3d2208914ba9cb8c74b1de48185dbef5909f6977

Author: ppwwyyxx

demo/predictor.py

@@ -21,7 +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.
         """
-        self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
+        self.metadata = MetadataCatalog.get(
+            cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else "__unused"
+        )
         self.cpu_device = torch.device("cpu")
         self.instance_mode = instance_mode
 

detectron2/utils/video_visualizer.py

@@ -87,7 +87,7 @@ def draw_instance_predictions(self, frame, predictions):
         ]
         colors = self._assign_colors(detected)
 
-        labels = _create_text_labels(classes, scores, self.metadata.thing_classes)
+        labels = _create_text_labels(classes, scores, self.metadata.get("thing_classes", None))
 
         if self._instance_mode == ColorMode.IMAGE_BW:
             # any() returns uint8 tensor

detectron2/utils/visualizer.py

@@ -686,20 +686,23 @@ def draw_and_connect_keypoints(self, keypoints):
             output (VisImage): image object with visualizations.
         """
         visible = {}
+        keypoint_names = self.metadata.get("keypoint_names")
         for idx, keypoint in enumerate(keypoints):
             # draw keypoint
             x, y, prob = keypoint
             if prob > _KEYPOINT_THRESHOLD:
                 self.draw_circle((x, y), color=_RED)
-                keypoint_name = self.metadata.keypoint_names[idx]
-                visible[keypoint_name] = (x, y)
-
-        for kp0, kp1, color in self.metadata.keypoint_connection_rules:
-            if kp0 in visible and kp1 in visible:
-                x0, y0 = visible[kp0]
-                x1, y1 = visible[kp1]
-                color = tuple(x / 255.0 for x in color)
-                self.draw_line([x0, x1], [y0, y1], color=color)
+                if keypoint_names:
+                    keypoint_name = keypoint_names[idx]
+                    visible[keypoint_name] = (x, y)
+
+        if self.metadata.get("keypoint_connection_rules"):
+            for kp0, kp1, color in self.metadata.keypoint_connection_rules:
+                if kp0 in visible and kp1 in visible:
+                    x0, y0 = visible[kp0]
+                    x1, y1 = visible[kp1]
+                    color = tuple(x / 255.0 for x in color)
+                    self.draw_line([x0, x1], [y0, y1], color=color)
 
         # draw lines from nose to mid-shoulder and mid-shoulder to mid-hip
         # Note that this strategy is specific to person keypoints.

tests/test_visualizer.py

@@ -102,3 +102,15 @@ def test_overlay_rotated_instances(self):
         v = Visualizer(img, self.metadata)
         output = v.overlay_instances(boxes=rotated_boxes, labels=labels).get_image()
         self.assertEqual(output.shape, img.shape)
+
+    def test_draw_no_metadata(self):
+        img, boxes, _, _, masks = self._random_data()
+        num_inst = len(boxes)
+        inst = Instances((img.shape[0], img.shape[1]))
+        inst.pred_classes = torch.randint(0, 80, size=(num_inst,))
+        inst.scores = torch.rand(num_inst)
+        inst.pred_boxes = torch.from_numpy(boxes)
+        inst.pred_masks = torch.from_numpy(np.asarray(masks))
+
+        v = Visualizer(img, MetadataCatalog.get("asdfasdf"))
+        v.draw_instance_predictions(inst)