Update app.py

app.py

@@ -2,44 +2,76 @@ import os
 import glob
 import shutil
 import tempfile
+
 import gradio as gr
 import cv2
 from PIL import Image
 from ultralytics import YOLO
 
-# Keep Ultralytics config in writable location
+# ------------------------------------------------------------------
+# HF Spaces requirement: Ultralytics config must be writable
+# ------------------------------------------------------------------
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
 
+# ------------------------------------------------------------------
+# Load YOLOv8 segmentation model (best.pt in same folder)
+# ------------------------------------------------------------------
 model = YOLO("best.pt")
 
+
+def unwrap_image(item):
+    """
+    Normalize Gradio Gallery item to a PIL.Image.
+
+    Gradio Gallery may return:
+    - PIL.Image
+    - dict with {"data": PIL.Image, ...}
+    - tuple (e.g., (PIL.Image, metadata) or (name, PIL.Image))
+    """
+    # Case 1: dict
+    if isinstance(item, dict):
+        return item["data"]
+
+    # Case 2: tuple
+    if isinstance(item, tuple):
+        for v in item:
+            if hasattr(v, "save"):  # PIL.Image has .save()
+                return v
+        raise ValueError("Gallery tuple does not contain a PIL image")
+
+    # Case 3: already PIL.Image
+    return item
+
+
 def segment(gallery_items, conf=0.25, imgsz=640):
     """
-    gallery_items from gr.Gallery(type="pil") can be:
-      - list[dict] where dict has {"data": PIL.Image, ...}
-      - list[PIL.Image] depending on gradio version
-    We handle both safely.
+    Runs YOLOv8 segmentation on multiple images using the SAME
+    pipeline as the training notebook: model.predict(save=True).
+    Returns segmentation overlays as a flat list of PIL images.
     """
 
     if not gallery_items:
         return []
 
-    # Create temp workspace
+    # ------------------------------------------------------------------
+    # Create a temporary working directory
+    # ------------------------------------------------------------------
     workdir = tempfile.mkdtemp(prefix="yolo_seg_")
     src_dir = os.path.join(workdir, "src")
     os.makedirs(src_dir, exist_ok=True)
 
-    # Save incoming images
-    saved_paths = []
-    for i, item in enumerate(gallery_items):
-        img = item["data"] if isinstance(item, dict) else item
-        out_path = os.path.join(src_dir, f"img_{i}.jpg")
-        img.save(out_path)
-        saved_paths.append(out_path)
-
-    print(f"[DEBUG] Saved {len(saved_paths)} images to: {src_dir}")
-
     try:
-        # Run YOLO like your notebook
+        # --------------------------------------------------------------
+        # Save uploaded images to disk
+        # --------------------------------------------------------------
+        for i, item in enumerate(gallery_items):
+            img = unwrap_image(item)
+            out_path = os.path.join(src_dir, f"img_{i}.jpg")
+            img.save(out_path)
+
+        # --------------------------------------------------------------
+        # Run YOLO prediction (NOTEBOOK-EQUIVALENT)
+        # --------------------------------------------------------------
         preds = model.predict(
             source=src_dir,
             imgsz=int(imgsz),
@@ -48,41 +80,34 @@ def segment(gallery_items, conf=0.25, imgsz=640):
             save=True,
             project=workdir,
             name="preds",
-            exist_ok=True
+            exist_ok=True,
         )
 
-        # Ultralytics will save overlays into:
-        # {workdir}/preds/
+        # YOLO saves overlays into: {workdir}/preds/
         pred_dir = os.path.join(workdir, "preds")
-        print(f"[DEBUG] pred_dir = {pred_dir}")
-        print(f"[DEBUG] pred_dir exists? {os.path.exists(pred_dir)}")
-
-        # Find overlay images YOLO saved
-        files = []
-        for ext in ("*.jpg", "*.png", "*.jpeg", "*.bmp"):
-            files.extend(glob.glob(os.path.join(pred_dir, ext)))
-
-        files = sorted(files)
-        print(f"[DEBUG] overlays found: {len(files)}")
-        print(f"[DEBUG] sample files: {files[:5]}")
 
-        # Load overlays -> PIL
+        # --------------------------------------------------------------
+        # Load saved overlay images
+        # --------------------------------------------------------------
         outputs = []
-        for f in files:
-            im = cv2.imread(f)
-            if im is None:
-                print(f"[WARN] cv2 failed to read: {f}")
-                continue
-            im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
-            outputs.append(Image.fromarray(im))
+        for ext in ("*.jpg", "*.png", "*.jpeg", "*.bmp"):
+            for f in sorted(glob.glob(os.path.join(pred_dir, ext))):
+                im = cv2.imread(f)
+                if im is None:
+                    continue
+                im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
+                outputs.append(Image.fromarray(im))
 
         return outputs
 
     finally:
-        # Cleanup temp files
+        # Always clean up temp files
         shutil.rmtree(workdir, ignore_errors=True)
 
 
+# ------------------------------------------------------------------
+# Gradio Interface (HF-safe)
+# ------------------------------------------------------------------
 demo = gr.Interface(
     fn=segment,
     inputs=[
@@ -90,8 +115,14 @@ demo = gr.Interface(
         gr.Slider(0.01, 0.9, value=0.25, step=0.01, label="Confidence"),
         gr.Slider(320, 1280, value=640, step=32, label="Image size"),
     ],
-    outputs=gr.Gallery(label="Predicted overlays"),
-    title="YOLOv8 Segmentation (Multi-Image) — Notebook Equivalent",
+    outputs=gr.Gallery(label="YOLOv8 Segmentation Overlays"),
+    title="YOLOv8 Segmentation (Hugging Face Space)",
+    description="Multi-image YOLOv8 segmentation using the same predict(save=True) pipeline as training.",
 )
 
-demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+# IMPORTANT: disable SSR to avoid asyncio cleanup noise on HF
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    ssr_mode=False,
+)