Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import io
 import cv2
 import sys
-import math
 import json
 import torch
 import gradio as gr
@@ -10,27 +9,31 @@ import numpy as np
 import pandas as pd
 from PIL import Image
 from typing import List, Tuple, Optional, Dict
-
 from ultralytics import YOLO
 import supervision as sv
 from huggingface_hub import hf_hub_download
-import spaces  # ZeroGPU-compatible decorator
+import spaces
+import tempfile
+
+# ------------------------------------------------------------------
+# Silence Ultralytics config dir warning in read-only home directories
+# ------------------------------------------------------------------
+os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics")
 
 # -----------------------------
 # Defaults / configuration
 # -----------------------------
-REPO_ID = "edeler/ICC"          # <- your HF repo with weights
-WEIGHTS_FILENAME = "best.pt"     # <- adjust if different
+REPO_ID = "edeler/ICC"          # your HF repo with weights
+WEIGHTS_FILENAME = "best.pt"     # adjust if different
 LOCAL_MODEL_DIR = "./models/ICC"
-EXAMPLES_DIR = "."               # scan current repo root for demo images
+EXAMPLES_DIR = "."               # scan repo root for demo images
 
-# Reasonable defaults (tunable in UI)
 DEFAULT_CONF = 0.25
-DEFAULT_IOU = 0.50               # for both model-level and global cross-slice NMS
+DEFAULT_IOU = 0.50
 DEFAULT_SLICE_WH = 1024
-DEFAULT_OVERLAP = 128            # helps stitching at tile borders
+DEFAULT_OVERLAP = 128
 DEFAULT_THICKNESS = 3
-DEFAULT_LONG_EDGE = 4096         # optional downscale for very large WSIs/crops
+DEFAULT_LONG_EDGE = 4096
 
 # -----------------------------
 # Torch / device helpers
@@ -60,8 +63,7 @@ def load_model() -> Tuple[YOLO, Dict[int, str]]:
     weights_path = hf_hub_download(
         repo_id=REPO_ID,
         filename=WEIGHTS_FILENAME,
-        local_dir=LOCAL_MODEL_DIR,
-        local_dir_use_symlinks=False,  # safer in Spaces
+        local_dir=LOCAL_MODEL_DIR,  # no symlink arg (deprecated)
     )
     model = YOLO(weights_path)
     class_names = model.model.names if hasattr(model, "model") else model.names
@@ -72,7 +74,6 @@ def load_model() -> Tuple[YOLO, Dict[int, str]]:
 # Image utilities
 # -----------------------------
 def ensure_bgr(img: np.ndarray) -> np.ndarray:
-    # Gradio provides RGB; OpenCV expects BGR
     if img.ndim == 3 and img.shape[2] == 3:
         return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
     return img
@@ -102,16 +103,9 @@ def run_sliced_inference(
     overlap_h: int,
     device: str,
 ) -> sv.Detections:
-    """
-    Uses supervision.InferenceSlicer to run model across tiles,
-    then returns all detections (to be merged via cross-slice NMS).
-    """
-    # inner callback called by slicer
     @torch.inference_mode()
     def callback(tile_bgr: np.ndarray) -> sv.Detections:
-        # Ultralytics expects RGB
         tile_rgb = cv2.cvtColor(tile_bgr, cv2.COLOR_BGR2RGB)
-        # Predict with thresholds at model level (faster than filtering post-hoc)
         results = model.predict(
             source=tile_rgb,
             conf=conf,
@@ -120,9 +114,7 @@ def run_sliced_inference(
             verbose=False,
             half=half_precision_available(device),
         )
-        res = results[0]
-        det = sv.Detections.from_ultralytics(res)
-        return det
+        return sv.Detections.from_ultralytics(results[0])
 
     slicer = sv.InferenceSlicer(
         callback=callback,
@@ -131,7 +123,6 @@ def run_sliced_inference(
         overlap_ratio_wh=None,
     )
     detections = slicer(image_bgr)
-    # Cross-slice NMS to merge duplicates at tile seams
     detections = detections.with_nms(threshold=iou, class_agnostic=False)
     return detections
 
@@ -150,11 +141,11 @@ def make_labels(det: sv.Detections, names: Dict[int, str], show_labels: bool) ->
 def detections_to_dataframe(det: sv.Detections, names: Dict[int, str]) -> pd.DataFrame:
     if len(det) == 0:
         return pd.DataFrame(columns=["class_id", "class_name", "confidence", "x_min", "y_min", "x_max", "y_max"])
-    xyxy = det.xyxy  # (N,4)
-    data = []
+    xyxy = det.xyxy
+    rows = []
     for i in range(len(det)):
         cls = int(det.class_id[i])
-        data.append({
+        rows.append({
             "class_id": cls,
             "class_name": names.get(cls, str(cls)),
             "confidence": float(det.confidence[i]),
@@ -163,7 +154,7 @@ def detections_to_dataframe(det: sv.Detections, names: Dict[int, str]) -> pd.Dat
             "x_max": float(xyxy[i, 2]),
             "y_max": float(xyxy[i, 3]),
         })
-    return pd.DataFrame(data)
+    return pd.DataFrame(rows)
 
 def per_class_summary(df: pd.DataFrame) -> str:
     if df.empty:
@@ -176,7 +167,7 @@ def per_class_summary(df: pd.DataFrame) -> str:
 # -----------------------------
 # Gradio inference function
 # -----------------------------
-@spaces.GPU  # ensures ZeroGPU allocation during this call
+@spaces.GPU
 def detect_objects(
     image: np.ndarray,
     conf: float,
@@ -195,16 +186,13 @@ def detect_objects(
         if image is None:
             raise ValueError("Please upload or select an image.")
 
-        # Prepare image (BGR) and optional downscale
         image_bgr = ensure_bgr(image)
-        image_bgr, scale = maybe_downscale_long_edge(image_bgr, long_edge)
+        image_bgr, _ = maybe_downscale_long_edge(image_bgr, long_edge)
 
-        # Load model + names lazily
         progress(0.05, desc="Loading model…")
         model, names = load_model()
         device = get_device()
 
-        # Inference (sliced)
         progress(0.35, desc="Running sliced inference…")
         with torch.inference_mode():
             detections = run_sliced_inference(
@@ -219,59 +207,55 @@ def detect_objects(
                 device=device,
             )
 
-        # Optional class filtering (by names)
         if selected_classes:
             allow_ids = {cid for cid, cname in names.items() if cname in set(selected_classes)}
             if len(detections) > 0:
                 mask = np.array([int(c) in allow_ids for c in detections.class_id], dtype=bool)
                 detections = detections[mask]
 
-        # Create labels (optional)
         labels = make_labels(detections, names, show_labels)
 
-        # Annotate
         progress(0.65, desc="Annotating…")
         annotator = sv.BoxAnnotator(thickness=thickness)
         annotated = annotator.annotate(scene=image_bgr.copy(), detections=detections, labels=labels)
-
-        # Convert to RGB for display
         annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
         annotated_pil = Image.fromarray(annotated_rgb)
 
-        # Tabular output + downloadable CSV
         df = detections_to_dataframe(detections, names)
-        csv_bytes = df.to_csv(index=False).encode("utf-8")
         summary = per_class_summary(df)
 
+        # Create a temporary CSV file for robust downloads on older Gradio builds
+        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
+        df.to_csv(tmp.name, index=False)
+        csv_path = tmp.name
+
         progress(1.0, desc="Done.")
-        return annotated_pil, summary, df, csv_bytes
+        return annotated_pil, summary, df, csv_path
 
     except Exception as e:
-        # Return clean error in summary field
         empty_img = None
         empty_df = pd.DataFrame(columns=["class_id", "class_name", "confidence", "x_min", "y_min", "x_max", "y_max"])
-        return empty_img, f"Error: {repr(e)}", empty_df, b""
+        return empty_img, f"Error: {repr(e)}", empty_df, None
 
 # -----------------------------
 # UI / App
 # -----------------------------
 def discover_examples(root: str) -> List[str]:
     exts = {".jpg", ".jpeg", ".png"}
-    paths = []
+    out = []
     try:
         for fname in os.listdir(root):
             if os.path.splitext(fname.lower())[1] in exts:
-                paths.append(os.path.join(root, fname))
+                out.append(os.path.join(root, fname))
     except Exception:
         pass
-    # Keep at most a handful to keep UI tidy
-    return sorted(paths)[:8]
+    return sorted(out)[:8]
 
 def reset_all():
     # image, output_img, summary, table, download
     return gr.update(value=None), gr.update(value=None), gr.update(value=""), pd.DataFrame(), None
 
-with gr.Blocks(title="Interstitial Cell of Cajal Detection and Quantification Tool", fill_height=True) as demo:
+with gr.Blocks(title="Interstitial Cell of Cajal Detection and Quantification Tool") as demo:
     gr.Markdown("<h1>Interstitial Cell of Cajal (ICC) Detection and Quantification</h1>"
                 "<p>YOLO-based tiled inference with cross-slice NMS. Adjust parameters under <em>Advanced Settings</em>.</p>")
 
@@ -296,8 +280,6 @@ with gr.Blocks(title="Interstitial Cell of Cajal Detection and Quantification To
                 thickness = gr.Slider(1, 8, value=DEFAULT_THICKNESS, step=1, label="Bounding box thickness")
                 show_labels = gr.Checkbox(value=True, label="Show class + confidence labels")
                 long_edge = gr.Slider(512, 8192, value=DEFAULT_LONG_EDGE, step=64, label="Optional downscale — max long edge (px)")
-                # Dynamic class list (populated on load)
-                # We attempt to load names now; if it fails (cold start), we show an empty multiselect.
                 try:
                     _, _names = load_model()
                     class_list = [v for _, v in sorted(_names.items())]
@@ -312,28 +294,25 @@ with gr.Blocks(title="Interstitial Cell of Cajal Detection and Quantification To
         with gr.Column(scale=1):
             output_img = gr.Image(label="Detection Result", interactive=False)
             detection_summary = gr.Textbox(label="Detection Summary", interactive=False)
+            # NOTE: remove unsupported 'height' kwarg for older Gradio
             detections_table = gr.Dataframe(
-                headers=["class_id", "class_name", "confidence", "x_min", "y_min", "x_max", "y_max"],
                 label="Detections (table)",
                 interactive=False,
-                wrap=True,
-                height=240,
             )
-            download_csv = gr.DownloadButton(label="Download detections as CSV", value=None, file_name="detections.csv")
+            # Use gr.File for robust downloads across Gradio versions
+            download_csv = gr.File(label="Download detections as CSV")
 
-    # Wire buttons
     predict.click(
         detect_objects,
         inputs=[input_img, conf, iou, slice_w, slice_h, overlap_w, overlap_h, thickness, show_labels, selected_classes, long_edge],
         outputs=[output_img, detection_summary, detections_table, download_csv],
     )
+
     clear.click(
         reset_all,
         inputs=None,
         outputs=[input_img, output_img, detection_summary, detections_table, download_csv],
     )
 
-# Recommended for Spaces stability / concurrency
 demo.queue(max_size=16, concurrency_count=1)
 demo.launch(server_name="0.0.0.0", server_port=7860, debug=False)
-