update space

app.py

@@ -14,26 +14,24 @@ deploying to Hugging Face Spaces).
 
 from __future__ import annotations
 
-import os
 import random
-from functools import lru_cache
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple
 
 import cv2
 import gradio as gr
 import numpy as np
 import pandas as pd
 import torch
-from datasets import Dataset, DatasetDict, IterableDataset, load_dataset
+from datasets import Dataset
 from PIL import Image
-from transformers import (
-    AutoImageProcessor,
-    AutoModelForImageClassification,
-)
-
 
-HF_REPO_ID = "raidium/curia"
-HF_DATASET_ID = "raidium/CuriaBench"
+from inference import (
+    HF_DATASET_ID,
+    load_curia_dataset,
+    load_id_to_labels,
+    to_numpy_image,
+    infer_image,
+)
 
 
 # ---------------------------------------------------------------------------
@@ -41,6 +39,7 @@ HF_DATASET_ID = "raidium/CuriaBench"
 # ---------------------------------------------------------------------------
 
 HEAD_OPTIONS: List[Tuple[str, str]] = [
+    ("abdominal-trauma", "Active Extravasation"),
     ("anatomy-ct", "Anatomy CT"),
     ("anatomy-mri", "Anatomy MRI"),
     ("atlas-stroke", "Atlas Stroke"),
@@ -53,6 +52,8 @@ HEAD_OPTIONS: List[Tuple[str, str]] = [
     ("kneeMRI", "Knee MRI"),
     ("luna16-3D", "LUNA16 3D"),
     ("neural_foraminal_narrowing", "Neural Foraminal Narrowing"),
+    ("spinal_canal_stenosis", "Spinal Canal Stenosis"),
+    ("subarticular_stenosis", "Subarticular Stenosis"),
     ("oasis", "OASIS"),
 ]
 
@@ -65,6 +66,7 @@ HEADS_REQUIRING_MASK: set[str] = {
     "kits",
     "kneeMRI",
     "luna16-3D",
+    "neural_foraminal_narrowing",
     "spinal_canal_stenosis",
     "subarticular_stenosis",
 }
@@ -115,60 +117,6 @@ DEFAULT_WINDOWINGS: Dict[str, Optional[Dict[str, int]]] = {
 # ---------------------------------------------------------------------------
 
 
-@lru_cache(maxsize=1)
-def load_processor() -> AutoImageProcessor:
-    token = os.environ.get("HF_TOKEN")
-    return AutoImageProcessor.from_pretrained(HF_REPO_ID, trust_remote_code=True, token=token)
-
-
-@lru_cache(maxsize=len(HEAD_OPTIONS))
-def load_model(head: str) -> AutoModelForImageClassification:
-    token = os.environ.get("HF_TOKEN")
-    model = AutoModelForImageClassification.from_pretrained(
-        HF_REPO_ID,
-        trust_remote_code=True,
-        subfolder=head,
-        token=token,
-    )
-    model.eval()
-    return model
-
-
-@lru_cache(maxsize=len(DATASET_OPTIONS))
-def load_curia_dataset(subset: str) -> Any:
-    token = os.environ.get("HF_TOKEN")
-    ds = load_dataset(
-        HF_DATASET_ID,
-        subset,
-        split="test",
-        token=token,
-    )
-    if isinstance(ds, DatasetDict):
-        return ds["test"]
-    return ds
-
-
-def to_numpy_image(image: Any) -> np.ndarray:
-    """Convert dataset or user-provided imagery to a float32 numpy array."""
-
-    if isinstance(image, np.ndarray):
-        arr = image
-    elif isinstance(image, Image.Image):
-        arr = np.array(image)
-    else:
-        # Some datasets provide nested dicts or lists – attempt to coerce.
-        arr = np.array(image)
-
-    if arr.ndim == 3 and arr.shape[-1] == 3:
-        # Convert RGB to grayscale by averaging channels
-        arr = arr.mean(axis=-1)
-
-    if arr.ndim != 2:
-        raise ValueError("Expected a 2D image (H, W). Please provide a single axial/coronal/sagittal slice.")
-
-    return arr.astype(np.float32)
-
-
 def apply_windowing(image: np.ndarray, subset: str) -> np.ndarray:
     """Apply CT windowing based on the dataset.
 
@@ -223,26 +171,8 @@ def to_display_image(image: np.ndarray) -> np.ndarray:
     return arr
 
 
-def coerce_mask_array(mask: Any) -> Optional[np.ndarray]:
-    if mask is None:
-        return None
-
-    try:
-        arr = np.array(mask)
-    except Exception:
-        return None
-
-    if arr.size == 0:
-        return None
-    return arr
-
-
-def prepare_mask_tensor(mask: Any, height: int, width: int) -> Optional[torch.Tensor]:
-    mask_array = coerce_mask_array(mask)
-    if mask_array is None:
-        return None
-
-    arr = np.squeeze(mask_array)
+def prepare_mask_tensor(mask: np.ndarray, height: int, width: int) -> Optional[torch.Tensor]:
+    arr = np.squeeze(mask)
     if arr.ndim == 2:
         arr = arr.reshape(1, height, width)
     else:
@@ -301,16 +231,17 @@ def apply_contour_overlay(
     return output
 
 
-def render_image_with_mask_info(image: np.ndarray, mask: Any) -> Tuple[np.ndarray, Optional[str]]:
+def render_image_with_mask_info(image: np.ndarray, mask: Any) -> np.ndarray:
     display = to_display_image(image)
     if mask is None:
-        return display, None
+        return display
 
     try:
         overlaid = apply_contour_overlay(display, mask)
-        return overlaid, ""
+        return overlaid
     except Exception:
-        return display, "Mask provided but could not be visualised."
+        gr.Warning("Mask provided but could not be visualised.")
+        return display
 
 
 def dataset_class_metadata(dataset: Dataset) -> Tuple[List[int], Dict[int, str]]:
@@ -341,39 +272,6 @@ def pick_random_indices(dataset: Dataset, target: Optional[int]) -> int:
     return random.choice(indices)
 
 
-def format_probabilities(probs: torch.Tensor, id2label: Dict[int, str]) -> pd.DataFrame:
-    """Return a dataframe sorted by probability desc."""
-
-    values = probs.detach().cpu().numpy()
-    rows = [
-        {"class_id": idx, "label": id2label.get(idx, str(idx)), "probability": float(val)}
-        for idx, val in enumerate(values)
-    ]
-    df = pd.DataFrame(rows)
-    df.sort_values("probability", ascending=False, inplace=True)
-    return df
-
-
-def infer_image(
-    image: np.ndarray,
-    head: str,
-) -> Tuple[str, pd.DataFrame]:
-    processor = load_processor()
-    model = load_model(head)
-    with torch.no_grad():
-        processed = processor(images=image, return_tensors="pt")
-        outputs = model(**processed)
-        print(outputs)
-        logits = outputs["logits"]
-        probs = torch.nn.functional.softmax(logits[0], dim=-1)
-
-    id2label = model.config.id2label or {}
-    df = format_probabilities(probs, id2label)
-    top_row = df.iloc[0]
-    prediction = f"{top_row['label']} (p={top_row['probability']:.3f})"
-    return prediction, df
-
-
 # ---------------------------------------------------------------------------
 # Gradio callbacks
 # ---------------------------------------------------------------------------
@@ -439,12 +337,12 @@ def parse_target_selection(selection: str) -> Optional[int]:
 def sample_dataset_example(
     subset: str,
     target_id: Optional[int],
-) -> Tuple[np.ndarray, str, Dict[str, Any]]:
+) -> Tuple[np.ndarray, Dict[str, Any]]:
     dataset = load_curia_dataset(subset)
     index = pick_random_indices(dataset, target_id)
     record = dataset[index]
     image = to_numpy_image(record["image"])
-    mask_array = coerce_mask_array(record.get("mask"))
+    mask_array = record.get("mask")
 
     meta = {
         "index": index,
@@ -452,7 +350,7 @@ def sample_dataset_example(
         "mask": mask_array,
     }
 
-    return image, f"Sample #{index}", meta
+    return image, meta
 
 
 def load_dataset_sample(
@@ -461,57 +359,69 @@ def load_dataset_sample(
     head: str,
 ) -> Tuple[
     Optional[np.ndarray],
-    str,
     pd.DataFrame,
     Dict[str, Any],
     Optional[Dict[str, Any]],
 ]:
     try:
         target_id = parse_target_selection(target_selection)
-        image, caption, meta = sample_dataset_example(subset, target_id)
+        image, meta = sample_dataset_example(subset, target_id)
 
         # Apply windowing only for display, keep raw image for model inference
         windowed_image = apply_windowing(image, subset)
-        display, mask_msg = render_image_with_mask_info(windowed_image, meta.get("mask"))
+        display = to_display_image(windowed_image)
+        if meta.get("mask") is not None:
+            display = apply_contour_overlay(display, meta.get("mask"))
 
         target = meta.get("target")
-        meta_text = caption
-        if target is not None:
-            meta_text += f" | target={target}"
-        status = "Image loaded. Click 'Run inference' to compute predictions."
-        if mask_msg:
-            status += f" {mask_msg}"
-        meta_text = status + "\n\n" + meta_text
-
         # Generate ground truth display
         ground_truth_update = gr.update(visible=False)
         if target is not None:
-            model = load_model(head)
-            id2label = model.config.id2label or {}
+            # Use id_to_labels.json mapping
+            id2label = load_id_to_labels().get(head, {})
             label_name = id2label.get(target, str(target))
             ground_truth_update = gr.update(value=f"**Ground Truth:** {label_name} (class {target})", visible=True)
 
         return (
             display,
-            meta_text,
             pd.DataFrame(),
             ground_truth_update,
             {"image": image, "mask": meta.get("mask")},  # Store raw image for inference
         )
     except Exception as exc:  # pragma: no cover - surfaced in UI
-        return None, f"Failed to load sample: {exc}", pd.DataFrame(), gr.update(visible=False), None
+        gr.Warning(f"Failed to load sample: {exc}")
+        return None, pd.DataFrame(), gr.update(visible=False), None
+
+
+def format_probabilities(probs: torch.Tensor, id2label: Dict[int, str]) -> pd.DataFrame:
+    """Return a dataframe sorted by probability desc."""
+
+    values = probs.detach().cpu().numpy()
+    rows = [
+        {"class_id": idx, "label": id2label.get(idx, str(idx)), "probability": float(val)}
+        for idx, val in enumerate(values)
+    ]
+    df = pd.DataFrame(rows)
+    df.sort_values("probability", ascending=False, inplace=True)
+    return df
 
 
 def run_inference(
-    sample_state: Optional[Dict[str, Any]],
+    image_state: Optional[Dict[str, Any]],
     head: str,
 ) -> Tuple[str, pd.DataFrame]:
-    if not sample_state or "image" not in sample_state:
+    if not image_state or "image" not in image_state:
         return "Load a dataset sample or upload an image first.", pd.DataFrame()
 
     try:
-        image = sample_state["image"]
-        prediction, df = infer_image(image, head)
+        image = image_state["image"]
+        probs = infer_image(image, head)
+
+        # Use id_to_labels.json mapping, fall back to model config if not available
+        id2label = load_id_to_labels().get(head, {})
+        df = format_probabilities(probs, id2label)
+        top_row = df.iloc[0]
+        prediction = f"{top_row['label']} (p={top_row['probability']:.3f})"
         result_text = f"**Prediction:** {prediction}"
         return result_text, df
     except Exception as exc:  # pragma: no cover - surfaced in UI
@@ -605,7 +515,7 @@ def build_demo() -> gr.Blocks:
         gr.Markdown("---")
 
         infer_btn = gr.Button("Run inference", variant="primary")
-
+        status_text = gr.Markdown()
         with gr.Row():
             with gr.Column():
                 image_display = gr.Image(label="Image", interactive=False, type="numpy")
@@ -613,7 +523,7 @@ def build_demo() -> gr.Blocks:
 
             with gr.Column():
                 gr.Markdown("### Predictions")
-                status_text = gr.Markdown()
+                main_prediction = gr.Markdown()
                 prediction_probs = gr.Dataframe(headers=["class_id", "label", "probability"])
 
         image_state = gr.State()
@@ -640,7 +550,6 @@ def build_demo() -> gr.Blocks:
             inputs=[dataset_dropdown, class_dropdown, head_dropdown],
             outputs=[
                 image_display,
-                status_text,
                 prediction_probs,
                 ground_truth_display,
                 image_state,
@@ -662,7 +571,7 @@ def build_demo() -> gr.Blocks:
         infer_btn.click(
             fn=run_inference,
             inputs=[image_state, head_dropdown],
-            outputs=[status_text, prediction_probs],
+            outputs=[main_prediction, prediction_probs],
         )
 
         gr.Markdown(

id_to_labels.json

@@ -0,0 +1,173 @@
+{
+    "covidx-ct": {
+        "0": "normal",
+        "1": "pneumonia",
+        "2": "COVID-19"
+    },
+    "deep-lesion-site": {
+        "0": "abdomen",
+        "1": "bone",
+        "2": "kidney",
+        "3": "liver",
+        "4": "lung",
+        "5": "mediastinum",
+        "6": "pelvis",
+        "7": "soft_tissue"
+    },
+    "kits": {
+        "0": "lesion_kidney",
+        "1": "cyst"
+    },
+    "kneeMRI": {
+        "0": "healthy",
+        "1": "partially injured",
+        "2": "completely ruptured"
+    },
+    "luna16-3D": {
+        "0": "benign",
+        "1": "malignant"
+    },
+    "oasis": {
+        "0": "nondemented",
+        "1": "signs of dementia"
+    },
+    "abdominal-trauma": {
+        "0": "no active extravasation",
+        "1": "active extravasation"
+    },
+    "ich": {
+        "0": "no hemorrhage",
+        "1": "intracranial hemorrhage"
+    },
+    "neural_foraminal_narrowing": {
+        "0": "normal/mild",
+        "1": "moderate",
+        "2": "severe"
+    },
+    "spinal_canal_stenosis": {
+        "0": "normal/mild",
+        "1": "moderate",
+        "2": "severe"
+    },
+    "subarticular_stenosis": {
+        "0": "normal/mild",
+        "1": "moderate",
+        "2": "severe"
+    },
+    "anatomy-ct": {
+        "0": "adrenal_gland_left",
+        "1": "adrenal_gland_right",
+        "2": "aorta",
+        "3": "autochthon_left",
+        "4": "autochthon_right",
+        "5": "brain",
+        "6": "clavicula_left",
+        "7": "clavicula_right",
+        "8": "colon",
+        "9": "duodenum",
+        "10": "esophagus",
+        "11": "face",
+        "12": "femur_left",
+        "13": "femur_right",
+        "14": "gallbladder",
+        "15": "gluteus_maximus_left",
+        "16": "gluteus_maximus_right",
+        "17": "gluteus_medius_left",
+        "18": "gluteus_medius_right",
+        "19": "gluteus_minimus_left",
+        "20": "gluteus_minimus_right",
+        "21": "heart",
+        "22": "hip_left",
+        "23": "hip_right",
+        "24": "humerus_left",
+        "25": "humerus_right",
+        "26": "iliac_artery_left",
+        "27": "iliac_artery_right",
+        "28": "iliac_vena_left",
+        "29": "iliac_vena_right",
+        "30": "iliopsoas_left",
+        "31": "iliopsoas_right",
+        "32": "inferior_vena_cava",
+        "33": "kidney_left",
+        "34": "kidney_right",
+        "35": "liver",
+        "36": "lung_left",
+        "37": "lung_right",
+        "38": "pancreas",
+        "39": "portal_vein",
+        "40": "pulmonary_artery",
+        "41": "rib_left",
+        "42": "rib_right",
+        "43": "sacrum",
+        "44": "scapula_left",
+        "45": "scapula_right",
+        "46": "small_bowel",
+        "47": "spleen",
+        "48": "stomach",
+        "49": "trachea",
+        "50": "urinary_bladder",
+        "51": "vertebrae_cervical",
+        "52": "vertebrae_dorsal",
+        "53": "vertebrae_lumbar"
+    },
+    "anatomy-mri": {
+        "0": "adrenal_gland_left",
+        "1": "adrenal_gland_right",
+        "2": "aorta",
+        "3": "autochthon_left",
+        "4": "autochthon_right",
+        "5": "brain",
+        "6": "colon",
+        "7": "duodenum",
+        "8": "esophagus",
+        "9": "femur_left",
+        "10": "femur_right",
+        "11": "fibula",
+        "12": "gallbladder",
+        "13": "gluteus_maximus_left",
+        "14": "gluteus_maximus_right",
+        "15": "gluteus_medius_left",
+        "16": "gluteus_medius_right",
+        "17": "heart",
+        "18": "hip_left",
+        "19": "hip_right",
+        "20": "humerus_left",
+        "21": "humerus_right",
+        "22": "iliac_artery_left",
+        "23": "iliac_artery_right",
+        "24": "iliac_vena_left",
+        "25": "iliac_vena_right",
+        "26": "iliopsoas_left",
+        "27": "iliopsoas_right",
+        "28": "inferior_vena_cava",
+        "29": "intervertebral_discs",
+        "30": "kidney_left",
+        "31": "kidney_right",
+        "32": "liver",
+        "33": "lung_left",
+        "34": "lung_right",
+        "35": "pancreas",
+        "36": "portal_vein_and_splenic_vein",
+        "37": "prostate",
+        "38": "quadriceps_femoris_left",
+        "39": "quadriceps_femoris_right",
+        "40": "sacrum",
+        "41": "sartorius_left",
+        "42": "sartorius_right",
+        "43": "small_bowel",
+        "44": "spinal_cord",
+        "45": "spleen",
+        "46": "stomach",
+        "47": "thigh_medial_compartment_left",
+        "48": "thigh_medial_compartment_right",
+        "49": "thigh_posterior_compartment_left",
+        "50": "thigh_posterior_compartment_right",
+        "51": "tibia",
+        "52": "urinary_bladder",
+        "53": "vertebrae"
+    },
+    "emidec-classification-mask": {
+        "0": "healthy",
+        "1": "infarction"
+    }
+}

inference.py

@@ -0,0 +1,99 @@
+"""Model and dataset loading, inference, and label extraction functions."""
+
+from __future__ import annotations
+
+import json
+import os
+from functools import lru_cache
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+from datasets import Dataset, DatasetDict, load_dataset
+from PIL import Image
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+)
+
+
+HF_REPO_ID = "raidium/curia"
+HF_DATASET_ID = "raidium/CuriaBench"
+
+
+@lru_cache(maxsize=1)
+def load_id_to_labels() -> Dict[str, Dict[str, str]]:
+    """Load the id_to_labels.json mapping file."""
+    json_path = os.path.join(os.path.dirname(__file__), "id_to_labels.json")
+    with open(json_path, "r") as f:
+        data = json.load(f)
+        # convert string keys to integers
+        for head in data:
+            data[head] = {int(k): v for k, v in data[head].items()}
+        return data
+
+
+@lru_cache(maxsize=1)
+def load_processor() -> AutoImageProcessor:
+    token = os.environ.get("HF_TOKEN")
+    return AutoImageProcessor.from_pretrained(HF_REPO_ID, trust_remote_code=True, token=token)
+
+
+@lru_cache(maxsize=None)
+def load_model(head: str) -> AutoModelForImageClassification:
+    token = os.environ.get("HF_TOKEN")
+    model = AutoModelForImageClassification.from_pretrained(
+        HF_REPO_ID,
+        trust_remote_code=True,
+        subfolder=head,
+        token=token,
+    )
+    model.eval()
+    return model
+
+
+@lru_cache(maxsize=None)
+def load_curia_dataset(subset: str) -> Any:
+    token = os.environ.get("HF_TOKEN")
+    ds = load_dataset(
+        HF_DATASET_ID,
+        subset,
+        split="test",
+        token=token,
+    )
+    if isinstance(ds, DatasetDict):
+        return ds["test"]
+    return ds
+
+
+def to_numpy_image(image: Any) -> np.ndarray:
+    """Convert dataset or user-provided imagery to a float32 numpy array."""
+
+    if isinstance(image, np.ndarray):
+        arr = image
+    elif isinstance(image, Image.Image):
+        arr = np.array(image)
+    else:
+        # Some datasets provide nested dicts or lists – attempt to coerce.
+        arr = np.array(image)
+
+    if arr.ndim == 3 and arr.shape[-1] == 3:
+        # Convert RGB to grayscale by averaging channels
+        arr = arr.mean(axis=-1)
+
+    return arr.astype(np.float32)
+
+
+def infer_image(
+    image: np.ndarray,
+    head: str,
+) -> torch.Tensor:
+    processor = load_processor()
+    model = load_model(head)
+    with torch.no_grad():
+        processed = processor(images=image, return_tensors="pt")
+        outputs = model(**processed)
+        logits = outputs["logits"]
+        probs = torch.nn.functional.softmax(logits[0], dim=-1)
+    return probs

requirements.txt

@@ -6,4 +6,4 @@ pandas>=2.2.0
 numpy>=1.26.0
 pillow>=10.2.0
 opencv-python>=4.8.0
-
+torchvision