Update app.py

app.py

@@ -22,8 +22,54 @@ def fig2img(fig):
     return img
 
 def plot(annotations, prompt_process, mask_random_color=True, better_quality=True, retina=True, with_contours=True):
-    # ... (keep the existing plot function as is)
-    # This function doesn't need modification for our purposes
+    for ann in annotations:
+        image = ann.orig_img[..., ::-1]  # BGR to RGB
+        original_h, original_w = ann.orig_shape
+        fig = plt.figure(figsize=(original_w / 100, original_h / 100))
+        plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
+        plt.margins(0, 0)
+        plt.gca().xaxis.set_major_locator(plt.NullLocator())
+        plt.gca().yaxis.set_major_locator(plt.NullLocator())
+        plt.imshow(image)
+
+        if ann.masks is not None:
+            masks = ann.masks.data
+            if better_quality:
+                if isinstance(masks[0], torch.Tensor):
+                    masks = np.array(masks.cpu())
+                for i, mask in enumerate(masks):
+                    mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
+                    masks[i] = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8))
+
+            prompt_process.fast_show_mask(
+                masks,
+                plt.gca(),
+                random_color=mask_random_color,
+                bbox=None,
+                points=None,
+                pointlabel=None,
+                retinamask=retina,
+                target_height=original_h,
+                target_width=original_w,
+            )
+
+            if with_contours:
+                contour_all = []
+                temp = np.zeros((original_h, original_w, 1))
+                for i, mask in enumerate(masks):
+                    mask = mask.astype(np.uint8)
+                    if not retina:
+                        mask = cv2.resize(mask, (original_w, original_h), interpolation=cv2.INTER_NEAREST)
+                    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+                    contour_all.extend(iter(contours))
+                cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2)
+                color = np.array([0 / 255, 0 / 255, 1.0, 0.8])
+                contour_mask = temp / 255 * color.reshape(1, 1, -1)
+                plt.imshow(contour_mask)
+
+        plt.axis("off")
+        plt.close()
+        return fig2img(fig)
 
 def segment_image(input_image, object_name):
     try: