Update app.py

app.py

@@ -22,54 +22,8 @@ def fig2img(fig):
     return img
 
 def plot(annotations, prompt_process, mask_random_color=True, better_quality=True, retina=True, with_contours=True):
-    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)
+    # ... (keep the existing plot function as is)
+    # This function doesn't need modification for our purposes
 
 def segment_image(input_image, object_name):
     try:
@@ -78,8 +32,8 @@ def segment_image(input_image, object_name):
         
         input_image = Image.fromarray(input_image).convert("RGB")
         
-        # Run FastSAM model
-        everything_results = model(input_image, retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
+        # Run FastSAM model with adjusted parameters
+        everything_results = model(input_image, retina_masks=True, imgsz=1024, conf=0.25, iou=0.7)
         
         # Prepare a Prompt Process object
         prompt_process = FastSAMPrompt(input_image, everything_results, device=device)
@@ -90,6 +44,20 @@ def segment_image(input_image, object_name):
         if not results:
             return input_image, f"Could not find '{object_name}' in the image."
         
+        # Post-process the masks
+        for ann in results:
+            if ann.masks is not None:
+                masks = ann.masks.data
+                if isinstance(masks[0], torch.Tensor):
+                    masks = np.array(masks.cpu())
+                for i, mask in enumerate(masks):
+                    # Apply more aggressive morphological operations
+                    kernel = np.ones((5,5), np.uint8)
+                    mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
+                    mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, kernel)
+                    masks[i] = cv2.dilate(mask, kernel, iterations=2)
+                ann.masks.data = masks
+        
         # Plot the results
         result_image = plot(annotations=results, prompt_process=prompt_process)