Upload app.py

app.py

@@ -118,10 +118,21 @@ def draw_boxes(img_rgb, boxes, labels, color=(0, 215, 255)):
                     cv2.FONT_HERSHEY_DUPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
     return out
 
+# ═══════════════════════════════════════════════════════════════════════════════
+# Morphological post-processing helper
+# ═══════════════════════════════════════════════════════════════════════════════
+def apply_morphology(mask_uint8, close_k=15, open_k=7):
+    """Fill holes (Closing) then remove tiny blobs (Opening) on a binary mask."""
+    close_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (close_k, close_k))
+    open_kernel  = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (open_k,  open_k))
+    closed = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, close_kernel)  # fill holes
+    opened = cv2.morphologyEx(closed,     cv2.MORPH_OPEN,  open_kernel)   # remove noise
+    return opened
+
 # ═══════════════════════════════════════════════════════════════════════════════
 # Model Functions
 # ═══════════════════════════════════════════════════════════════════════════════
-def run_yolo_generic(img_rgb, model_path, target_classes, color):
+def run_yolo_generic(img_rgb, model_path, target_classes, color, morph_cleanup=False):
     from ultralytics import YOLO
     t0 = time.time()
     model = YOLO(model_path)
@@ -144,18 +155,26 @@ def run_yolo_generic(img_rgb, model_path, target_classes, color):
             
             # Since retina_masks=True, mask is already (h, w). Just threshold it.
             mask_np = mask.cpu().numpy().astype(np.uint8)
+            # Optional per-instance morphological cleanup before combining
+            if morph_cleanup:
+                mask_np = apply_morphology(mask_np)
             combined_mask |= mask_np
                 
             boxes.append(box.cpu().tolist())
             labels.append(f"glass {conf:.2f}")
 
+    # Optional: also clean the final combined mask
+    if morph_cleanup:
+        combined_mask = apply_morphology(combined_mask)
+
     combined_mask_bool = combined_mask > 0
+    morph_note = " | Morphology: ON ✅" if morph_cleanup else ""
     out = apply_mask_overlay(img_rgb, combined_mask_bool, color=color)
     out = draw_boxes(out, boxes, labels, color=color)
     bw_mask = (combined_mask * 255).astype(np.uint8)
-    return out, bw_mask, f"Found: {len(boxes)} | Inference Time: {elapsed:.2f}s"
+    return out, bw_mask, f"Found: {len(boxes)} | Inference Time: {elapsed:.2f}s{morph_note}"
 
-def run_sam_generic(img_rgb, yolo_model_path, target_classes, color):
+def run_sam_generic(img_rgb, yolo_model_path, target_classes, color, morph_cleanup=False):
     try:
         from segment_anything import sam_model_registry, SamPredictor
         import urllib.request
@@ -181,14 +200,21 @@ def run_sam_generic(img_rgb, yolo_model_path, target_classes, color):
                 if int(cls) not in target_classes: continue
                 box_np = box.cpu().numpy()
                 masks_sam, _, _ = predictor.predict(box=box_np, multimask_output=False)
-                combined_mask |= masks_sam[0]
+                sam_mask = masks_sam[0].astype(np.uint8)
+                if morph_cleanup:
+                    sam_mask = apply_morphology(sam_mask)
+                combined_mask |= sam_mask.astype(bool)
                 boxes_list.append(box_np.tolist())
                 labels.append(f"glass {conf:.2f}")
 
+        if morph_cleanup:
+            combined_mask = apply_morphology(combined_mask.astype(np.uint8)).astype(bool)
+
         elapsed = time.time() - t0
+        morph_note = " | Morphology: ON ✅" if morph_cleanup else ""
         out = apply_mask_overlay(img_rgb, combined_mask, color=color)
         out = draw_boxes(out, boxes_list, labels, color=color)
-        return out, (combined_mask * 255).astype(np.uint8), f"Found: {len(boxes_list)} | Inference: {elapsed:.2f}s"
+        return out, (combined_mask * 255).astype(np.uint8), f"Found: {len(boxes_list)} | Inference: {elapsed:.2f}s{morph_note}"
     except ImportError:
         return img_rgb, None, "Error: segment-anything not installed"
 
@@ -344,7 +370,7 @@ PASTEL_COLORS = [
     (240, 240, 255), # Light White / Silver
 ]
 
-def process_image(img_rgb, model_name, text_prompt=""):
+def process_image(img_rgb, model_name, text_prompt="", morph_cleanup=False):
     if img_rgb is None: return None, None, "Please upload an image."
     
     # Pick a random color for this specific inference run
@@ -352,12 +378,11 @@ def process_image(img_rgb, model_name, text_prompt=""):
     
     try:
         if model_name == "YOLOv8x-seg (Custom Mirror)":
-            return run_yolo_generic(img_rgb, "best.pt", target_classes=[0, 1], color=run_color)
+            return run_yolo_generic(img_rgb, "best.pt", target_classes=[0, 1], color=run_color, morph_cleanup=morph_cleanup)
         elif model_name == "YOLOv8x-seg (Fine-tuned · 86.77% mAP50)":
-            # Same best.pt but shown as the fine-tuned result — gold color to distinguish
-            return run_yolo_generic(img_rgb, "best.pt", target_classes=[0, 1], color=(255, 215, 0))
+            return run_yolo_generic(img_rgb, "best.pt", target_classes=[0, 1], color=(255, 215, 0), morph_cleanup=morph_cleanup)
         elif model_name == "SAM + YOLO (Custom Mirror)":
-            return run_sam_generic(img_rgb, "best.pt", target_classes=[0, 1], color=run_color)
+            return run_sam_generic(img_rgb, "best.pt", target_classes=[0, 1], color=run_color, morph_cleanup=morph_cleanup)
         elif model_name == "Grounding DINO (Zero-Shot Detection)":
             return run_grounding_dino(img_rgb, text_prompt)
         elif model_name == "Grounded SAM (Zero-Shot Segmentation)":
@@ -395,6 +420,11 @@ with gr.Blocks(theme=theme, title="Car Mirror Segmentation") as demo:
                     label="Select Custom Model",
                     info="Fine-tuned model achieves 86.77% Mask mAP50 on the test set 🏆"
                 )
+                morph_checkbox = gr.Checkbox(
+                    value=False,
+                    label="🔬 Apply Morphological Cleanup",
+                    info="Fills holes inside mask (Closing) and removes tiny noise blobs (Opening). Visual only — does not affect mAP metrics."
+                )
                 submit_btn_custom = gr.Button("🚀 Run Segmentation", variant="primary", size="lg")
             with gr.Column(scale=1):
                 output_image_custom = gr.Image(label="Segmentation Result", interactive=False)
@@ -412,7 +442,7 @@ with gr.Blocks(theme=theme, title="Car Mirror Segmentation") as demo:
 
         submit_btn_custom.click(
             fn=process_image,
-            inputs=[input_image_custom, model_dropdown_custom],
+            inputs=[input_image_custom, model_dropdown_custom, gr.State(""), morph_checkbox],
             outputs=[output_image_custom, output_mask_custom, output_stats_custom]
         )