Upload app.py

app.py

@@ -4,42 +4,17 @@ import cv2
 import time
 import torch
 import warnings
-import glob
-import os
-import zipfile
-import shutil
 from PIL import Image
 
 warnings.filterwarnings("ignore")
 
-# ─── ZIP Extraction on Startup ────────────────────────────────────────────────
-# This runs ONCE when the app starts, extracts test.zip into clean folders.
-def extract_zip_if_needed(zip_name, extract_to):
-    if os.path.exists(zip_name) and not os.path.exists(extract_to):
-        print(f"Extracting {zip_name} -> {extract_to} ...")
-        with zipfile.ZipFile(zip_name, 'r') as zf:
-            zf.extractall(extract_to)
-        print("Extraction done!")
-
-extract_zip_if_needed("test.zip",  "test_images")
-extract_zip_if_needed("cars.zip",  "car_images")
-
-def get_images_from_folder(folder, limit=10):
-    """Recursively finds all jpg/png files inside a folder."""
-    found = []
-    for ext in ["*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"]:
-        found += glob.glob(os.path.join(folder, "**", ext), recursive=True)
-    found = sorted(found)[:limit]
-    if not found:
-        return [["car.jpeg"]]
-    return [[f] for f in found]
-
-# ─── Global Settings ──────────────────────────────────────────────────────────
+# Global Settings
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-CONF   = 0.45
+CONF = 0.45
 
 # ─── Helpers ──────────────────────────────────────────────────────────────────
 def apply_mask_overlay(img_rgb, mask_bool, color=(0, 220, 100), alpha=0.45):
+    """Blends binary mask with image as a colored overlay."""
     overlay = img_rgb.copy()
     overlay[mask_bool] = color
     return cv2.addWeighted(overlay, alpha, img_rgb, 1 - alpha, 0)
@@ -48,13 +23,16 @@ def draw_boxes(img_rgb, boxes, labels, color=(0, 220, 100)):
     out = img_rgb.copy()
     for box, label in zip(boxes, labels):
         x1, y1, x2, y2 = map(int, box)
+        # Bounding box ki thickness 2 se 1 kar di
         cv2.rectangle(out, (x1, y1), (x2, y2), color, 1)
+        # Font size 0.6 se 0.4 aur thickness 1 kar di taa ke small cars clear nazar aayein
         cv2.putText(out, label, (x1, max(y1-5, 10)),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1)
     return out
 
 # ─── Model Functions ──────────────────────────────────────────────────────────
 def run_yolo_generic(img_rgb, model_path, target_classes, color):
+    # pyrefly: ignore [missing-import]
     from ultralytics import YOLO
     t0 = time.time()
     model = YOLO(model_path)
@@ -90,10 +68,11 @@ def run_yolo_generic(img_rgb, model_path, target_classes, color):
 def run_sam_generic(img_rgb, yolo_model_path, target_classes, color):
     try:
         from segment_anything import sam_model_registry, SamPredictor
-        import urllib.request
+        import urllib.request, os
 
         CKPT = "sam_vit_b_01ec64.pth"
         URL  = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
+
         if not os.path.exists(CKPT):
             urllib.request.urlretrieve(URL, CKPT)
 
@@ -104,8 +83,8 @@ def run_sam_generic(img_rgb, yolo_model_path, target_classes, color):
         predictor.set_image(img_rgb)
 
         from ultralytics import YOLO as _YOLO
-        yolo     = _YOLO(yolo_model_path)
-        yolo_res = yolo(img_rgb, conf=CONF, verbose=False)[0]
+        yolo      = _YOLO(yolo_model_path)
+        yolo_res  = yolo(img_rgb, conf=CONF, verbose=False)[0]
 
         h, w = img_rgb.shape[:2]
         combined_mask = np.zeros((h, w), dtype=bool)
@@ -120,7 +99,10 @@ 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)
+                masks_sam, _, _ = predictor.predict(
+                    box=box_np,
+                    multimask_output=False
+                )
                 combined_mask |= masks_sam[0]
                 boxes_list.append(box_np.tolist())
                 class_name = "mirror" if 0 in target_classes or 1 in target_classes else "car"
@@ -140,7 +122,9 @@ def run_sam_generic(img_rgb, yolo_model_path, target_classes, color):
         return placeholder, "Error: segment-anything not installed"
 
 def run_maskrcnn(img_rgb):
+    # pyrefly: ignore [missing-import]
     import torchvision
+    # pyrefly: ignore [missing-import]
     from torchvision.models.detection import maskrcnn_resnet50_fpn, MaskRCNN_ResNet50_FPN_Weights
     t0 = time.time()
 
@@ -159,6 +143,7 @@ def run_maskrcnn(img_rgb):
     h, w    = img_rgb.shape[:2]
     combined_mask = np.zeros((h, w), dtype=bool)
     boxes, labels = [], []
+
     COCO_LABELS = weights.meta["categories"]
 
     for mask, box, label, score in zip(
@@ -199,10 +184,13 @@ def run_segformer(img_rgb):
         logits = model(**inputs).logits
 
     h, w = img_rgb.shape[:2]
-    upsampled = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
+    upsampled = F.interpolate(
+        logits, size=(h, w), mode="bilinear", align_corners=False
+    )
     seg_map   = upsampled.argmax(dim=1)[0].cpu().numpy()
     car_mask  = seg_map == CAR_IDX
-    elapsed   = time.time() - t0
+
+    elapsed = time.time() - t0
 
     out  = apply_mask_overlay(img_rgb, car_mask, color=(255, 180, 50))
     contours, _ = cv2.findContours(
@@ -212,43 +200,52 @@ def run_segformer(img_rgb):
     n = len(contours)
     return out, f"Car regions: {n} | Inference Time: {elapsed:.2f}s"
 
-# ─── Gradio Process Function ───────────────────────────────────────────────────
+# ─── Gradio Interface ─────────────────────────────────────────────────────────
 def process_image(img_rgb, model_name):
     if img_rgb is None:
         return 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))
+            
         elif model_name == "YOLOv8x (Pretrained Car)":
             return run_yolo_generic(img_rgb, "yolov8x-seg.pt", target_classes=[2], color=(0, 200, 255))
+            
         elif model_name == "SAM + YOLO (Custom Mirror)":
             return run_sam_generic(img_rgb, "best.pt", target_classes=[0, 1], color=(255, 80, 160))
+            
         elif model_name == "SAM + YOLO (Pretrained Car)":
             return run_sam_generic(img_rgb, "yolov8x-seg.pt", target_classes=[2], color=(200, 80, 255))
+            
         elif model_name == "Mask R-CNN (Pretrained Car)":
             return run_maskrcnn(img_rgb)
+            
         elif model_name == "SegFormer (Pretrained Car)":
             return run_segformer(img_rgb)
+            
         else:
             return img_rgb, "Model not recognized."
     except Exception as e:
         return img_rgb, f"Error: {str(e)}"
 
-# ─── Build Example Lists (from extracted zip folders) ─────────────────────────
-car_examples    = get_images_from_folder("car_images")
-mirror_examples = get_images_from_folder("test_images")
-
-# ─── Gradio UI ────────────────────────────────────────────────────────────────
-theme = gr.themes.Soft(primary_hue="blue", secondary_hue="indigo")
-
-with gr.Blocks(theme=theme, title="Car vs Mirror Segmentation") as demo:
-    gr.Markdown("""
-    # 🚗 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)"):
+# Define the UI theme and layout
+theme = gr.themes.Soft(
+    primary_hue="blue",
+    secondary_hue="indigo",
+)
+
+with gr.Blocks(theme=theme, title="Car and Mirror Segmentation") as demo:
+    gr.Markdown(
+        """
+        # Car and Mirror Segmentation
+        """
+    )
+    
+    # ==========================================
+    # TAB 1: PRETRAINED FULL CAR MODELS
+    # ==========================================
+    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")
@@ -263,20 +260,27 @@ with gr.Blocks(theme=theme, title="Car vs Mirror 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_stats_car = gr.Textbox(label="Detection Statistics", interactive=False)
-
+                
+        import glob
+        import os
+        car_imgs = [f for f in glob.glob("car Images/*") if os.path.isfile(f)]
+        car_list = [[img] for img in car_imgs[:10]] # Pass only images to show thumbnails
+        if not car_list:
+            car_list = [["car.jpeg"]]
+            
         gr.Examples(
-            examples=car_examples,
+            examples=car_list,
             inputs=[input_image_car],
             examples_per_page=10,
             outputs=[output_image_car, output_stats_car],
             fn=process_image,
             cache_examples=False,
-            label="📸 Click any car image to load it"
+            label="Click any image below to test (Side by side)"
         )
 
         submit_btn_car.click(
@@ -285,8 +289,10 @@ with gr.Blocks(theme=theme, title="Car vs Mirror Segmentation") as demo:
             outputs=[output_image_car, output_stats_car]
         )
 
-    # ── TAB 2: Car Mirrors (Custom) ───────────────────────────────────────────
-    with gr.Tab("🪞 Test Car Mirrors (Custom Models)"):
+    # ==========================================
+    # TAB 2: CUSTOM MIRROR MODELS
+    # ==========================================
+    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")
@@ -299,22 +305,27 @@ with gr.Blocks(theme=theme, title="Car vs Mirror Segmentation") as demo:
                     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")
-
+                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)
-
+                
+        mirror_imgs = [f for f in glob.glob("test car windows/*") if os.path.isfile(f)]
+        mirror_list = [[img] for img in mirror_imgs[:10]] # Pass only images to show thumbnails
+        if not mirror_list:
+            mirror_list = [["car.jpeg"]]
+            
         gr.Examples(
-            examples=mirror_examples,
+            examples=mirror_list,
             inputs=[input_image_mirror],
             examples_per_page=10,
             outputs=[output_image_mirror, output_stats_mirror],
             fn=process_image,
             cache_examples=False,
-            label="📸 Click any mirror image to load it"
+            label="Click any image below to test (Side by side)"
         )
-
+        
         submit_btn_mirror.click(
             fn=process_image,
             inputs=[input_image_mirror, model_dropdown_mirror],