Update visualization.py

visualization.py

@@ -7,8 +7,17 @@ from constants import COLORS
 from utils import fig2img
 
 
+def visualize_mask(mask, img, alpha=0.5):
+    mask = mask.cpu().squeeze().numpy()
+    img = img.cpu().squeeze().permute(1, 2, 0).numpy()
+    plt.imshow(img)
+    plt.imshow(mask, alpha=alpha)
+    plt.axis("off")
+    return fig2img(plt.gcf())
+
+
 def visualize_prediction(
-    pil_img, output_dict, threshold=0.7, id2label=None
+    pil_img, output_dict, threshold=0.7, id2label=None, display_mask=False, mask=None
 ):
     keep = output_dict["scores"] > threshold
     boxes = output_dict["boxes"][keep].tolist()
@@ -19,6 +28,8 @@ def visualize_prediction(
 
     fig, ax = plt.subplots(figsize=(12, 12))
     ax.imshow(pil_img)
+    if display_mask:
+        ax.imshow(mask, alpha=0.5)
     colors = COLORS * 100
     for score, (xmin, ymin, xmax, ymax), label, color in zip(
         scores, boxes, labels, colors