binary mask threshold

Summary:
Change area_threshold to 0.
Add some tests.
Pull Request resolved: fairinternal/detectron2#412

Differential Revision: D21813820

Pulled By: ppwwyyxx

fbshipit-source-id: 6f170a6aef0d88a14d438f8ea1a192dffa3560e2

Author: ppwwyyxx

demo/predictor.py

@@ -133,7 +133,7 @@ class AsyncPredictor:
     """
     A predictor that runs the model asynchronously, possibly on >1 GPUs.
     Because rendering the visualization takes considerably amount of time,
-    this helps improve throughput when rendering videos.
+    this helps improve throughput a little bit when rendering videos.
     """
 
     class _StopToken:

detectron2/utils/visualizer.py

@@ -14,6 +14,7 @@
 from matplotlib.backends.backend_agg import FigureCanvasAgg
 from PIL import Image
 
+from detectron2.data import MetadataCatalog
 from detectron2.structures import BitMasks, Boxes, BoxMode, Keypoints, PolygonMasks, RotatedBoxes
 
 from .colormap import random_color
@@ -306,7 +307,7 @@ def get_image(self):
 
 
 class Visualizer:
-    def __init__(self, img_rgb, metadata, scale=1.0, instance_mode=ColorMode.IMAGE):
+    def __init__(self, img_rgb, metadata=None, scale=1.0, instance_mode=ColorMode.IMAGE):
         """
         Args:
             img_rgb: a numpy array of shape (H, W, C), where H and W correspond to
@@ -317,6 +318,8 @@ def __init__(self, img_rgb, metadata, scale=1.0, instance_mode=ColorMode.IMAGE):
             metadata (MetadataCatalog): image metadata.
         """
         self.img = np.asarray(img_rgb).clip(0, 255).astype(np.uint8)
+        if metadata is None:
+            metadata = MetadataCatalog.get("__nonexist__")
         self.metadata = metadata
         self.output = VisImage(self.img, scale=scale)
         self.cpu_device = torch.device("cpu")
@@ -946,7 +949,7 @@ def draw_line(self, x_data, y_data, color, linestyle="-", linewidth=None):
         return self.output
 
     def draw_binary_mask(
-        self, binary_mask, color=None, *, edge_color=None, text=None, alpha=0.5, area_threshold=4096
+        self, binary_mask, color=None, *, edge_color=None, text=None, alpha=0.5, area_threshold=0
     ):
         """
         Args:
@@ -967,8 +970,6 @@ def draw_binary_mask(
         if color is None:
             color = random_color(rgb=True, maximum=1)
         color = mplc.to_rgb(color)
-        if area_threshold is None:
-            area_threshold = 4096
 
         has_valid_segment = False
         binary_mask = binary_mask.astype("uint8")  # opencv needs uint8
@@ -979,7 +980,7 @@ def draw_binary_mask(
             # draw polygons for regular masks
             for segment in mask.polygons:
                 area = mask_util.area(mask_util.frPyObjects([segment], shape2d[0], shape2d[1]))
-                if area < area_threshold:
+                if area < (area_threshold or 0):
                     continue
                 has_valid_segment = True
                 segment = segment.reshape(-1, 2)

tests/test_visualizer.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import unittest
+import cv2
 import torch
 
 from detectron2.data import MetadataCatalog
@@ -141,3 +142,17 @@ def test_draw_no_metadata(self):
 
         v = Visualizer(img, MetadataCatalog.get("asdfasdf"))
         v.draw_instance_predictions(inst)
+
+    def test_draw_binary_mask(self):
+        img, boxes, _, _, masks = self._random_data()
+        img[:, :, 0] = 0  # remove red color
+        mask = masks[0]
+        mask_with_hole = np.zeros_like(mask).astype("uint8")
+        mask_with_hole = cv2.rectangle(mask_with_hole, (10, 10), (50, 50), 1, 5)
+
+        for m in [mask, mask_with_hole]:
+            v = Visualizer(img)
+            o = v.draw_binary_mask(m, color="red", text="test")
+            o = o.get_image().astype("float32")
+            # red color is drawn on the image
+            self.assertTrue(o[:, :, 0].sum() > 0)