Upload app.py

app.py

@@ -148,7 +148,8 @@ def run_yolo_generic(img_rgb, model_path, target_classes, color):
 
     out = apply_mask_overlay(img_rgb, combined_mask, color=color)
     out = draw_boxes(out, boxes, labels, color=color)
-    return out, f"Found: {len(boxes)} | Inference Time: {elapsed:.2f}s"
+    bw_mask = (combined_mask * 255).astype(np.uint8)
+    return out, bw_mask, f"Found: {len(boxes)} | Inference Time: {elapsed:.2f}s"
 
 def run_sam_generic(img_rgb, yolo_model_path, target_classes, color):
     try:
@@ -190,14 +191,15 @@ def run_sam_generic(img_rgb, yolo_model_path, target_classes, color):
         elapsed = time.time() - t0
         out = apply_mask_overlay(img_rgb, combined_mask, color=color)
         out = draw_boxes(out, boxes_list, labels, color=color)
-        return out, f"Found: {len(boxes_list)} | Inference Time: {elapsed:.2f}s"
+        bw_mask = (combined_mask * 255).astype(np.uint8)
+        return out, bw_mask, f"Found: {len(boxes_list)} | Inference Time: {elapsed:.2f}s"
 
     except ImportError:
         h, w = img_rgb.shape[:2]
         placeholder = img_rgb.copy()
         cv2.putText(placeholder, "SAM: pip install segment-anything",
                     (20, h//2), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,80,160), 2)
-        return placeholder, "Error: segment-anything not installed"
+        return placeholder, None, "Error: segment-anything not installed"
 
 def run_maskrcnn(img_rgb):
     import torchvision
@@ -231,7 +233,8 @@ def run_maskrcnn(img_rgb):
         labels.append(f"car {score:.2f}")
     out = apply_mask_overlay(img_rgb, combined_mask, color=(80, 120, 255))
     out = draw_boxes(out, boxes, labels, color=(80, 120, 255))
-    return out, f"Cars found: {len(boxes)} | Inference Time: {elapsed:.2f}s"
+    bw_mask = (combined_mask * 255).astype(np.uint8)
+    return out, bw_mask, f"Cars found: {len(boxes)} | Inference Time: {elapsed:.2f}s"
 
 def run_segformer(img_rgb):
     from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
@@ -256,14 +259,15 @@ def run_segformer(img_rgb):
         car_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
     )
     out = cv2.drawContours(out, contours, -1, (255, 180, 50), 2)
-    return out, f"Car regions: {len(contours)} | Inference Time: {elapsed:.2f}s"
+    bw_mask = (car_mask * 255).astype(np.uint8)
+    return out, bw_mask, f"Car regions: {len(contours)} | Inference Time: {elapsed:.2f}s"
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # Gradio Process Function
 # ═══════════════════════════════════════════════════════════════════════════════
 def process_image(img_rgb, model_name):
     if img_rgb is None:
-        return None, "Please upload an image."
+        return None, None, "Please upload an image."
     try:
         if model_name == "YOLOv8x-seg (Custom Mirror)":
             return run_yolo_generic(img_rgb, "best.pt", target_classes=[0, 1], color=(50, 220, 100))
@@ -278,22 +282,65 @@ def process_image(img_rgb, model_name):
         elif model_name == "SegFormer (Pretrained Car)":
             return run_segformer(img_rgb)
         else:
-            return img_rgb, "Model not recognized."
+            return img_rgb, None, "Model not recognized."
     except Exception as e:
-        return img_rgb, f"Error: {str(e)}"
+        return img_rgb, None, f"Error: {str(e)}"
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # Gradio UI
 # ═══════════════════════════════════════════════════════════════════════════════
 theme = gr.themes.Soft(primary_hue="blue", secondary_hue="indigo")
 
-with gr.Blocks(theme=theme, title="Car and Windows Segmentation") as demo:
+with gr.Blocks(theme=theme, title="Car vs Mirror Segmentation") as demo:
     gr.Markdown("""
-    #  Car and Windows Segmentation
+    # 🚗 Car vs Mirror Segmentation Comparison
+    Compare your custom trained **Car Mirror** models against pretrained **Full Car** models.
     """)
 
-    # ── TAB 1: Full Car (Pretrained) ──
-    with gr.Tab(" Test Full Cars (Pretrained Models)"):
+    # ── TAB 1: Car Mirrors (Custom) - NOW THE DEFAULT TAB ──
+    with gr.Tab("🪞 Test Car Mirrors (Custom Models)"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image_mirror = gr.Image(type="numpy", label="Upload Mirror Image")
+                model_dropdown_mirror = gr.Dropdown(
+                    choices=[
+                        "YOLOv8x-seg (Custom Mirror)",
+                        "SAM + YOLO (Custom Mirror)"
+                    ],
+                    value="YOLOv8x-seg (Custom Mirror)",
+                    label="Select Custom Model",
+                    info="These models are specifically trained to detect car mirrors."
+                )
+                submit_btn_mirror = gr.Button("🚀 Run Segmentation", variant="primary", size="lg")
+            with gr.Column(scale=1):
+                output_image_mirror = gr.Image(label="Segmentation Result", interactive=False)
+                output_mask_mirror = gr.Image(label="Binary Mask (White=Object, Black=Background)", interactive=False)
+                output_stats_mirror = gr.Textbox(label="Detection Statistics", interactive=False)
+
+        if mirror_examples:
+            gr.Markdown("### 📸 Click any mirror image below to load it")
+            mirror_gallery = gr.Gallery(
+                value=mirror_examples,
+                columns=10,
+                height=120,
+                object_fit="cover",
+                allow_preview=False,
+                show_label=False
+            )
+            
+            def load_mirror_img(evt: gr.SelectData):
+                return mirror_examples[evt.index]
+                
+            mirror_gallery.select(fn=load_mirror_img, inputs=None, outputs=input_image_mirror)
+
+        submit_btn_mirror.click(
+            fn=process_image,
+            inputs=[input_image_mirror, model_dropdown_mirror],
+            outputs=[output_image_mirror, output_mask_mirror, output_stats_mirror]
+        )
+
+    # ── TAB 2: Full Car (Pretrained) ──
+    with gr.Tab("🚘 Test Full Cars (Pretrained Models)"):
         with gr.Row():
             with gr.Column(scale=1):
                 input_image_car = gr.Image(type="numpy", label="Upload Car Image")
@@ -308,16 +355,19 @@ with gr.Blocks(theme=theme, title="Car and Windows Segmentation") as demo:
                     label="Select Pretrained Model",
                     info="These models are pretrained from the internet to detect full cars."
                 )
-                submit_btn_car = gr.Button(" Run Segmentation", variant="primary", size="lg")
+                submit_btn_car = gr.Button("🚀 Run Segmentation", variant="primary", size="lg")
             with gr.Column(scale=1):
                 output_image_car = gr.Image(label="Segmentation Result", interactive=False)
+                output_mask_car = gr.Image(label="Binary Mask (White=Object, Black=Background)", interactive=False)
                 output_stats_car = gr.Textbox(label="Detection Statistics", interactive=False)
 
         if car_examples:
-            gr.Markdown("###  Click any car image below to load it")
+            gr.Markdown("### 📸 Click any car image below to load it")
             car_gallery = gr.Gallery(
                 value=car_examples,
                 columns=10,
+                height=120,
+                object_fit="cover",
                 allow_preview=False,
                 show_label=False
             )
@@ -330,46 +380,7 @@ with gr.Blocks(theme=theme, title="Car and Windows Segmentation") as demo:
         submit_btn_car.click(
             fn=process_image,
             inputs=[input_image_car, model_dropdown_car],
-            outputs=[output_image_car, output_stats_car]
-        )
-
-    # ── TAB 2: Car Mirrors (Custom) ──
-    with gr.Tab(" Test Car Mirrors (Custom Models)"):
-        with gr.Row():
-            with gr.Column(scale=1):
-                input_image_mirror = gr.Image(type="numpy", label="Upload Mirror Image")
-                model_dropdown_mirror = gr.Dropdown(
-                    choices=[
-                        "YOLOv8x-seg (Custom Mirror)",
-                        "SAM + YOLO (Custom Mirror)"
-                    ],
-                    value="YOLOv8x-seg (Custom Mirror)",
-                    label="Select Custom Model",
-                    info="These models are specifically trained to detect car mirrors."
-                )
-                submit_btn_mirror = gr.Button("🚀 Run Segmentation", variant="primary", size="lg")
-            with gr.Column(scale=1):
-                output_image_mirror = gr.Image(label="Segmentation Result", interactive=False)
-                output_stats_mirror = gr.Textbox(label="Detection Statistics", interactive=False)
-
-        if mirror_examples:
-            gr.Markdown("###  Click any mirror image below to load it")
-            mirror_gallery = gr.Gallery(
-                value=mirror_examples,
-                columns=10,
-                allow_preview=False,
-                show_label=False
-            )
-            
-            def load_mirror_img(evt: gr.SelectData):
-                return mirror_examples[evt.index]
-                
-            mirror_gallery.select(fn=load_mirror_img, inputs=None, outputs=input_image_mirror)
-
-        submit_btn_mirror.click(
-            fn=process_image,
-            inputs=[input_image_mirror, model_dropdown_mirror],
-            outputs=[output_image_mirror, output_stats_mirror]
+            outputs=[output_image_car, output_mask_car, output_stats_car]
         )
 
 if __name__ == "__main__":