Update app.py

app.py

@@ -1,109 +1,339 @@
 import os
+import io
 import cv2
+import sys
+import math
+import json
 import torch
 import gradio as gr
 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 PIL import Image
-from huggingface_hub import snapshot_download
-import spaces  # Import spaces for ZeroGPU compatibility
-
-# Adjustable parameters for detection
-CONFIDENCE_THRESHOLD = 0.1  # Confidence threshold for detections
-NMS_THRESHOLD = 0         # IoU threshold for non-maximum suppression
-SLICE_WIDTH = 1024          # Width of each slice
-SLICE_HEIGHT = 1024         # Height of each slice
-OVERLAP_WIDTH = 0         # Overlap width between slices
-OVERLAP_HEIGHT = 0        # Overlap height between slices
-ANNOTATION_COLOR = sv.Color.RED  # Red in BGR format for OpenCV
-ANNOTATION_THICKNESS = 4     # Thickness of bounding box lines
-
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-# Download YOLO model weights from Hugging Face Hub
-repo_id = 'edeler/ICC'  # Replace with your Hugging Face repository ID
-model_dir = snapshot_download(repo_id, local_dir='./models/ICC')
-model_path = os.path.join(model_dir, "best.pt")  # Adjust if filename differs
-model = YOLO(model_path).to(device)
-
-# Define the detection function for Gradio
-@spaces.GPU  # Decorator to allocate GPU for ZeroGPU-enabled Spaces
-def detect_objects(image: np.ndarray) -> (Image.Image, str):
-    # Ensure the image is in BGR format if provided by PIL (Gradio gives us an RGB image)
-    if image.shape[-1] == 3:
-        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    
-    # Define callback function for slice-based inference
-    def callback(image_slice: np.ndarray) -> sv.Detections:
-        # Run inference on each slice
-        result = model(image_slice)[0]
-        # Convert detections to `sv.Detections` format for further processing
-        detections = sv.Detections.from_ultralytics(result)
-        # Filter detections based on confidence threshold
-        return detections[detections.confidence >= CONFIDENCE_THRESHOLD]
-
-    # Initialize InferenceSlicer with adjustable slice dimensions and overlap settings
+from huggingface_hub import hf_hub_download
+import spaces  # ZeroGPU-compatible decorator
+
+# -----------------------------
+# Defaults / configuration
+# -----------------------------
+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
+
+# Reasonable defaults (tunable in UI)
+DEFAULT_CONF = 0.25
+DEFAULT_IOU = 0.50               # for both model-level and global cross-slice NMS
+DEFAULT_SLICE_WH = 1024
+DEFAULT_OVERLAP = 128            # helps stitching at tile borders
+DEFAULT_THICKNESS = 3
+DEFAULT_LONG_EDGE = 4096         # optional downscale for very large WSIs/crops
+
+# -----------------------------
+# Torch / device helpers
+# -----------------------------
+def get_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+def half_precision_available(dev: str) -> bool:
+    return dev == "cuda"
+
+# -----------------------------
+# Model loading (lazy, cached)
+# -----------------------------
+_MODEL: Optional[YOLO] = None
+_CLASS_NAMES: Optional[Dict[int, str]] = None
+
+def load_model() -> Tuple[YOLO, Dict[int, str]]:
+    global _MODEL, _CLASS_NAMES
+    if _MODEL is not None and _CLASS_NAMES is not None:
+        return _MODEL, _CLASS_NAMES
+
+    os.makedirs(LOCAL_MODEL_DIR, exist_ok=True)
+    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
+    )
+    model = YOLO(weights_path)
+    class_names = model.model.names if hasattr(model, "model") else model.names
+    _MODEL, _CLASS_NAMES = model, class_names
+    return _MODEL, _CLASS_NAMES
+
+# -----------------------------
+# 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
+
+def maybe_downscale_long_edge(img: np.ndarray, target_long_edge: int) -> Tuple[np.ndarray, float]:
+    h, w = img.shape[:2]
+    long_edge = max(h, w)
+    if long_edge <= target_long_edge:
+        return img, 1.0
+    scale = target_long_edge / float(long_edge)
+    new_w = max(1, int(round(w * scale)))
+    new_h = max(1, int(round(h * scale)))
+    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
+    return resized, scale
+
+# -----------------------------
+# Inference (sliced + global NMS)
+# -----------------------------
+def run_sliced_inference(
+    image_bgr: np.ndarray,
+    model: YOLO,
+    conf: float,
+    iou: float,
+    slice_w: int,
+    slice_h: int,
+    overlap_w: int,
+    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,
+            iou=iou,
+            device=device,
+            verbose=False,
+            half=half_precision_available(device),
+        )
+        res = results[0]
+        det = sv.Detections.from_ultralytics(res)
+        return det
+
     slicer = sv.InferenceSlicer(
         callback=callback,
-        slice_wh=(SLICE_WIDTH, SLICE_HEIGHT),
-        overlap_wh=(OVERLAP_WIDTH, OVERLAP_HEIGHT),
-        overlap_ratio_wh=None
+        slice_wh=(slice_w, slice_h),
+        overlap_wh=(overlap_w, overlap_h),
+        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
 
-    # Perform slicing-based inference on the entire image
-    detections = slicer(image)
+# -----------------------------
+# Annotation + tabular export
+# -----------------------------
+def make_labels(det: sv.Detections, names: Dict[int, str], show_labels: bool) -> Optional[List[str]]:
+    if not show_labels:
+        return None
+    labels = []
+    for cls, conf in zip(det.class_id, det.confidence):
+        cname = names.get(int(cls), str(cls))
+        labels.append(f"{cname} {conf:.2f}")
+    return labels
 
-    # Apply Non-Maximum Suppression (NMS) to the detections to avoid duplicate boxes
-    detections = detections.with_nms(threshold=NMS_THRESHOLD, class_agnostic=False)
+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 = []
+    for i in range(len(det)):
+        cls = int(det.class_id[i])
+        data.append({
+            "class_id": cls,
+            "class_name": names.get(cls, str(cls)),
+            "confidence": float(det.confidence[i]),
+            "x_min": float(xyxy[i, 0]),
+            "y_min": float(xyxy[i, 1]),
+            "x_max": float(xyxy[i, 2]),
+            "y_max": float(xyxy[i, 3]),
+        })
+    return pd.DataFrame(data)
 
-    # Count total detections after NMS
-    total_detections = len(detections)
+def per_class_summary(df: pd.DataFrame) -> str:
+    if df.empty:
+        return "No detections."
+    counts = df.groupby("class_name").size().sort_values(ascending=False)
+    parts = [f"{k}: {v}" for k, v in counts.items()]
+    total = int(counts.sum())
+    return f"Total detections: {total} | " + " | ".join(parts)
 
-    # Initialize an annotator for bounding boxes with specified color and thickness
-    box_annotator = sv.OrientedBoxAnnotator(color=ANNOTATION_COLOR, thickness=ANNOTATION_THICKNESS)
+# -----------------------------
+# Gradio inference function
+# -----------------------------
+@spaces.GPU  # ensures ZeroGPU allocation during this call
+def detect_objects(
+    image: np.ndarray,
+    conf: float,
+    iou: float,
+    slice_w: int,
+    slice_h: int,
+    overlap_w: int,
+    overlap_h: int,
+    thickness: int,
+    show_labels: bool,
+    selected_classes: List[str],
+    long_edge: int,
+    progress=gr.Progress(track_tqdm=True),
+) -> tuple:
+    try:
+        if image is None:
+            raise ValueError("Please upload or select an image.")
 
-    # Annotate the image with bounding boxes after NMS
-    annotated_img = box_annotator.annotate(scene=image.copy(), detections=detections)
+        # Prepare image (BGR) and optional downscale
+        image_bgr = ensure_bgr(image)
+        image_bgr, scale = maybe_downscale_long_edge(image_bgr, long_edge)
 
-    # Convert annotated image to RGB for Gradio display (PIL expects RGB)
-    annotated_img_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
+        # Load model + names lazily
+        progress(0.05, desc="Loading model…")
+        model, names = load_model()
+        device = get_device()
 
-    # Return the annotated image and the total count of detections
-    return Image.fromarray(annotated_img_rgb), f"Total Detections: {total_detections}"
+        # Inference (sliced)
+        progress(0.35, desc="Running sliced inference…")
+        with torch.inference_mode():
+            detections = run_sliced_inference(
+                image_bgr=image_bgr,
+                model=model,
+                conf=conf,
+                iou=iou,
+                slice_w=slice_w,
+                slice_h=slice_h,
+                overlap_w=overlap_w,
+                overlap_h=overlap_h,
+                device=device,
+            )
 
-# Reset function for Gradio UI
-def gradio_reset():
-    return gr.update(value=None), gr.update(value=None), gr.update(value="")
+        # 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]
 
-# Set up Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("<h1>Interstitial Cell of Cajal Detection and Quantification Tool</h1>")
+        # Create labels (optional)
+        labels = make_labels(detections, names, show_labels)
 
-    # Define input image before using it in the Examples section
-    input_img = gr.Image(label="Upload an Image", type="numpy", interactive=True)
-    
-    # Move the Examples section here, right after the title
-    with gr.Accordion("Select an Example Image"):
-        example_root = os.path.dirname(__file__)  # Root directory
-        example_images = [os.path.join(example_root, file) for file in os.listdir(example_root) if file.endswith(".jpg")]
-        gr.Examples(
-            examples=example_images,
-            inputs=[input_img],
-        )
+        # 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)
+
+        progress(1.0, desc="Done.")
+        return annotated_pil, summary, df, csv_bytes
+
+    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""
+
+# -----------------------------
+# UI / App
+# -----------------------------
+def discover_examples(root: str) -> List[str]:
+    exts = {".jpg", ".jpeg", ".png"}
+    paths = []
+    try:
+        for fname in os.listdir(root):
+            if os.path.splitext(fname.lower())[1] in exts:
+                paths.append(os.path.join(root, fname))
+    except Exception:
+        pass
+    # Keep at most a handful to keep UI tidy
+    return sorted(paths)[: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:
+    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>")
 
     with gr.Row():
-        with gr.Column():
-            clear = gr.Button(value="Clear")
-            predict = gr.Button(value="Detect", variant="primary")
-            
-        with gr.Column():
+        with gr.Column(scale=1):
+            input_img = gr.Image(label="Upload an Image", type="numpy", interactive=True)
+
+            with gr.Accordion("Examples", open=False):
+                examples = discover_examples(EXAMPLES_DIR)
+                if examples:
+                    gr.Examples(examples=examples, inputs=[input_img])
+                else:
+                    gr.Markdown("_No example images found in repository root._")
+
+            with gr.Accordion("Advanced Settings", open=False):
+                conf = gr.Slider(0.0, 1.0, value=DEFAULT_CONF, step=0.01, label="Confidence threshold (conf)")
+                iou = gr.Slider(0.0, 1.0, value=DEFAULT_IOU, step=0.01, label="IoU threshold (NMS)")
+                slice_w = gr.Slider(256, 2048, value=DEFAULT_SLICE_WH, step=64, label="Slice width (px)")
+                slice_h = gr.Slider(256, 2048, value=DEFAULT_SLICE_WH, step=64, label="Slice height (px)")
+                overlap_w = gr.Slider(0, 512, value=DEFAULT_OVERLAP, step=16, label="Horizontal overlap (px)")
+                overlap_h = gr.Slider(0, 512, value=DEFAULT_OVERLAP, step=16, label="Vertical overlap (px)")
+                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())]
+                except Exception:
+                    class_list = []
+                selected_classes = gr.CheckboxGroup(choices=class_list, label="Filter by class (optional)")
+
+            with gr.Row():
+                predict = gr.Button(value="Detect", variant="primary")
+                clear = gr.Button(value="Clear")
+
+        with gr.Column(scale=1):
             output_img = gr.Image(label="Detection Result", interactive=False)
-            detection_count = gr.Textbox(label="Detection Summary", interactive=False)
+            detection_summary = gr.Textbox(label="Detection Summary", interactive=False)
+            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")
+
+    # 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],
+    )
 
-    # Define button actions
-    clear.click(gradio_reset, inputs=None, outputs=[input_img, output_img, detection_count])
-    predict.click(detect_objects, inputs=[input_img], outputs=[output_img, detection_count])
+# 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)
 
-# Launch Gradio app
-demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)