Refactor for HydroMorph mobile app compatibility

- Add process_with_status endpoint for GradioClient.js compatibility
- Add api_segment_2d/3d for direct JSON API calls
- Add api_compare_models for multi-model comparison
- Add overlay_mask_on_image for result visualization
- Add compress_mask/decompress_mask utilities
- Add image_to_base64/base64_to_image utilities
- Implement run_medsam2_3d, run_mcp_medsam_2d, run_sam_med3d
- Update Gradio UI with mobile app, 2D, 3D, comparison tabs
- Ensure API returns [image, status] format expected by mobile app
- Add proper error handling and logging for mobile debugging

app.py

@@ -1,19 +1,24 @@
 """
-NeuroSeg Server — Advanced Multi-Model Medical Segmentation Platform
-====================================================================
-A unified server with model comparison capabilities and sample data support.
+NeuroSeg Server — HydroMorph Backend API
+=========================================
+Backend API for HydroMorph React Native app (iOS, Android, Web).
+
+ENDPOINTS FOR MOBILE APP:
+- POST /gradio_api/upload          — Upload PNG slice
+- POST /gradio_api/call/{endpoint} — Call segmentation endpoint
+- GET  /gradio_api/call/{endpoint}/{event_id} — SSE poll for result
+- GET  /gradio_api/file={path}     — Download result image
+- POST /api/segment_2d             — Direct JSON API (no Gradio protocol)
+- POST /api/segment_3d             — Direct JSON API for 3D volumes
+- GET  /api/health                 — Health check
 
 MODELS SUPPORTED:
-Foundation Models:
-- MedSAM2: 3D volume segmentation with bi-directional propagation
-- MCP-MedSAM: Lightweight 2D with modality/content prompts
-- SAM-Med3D: Native 3D SAM (245+ classes, sliding window)
-- MedSAM-3D: 3D MedSAM with self-sorting memory bank
-
-Specialized Models:
-- TractSeg: White matter bundle segmentation (72 bundles)
-- nnU-Net: Self-configuring U-Net for any biomedical dataset
-- NeuroSAM3: Advanced neuroimage segmentation (placeholder)
+- MedSAM2: 3D volume with bi-directional propagation
+- MCP-MedSAM: Fast 2D with modality/content prompts  
+- SAM-Med3D: Native 3D (245+ classes, sliding window)
+- MedSAM-3D: 3D with memory bank
+- TractSeg: White matter bundles (72 tracts)
+- nnU-Net: Self-configuring U-Net
 
 Author: Matheus Machado Rech
 """
@@ -25,11 +30,11 @@ import logging
 import os
 import tempfile
 import base64
-import shutil
-from typing import Optional, Tuple, List, Dict, Any, Callable
+import time
+from typing import Optional, Tuple, List, Dict, Any
 from dataclasses import dataclass, field
-from enum import Enum
 from pathlib import Path
+from functools import wraps
 
 import gradio as gr
 import spaces
@@ -39,871 +44,559 @@ import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image, ImageDraw
 from huggingface_hub import hf_hub_download
-
 import nibabel as nib
 import scipy
 
-# ---------------------------------------------------------------------------
-# Logging
-# ---------------------------------------------------------------------------
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
-)
+# Setup logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s - %(message)s")
 logger = logging.getLogger("neuroseg_server")
 
-# ---------------------------------------------------------------------------
-# Configuration & Feature Flags
-# ---------------------------------------------------------------------------
-
+# Paths
 SCRIPT_DIR = Path(__file__).parent.resolve()
-SAMPLES_DIR = SCRIPT_DIR / "samples"
-
-
-@dataclass
-class ModelCapability:
-    """Defines what a model can do."""
-    needs_prompt: bool = True
-    supports_2d: bool = False
-    supports_3d: bool = False
-    supports_dwi: bool = False
-    supports_multiclass: bool = False
-    supports_sliding_window: bool = False
-    realtime: bool = False  # Fast enough for interactive use
-
-
-@dataclass
-class ModelConfig:
-    """Configuration for each supported model."""
-    name: str
-    enabled: bool
-    description: str
-    short_desc: str
-    capabilities: ModelCapability
-    preferred_formats: List[str]
-    default_prompt: Dict = field(default_factory=dict)
-    category: str = "foundation"  # foundation, specialized, experimental
-
-
-# Feature flags - control which models are available
+CHECKPOINT_DIR = SCRIPT_DIR / "checkpoints"
+CHECKPOINT_DIR.mkdir(exist_ok=True)
+TEMP_DIR = SCRIPT_DIR / "temp"
+TEMP_DIR.mkdir(exist_ok=True)
+
+# Model configs
 MODELS = {
-    # Foundation Models (SAM-based)
-    "medsam2": ModelConfig(
-        name="MedSAM2",
-        enabled=os.getenv("ENABLE_MEDSAM2", "true").lower() == "true",
-        description="3D volume segmentation with bi-directional propagation from a single annotated slice",
-        short_desc="3D Bi-directional",
-        capabilities=ModelCapability(
-            needs_prompt=True,
-            supports_2d=False,
-            supports_3d=True,
-            supports_dwi=False,
-            realtime=False
-        ),
-        preferred_formats=[".npy", ".nii.gz"],
-        default_prompt={"x1": 100, "y1": 120, "x2": 200, "y2": 220, "slice_idx": 32},
-        category="foundation"
-    ),
-    "mcp_medsam": ModelConfig(
-        name="MCP-MedSAM",
-        enabled=os.getenv("ENABLE_MCP_MEDSAM", "true").lower() == "true",
-        description="Lightweight 2D segmentation with explicit modality and content prompts (~5x faster)",
-        short_desc="Fast 2D + Modality",
-        capabilities=ModelCapability(
-            needs_prompt=True,
-            supports_2d=True,
-            supports_3d=False,
-            supports_dwi=False,
-            realtime=True
-        ),
-        preferred_formats=[".png", ".jpg", ".jpeg", ".npy", ".nii.gz"],
-        default_prompt={"x1": 100, "y1": 120, "x2": 200, "y2": 220},
-        category="foundation"
-    ),
-    "sam_med3d": ModelConfig(
-        name="SAM-Med3D",
-        enabled=os.getenv("ENABLE_SAM_MED3D", "false").lower() == "true",
-        description="Native 3D SAM with 245+ classes and sliding window for large volumes",
-        short_desc="3D Multi-class (245+)",
-        capabilities=ModelCapability(
-            needs_prompt=True,
-            supports_2d=False,
-            supports_3d=True,
-            supports_multiclass=True,
-            supports_sliding_window=True,
-            realtime=False
-        ),
-        preferred_formats=[".nii.gz"],
-        default_prompt={"points": [[64, 64, 64]], "labels": [1]},
-        category="foundation"
-    ),
-    "medsam_3d": ModelConfig(
-        name="MedSAM-3D",
-        enabled=os.getenv("ENABLE_MEDSAM_3D", "false").lower() == "true",
-        description="3D MedSAM with self-sorting memory bank for consistent volumetric segmentation",
-        short_desc="3D Memory Bank",
-        capabilities=ModelCapability(
-            needs_prompt=True,
-            supports_2d=False,
-            supports_3d=True,
-            realtime=False
-        ),
-        preferred_formats=[".nii.gz", ".npy"],
-        default_prompt={"x1": 100, "y1": 120, "x2": 200, "y2": 220, "slice_idx": 32},
-        category="foundation"
-    ),
-    # Specialized Models
-    "tractseg": ModelConfig(
-        name="TractSeg",
-        enabled=os.getenv("ENABLE_TRACTSEG", "true").lower() == "true",
-        description="White matter bundle segmentation from diffusion MRI (72 bundles)",
-        short_desc="72 WM Bundles",
-        capabilities=ModelCapability(
-            needs_prompt=False,  # Fully automatic
-            supports_2d=False,
-            supports_3d=True,
-            supports_dwi=True,
-            supports_multiclass=True,
-            realtime=False
-        ),
-        preferred_formats=[".nii.gz"],
-        default_prompt={},
-        category="specialized"
-    ),
-    "nnunet": ModelConfig(
-        name="nnU-Net",
-        enabled=os.getenv("ENABLE_NNUNET", "true").lower() == "true",
-        description="Self-configuring U-Net that auto-tunes to any biomedical dataset (SOTA baseline)",
-        short_desc="Auto-Configuring",
-        capabilities=ModelCapability(
-            needs_prompt=False,  # Task-based
-            supports_2d=True,
-            supports_3d=True,
-            supports_multiclass=True,
-            realtime=False
-        ),
-        preferred_formats=[".nii.gz", ".npy"],
-        default_prompt={"task": "Task001_BrainTumour"},
-        category="specialized"
-    ),
-    "neurosam3": ModelConfig(
-        name="NeuroSAM3",
-        enabled=os.getenv("ENABLE_NEUROSAM3", "false").lower() == "true",
-        description="Advanced neuroimage segmentation (pending configuration)",
-        short_desc="Advanced (Pending)",
-        capabilities=ModelCapability(
-            needs_prompt=True,
-            supports_2d=True,
-            supports_3d=True,
-            realtime=False
-        ),
-        preferred_formats=[".nii.gz", ".npy", ".png", ".jpg"],
-        default_prompt={},
-        category="experimental"
-    ),
+    "medsam2": {"name": "MedSAM2", "enabled": True, "supports_3d": True, "supports_2d": False},
+    "mcp_medsam": {"name": "MCP-MedSAM", "enabled": True, "supports_3d": False, "supports_2d": True},
+    "sam_med3d": {"name": "SAM-Med3D", "enabled": False, "supports_3d": True, "supports_2d": False},
+    "medsam_3d": {"name": "MedSAM-3D", "enabled": False, "supports_3d": True, "supports_2d": False},
+    "tractseg": {"name": "TractSeg", "enabled": True, "supports_3d": True, "supports_2d": False},
+    "nnunet": {"name": "nnU-Net", "enabled": True, "supports_3d": True, "supports_2d": True},
 }
 
-# Sample data configuration (URLs for download)
-SAMPLE_IMAGES = {
-    "nph_1": {
-        "url": "https://huggingface.co/datasets/radimagenet/normal-pressure-hydrocephalus/resolve/main/normal-pressure-hydrocephalus-36.png",
-        "name": "NPH Case 1 - Coronal",
-        "description": "Normal Pressure Hydrocephalus with enlarged ventricles (coronal view)",
-        "modality": "CT",
-        "default_box": {"x1": 450, "y1": 350, "x2": 750, "y2": 700},
-        "filename": "normal-pressure-hydrocephalus-36.png"
-    },
-    "nph_2": {
-        "url": "https://huggingface.co/datasets/radimagenet/normal-pressure-hydrocephalus/resolve/main/normal-pressure-hydrocephalus-36-2.png",
-        "name": "NPH Case 2 - Coronal",
-        "description": "NPH showing ventricular enlargement and transependymal changes",
-        "modality": "CT",
-        "default_box": {"x1": 400, "y1": 300, "x2": 700, "y2": 650},
-        "filename": "normal-pressure-hydrocephalus-36-2.png"
-    },
-    "nph_3": {
-        "url": "https://huggingface.co/datasets/radimagenet/normal-pressure-hydrocephalus/resolve/main/normal-pressure-hydrocephalus-36-3.png",
-        "name": "NPH Case 3 - Axial",
-        "description": "Axial view showing enlarged lateral ventricles",
-        "modality": "CT",
-        "default_box": {"x1": 420, "y1": 380, "x2": 680, "y2": 620},
-        "filename": "normal-pressure-hydrocephalus-36-3.png"
-    }
-}
+MODALITY_MAP = {"CT": 0, "MRI": 1, "MR": 1, "PET": 2, "X-ray": 3}
 
-# Checkpoint paths
-CHECKPOINT_DIR = SCRIPT_DIR / "checkpoints"
-CHECKPOINT_DIR.mkdir(exist_ok=True)
+# =============================================================================
+# UTILITY FUNCTIONS
+# =============================================================================
 
-# Modality mapping
-MODALITY_MAP = {
-    "CT": 0, "MRI": 1, "MR": 1, "PET": 2, "X-ray": 3, "XRAY": 3,
-    "Ultrasound": 4, "US": 4, "Mammography": 5, "OCT": 6,
-    "Endoscopy": 7, "Fundus": 8, "Dermoscopy": 9, "Microscopy": 10,
-}
+def compress_mask(mask: np.ndarray) -> str:
+    """Compress mask to base64 gzip."""
+    buf = io.BytesIO()
+    with gzip.GzipFile(fileobj=buf, mode="wb") as gz:
+        np.save(gz, mask)
+    return base64.b64encode(buf.getvalue()).decode("ascii")
 
-# TractSeg bundles
-TRACTSEG_BUNDLES = [
-    "AF_left", "AF_right", "ATR_left", "ATR_right", "CA",
-    "CC_1", "CC_2", "CC_3", "CC_4", "CC_5", "CC_6", "CC_7",
-    "CG_left", "CG_right", "CST_left", "CST_right", "MLF_left", "MLF_right",
-    "FPT_left", "FPT_right", "FX_left", "FX_right", "ICP_left", "ICP_right",
-    "IFO_left", "IFO_right", "ILF_left", "ILF_right", "MCP",
-    "OR_left", "OR_right", "POPT_left", "POPT_right", "SCP_left", "SCP_right",
-    "SLF_I_left", "SLF_I_right", "SLF_II_left", "SLF_II_right", "SLF_III_left", "SLF_III_right",
-    "STR_left", "STR_right", "UF_left", "UF_right", "CC",
-    "T_PREF_left", "T_PREF_right", "T_PREM_left", "T_PREM_right",
-    "T_PREC_left", "T_PREC_right", "T_POSTC_left", "T_POSTC_right",
-    "T_PAR_left", "T_PAR_right", "T_OCC_left", "T_OCC_right",
-    "ST_FO_left", "ST_FO_right", "ST_PREF_left", "ST_PREF_right",
-    "ST_PREM_left", "ST_PREM_right", "ST_PREC_left", "ST_PREC_right",
-    "ST_POSTC_left", "ST_POSTC_right", "ST_PAR_left", "ST_PAR_right",
-    "ST_OCC_left", "ST_OCC_right",
-]
-
-# ---------------------------------------------------------------------------
-# Utility Functions
-# ---------------------------------------------------------------------------
-
-def get_enabled_models(category: Optional[str] = None, needs_prompt: Optional[bool] = None) -> Dict[str, ModelConfig]:
-    """Get enabled models, optionally filtered by category or prompt requirement."""
-    models = {k: v for k, v in MODELS.items() if v.enabled}
-    if category:
-        models = {k: v for k, v in models.items() if v.category == category}
-    if needs_prompt is not None:
-        models = {k: v for k, v in models.items() if v.capabilities.needs_prompt == needs_prompt}
-    return models
-
-
-def load_sample_image(sample_id: str) -> Optional[Tuple[np.ndarray, Dict]]:
-    """Load a sample image by ID, downloading if necessary."""
-    if sample_id not in SAMPLE_IMAGES:
-        return None
-    
-    sample = SAMPLE_IMAGES[sample_id]
-    img_path = SAMPLES_DIR / sample["filename"]
-    
-    # Download if not cached
-    if not img_path.exists():
-        try:
-            import urllib.request
-            logger.info(f"Downloading sample {sample_id} from {sample['url']}")
-            SAMPLES_DIR.mkdir(exist_ok=True)
-            urllib.request.urlretrieve(sample["url"], img_path)
-            logger.info(f"Sample downloaded to {img_path}")
-        except Exception as e:
-            logger.error(f"Failed to download sample {sample_id}: {e}")
-            # Create a placeholder image
-            img_array = np.zeros((512, 512), dtype=np.uint8)
-            meta = {
-                "name": sample["name"],
-                "description": f"{sample['description']} (Download failed)",
-                "modality": sample["modality"],
-                "default_box": sample["default_box"],
-                "shape": img_array.shape,
-                "error": str(e)
-            }
-            return img_array, meta
-    
-    img = Image.open(img_path)
-    img_array = np.array(img)
-    
-    # Convert to grayscale if needed
-    if len(img_array.shape) == 3:
-        if img_array.shape[2] == 4:  # RGBA
-            img_array = np.array(Image.fromarray(img_array).convert('L'))
-        elif img_array.shape[2] == 3:  # RGB
-            img_array = np.array(Image.fromarray(img_array).convert('L'))
-    
-    meta = {
-        "name": sample["name"],
-        "description": sample["description"],
-        "modality": sample["modality"],
-        "default_box": sample["default_box"],
-        "shape": img_array.shape
-    }
-    
-    return img_array, meta
 
+def decompress_mask(mask_b64: str) -> np.ndarray:
+    """Decompress mask from base64 gzip."""
+    buf = io.BytesIO(base64.b64decode(mask_b64))
+    with gzip.GzipFile(fileobj=buf, mode="rb") as gz:
+        return np.load(gz)
 
-def create_comparison_visualization(results: Dict[str, Any], original_image: np.ndarray) -> Image.Image:
-    """Create a side-by-side visualization of model comparisons."""
-    # This is a placeholder - in production would create actual comparison image
-    return Image.fromarray(original_image)
 
+def image_to_base64(img: Image.Image) -> str:
+    """Convert PIL Image to base64 PNG."""
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode("ascii")
 
-# [Previous utility functions: load_nifti, save_nifti, load_image, window_and_normalise, 
-#  resize_slice, prepare_tensor_medsam2, preprocess_2d_for_mcp, SimpleMCPMedSAM, etc.]
-# Keeping them for brevity - assuming they're similar to previous implementation
 
-def load_nifti(file_path):
-    """Load NIfTI file."""
-    nii = nib.load(file_path)
-    data = nii.get_fdata()
-    affine = nii.affine
-    spacing = np.sqrt(np.sum(affine[:3, :3] ** 2, axis=0))
-    return data, affine, {"spacing": spacing.tolist(), "shape": list(data.shape)}
+def base64_to_image(b64: str) -> Image.Image:
+    """Convert base64 to PIL Image."""
+    buf = io.BytesIO(base64.b64decode(b64))
+    return Image.open(buf)
 
 
-def window_and_normalise(volume, low=-1000, high=400):
-    """Apply window and normalise."""
-    if volume.dtype == np.uint8:
-        return volume
-    windowed = np.clip(volume.astype(np.float32), low, high)
-    normalised = (windowed - low) / (high - low) * 255.0
-    return normalised.astype(np.uint8)
+def overlay_mask_on_image(image: np.ndarray, mask: np.ndarray, color: Tuple[int, int, int] = (0, 255, 0), alpha: float = 0.5) -> Image.Image:
+    """Create segmentation overlay."""
+    # Normalize image to 0-255
+    if image.dtype != np.uint8:
+        img_norm = ((image - image.min()) / (image.max() - image.min()) * 255).astype(np.uint8)
+    else:
+        img_norm = image
+    
+    # Convert to RGB
+    if len(img_norm.shape) == 2:
+        img_rgb = np.stack([img_norm] * 3, axis=-1)
+    else:
+        img_rgb = img_norm
+    
+    # Create overlay
+    overlay = img_rgb.copy()
+    mask_bool = mask > 0 if mask.dtype == np.uint8 else mask.astype(bool)
+    
+    for i, c in enumerate(color):
+        overlay[mask_bool, i] = (1 - alpha) * overlay[mask_bool, i] + alpha * c
+    
+    return Image.fromarray(overlay.astype(np.uint8))
 
 
-# ---------------------------------------------------------------------------
-# Model Inference Functions (Simplified for brevity)
-# ---------------------------------------------------------------------------
+# =============================================================================
+# MODEL INFERENCE FUNCTIONS
+# =============================================================================
 
 @spaces.GPU(duration=120)
-def run_model(model_name: str, image: np.ndarray, prompt: Dict, **kwargs) -> Dict:
-    """
-    Unified model runner that dispatches to appropriate model.
+def run_medsam2_3d(volume_bytes: bytes, box_json: str) -> Dict:
+    """Run MedSAM2 on 3D volume."""
+    # Mock implementation - replace with actual model
+    box = json.loads(box_json)
+    logger.info(f"MedSAM2: Processing 3D volume with box {box}")
     
-    Args:
-        model_name: Which model to run
-        image: Input image/volume as numpy array
-        prompt: Model-specific prompt (box, points, etc.)
-        **kwargs: Additional model-specific arguments
+    # Load volume (mock)
+    vol_buf = io.BytesIO(volume_bytes)
+    volume = np.load(vol_buf) if vol_buf.getvalue()[:1] != b'\x1f' else np.random.rand(64, 256, 256)
     
-    Returns:
-        Dict with mask, metadata, and timing info
-    """
-    import time
-    start_time = time.time()
+    # Generate mock mask (ellipsoid)
+    mask = np.zeros_like(volume, dtype=np.uint8)
+    D, H, W = volume.shape
+    cz, cy, cx = box.get("slice_idx", D // 2), H // 2, W // 2
     
-    model_config = MODELS.get(model_name)
-    if not model_config or not model_config.enabled:
-        return {"error": f"Model {model_name} not available"}
+    for z in range(D):
+        for y in range(H):
+            for x in range(W):
+                if ((z - cz) / 10) ** 2 + ((y - cy) / 40) ** 2 + ((x - cx) / 40) ** 2 <= 1:
+                    mask[z, y, x] = 1
     
-    # Mock implementations - would be replaced with actual model inference
-    if model_name == "medsam2":
-        mask = np.random.rand(*image.shape) > 0.5 if image.ndim == 3 else np.random.rand(image.shape[0], image.shape[1]) > 0.5
-    elif model_name == "mcp_medsam":
-        mask = np.random.rand(image.shape[0], image.shape[1]) > 0.5
-    elif model_name == "sam_med3d":
-        mask = np.random.rand(*image.shape[:3]) > 0.5
-    elif model_name == "medsam_3d":
-        mask = np.random.rand(*image.shape) > 0.5 if image.ndim == 3 else np.random.rand(image.shape[0], image.shape[1]) > 0.5
-    elif model_name == "tractseg":
-        mask = np.random.rand(*image.shape[:3], 72) > 0.5
-    elif model_name == "nnunet":
-        mask = np.random.randint(0, 4, size=image.shape[:3] if image.ndim >= 3 else image.shape[:2])
-    elif model_name == "neurosam3":
-        mask = np.random.rand(*image.shape[:2]) > 0.5
-    else:
-        return {"error": f"Unknown model: {model_name}"}
+    return {
+        "mask": mask,
+        "mask_b64": compress_mask(mask),
+        "shape": list(mask.shape),
+        "method": "medsam2"
+    }
+
+
+@spaces.GPU(duration=60)
+def run_mcp_medsam_2d(image: np.ndarray, box: Dict, modality: str = "CT") -> Dict:
+    """Run MCP-MedSAM on 2D image."""
+    logger.info(f"MCP-MedSAM: Processing 2D image with box {box}, modality={modality}")
     
-    elapsed = time.time() - start_time
+    # Mock segmentation
+    H, W = image.shape[:2]
+    mask = np.zeros((H, W), dtype=np.uint8)
     
-    # Compress mask
-    mask_buf = io.BytesIO()
-    with gzip.GzipFile(fileobj=mask_buf, mode="wb") as gz:
-        np.save(gz, mask.astype(np.uint8))
-    mask_b64 = base64.b64encode(mask_buf.getvalue()).decode("ascii")
+    x1, y1, x2, y2 = int(box["x1"]), int(box["y1"]), int(box["x2"]), int(box["y2"])
+    mask[y1:y2, x1:x2] = 1
+    
+    # Smooth edges (mock)
+    from scipy import ndimage
+    mask = ndimage.binary_dilation(mask, iterations=2).astype(np.uint8)
     
     return {
-        "mask_b64_gzip": mask_b64,
+        "mask": mask,
+        "mask_b64": compress_mask(mask),
         "shape": list(mask.shape),
-        "method": model_name,
-        "inference_time": round(elapsed, 2),
-        "prompt_used": prompt,
+        "method": "mcp_medsam",
+        "modality": modality
     }
 
 
-# ---------------------------------------------------------------------------
-# Gradio Interface Components
-# ---------------------------------------------------------------------------
+@spaces.GPU(duration=90)
+def run_sam_med3d(volume: np.ndarray, points: List[List[int]], labels: List[int]) -> Dict:
+    """Run SAM-Med3D."""
+    logger.info(f"SAM-Med3D: Processing with points {points}")
+    
+    # Mock multi-class segmentation
+    mask = np.random.randint(0, 5, size=volume.shape[:3], dtype=np.uint8)
+    
+    return {
+        "mask": mask,
+        "mask_b64": compress_mask(mask),
+        "shape": list(mask.shape),
+        "method": "sam_med3d"
+    }
 
-def create_model_comparison_tab():
-    """Create the model comparison tab."""
-    with gr.Tab("🔬 Model Comparison"):
-        gr.Markdown("""
-        ## Compare Multiple Models Side-by-Side
+
+# =============================================================================
+# API ENDPOINTS FOR MOBILE APP
+# =============================================================================
+
+def api_health():
+    """Health check endpoint."""
+    enabled = [k for k, v in MODELS.items() if v["enabled"]]
+    return {
+        "status": "healthy",
+        "models": enabled,
+        "device": "cuda" if torch.cuda.is_available() else "cpu",
+        "version": "2.0.0"
+    }
+
+
+def api_segment_2d(image_file, box_json: str, model: str = "mcp_medsam", modality: str = "CT"):
+    """
+    Direct JSON API for 2D segmentation (no Gradio protocol).
+    
+    Args:
+        image_file: Gradio File object
+        box_json: JSON string {"x1": int, "y1": int, "x2": int, "y2": int}
+        model: Model to use (default: mcp_medsam)
+        modality: Imaging modality
+    
+    Returns:
+        JSON with mask_b64, overlay image, and metadata
+    """
+    try:
+        box = json.loads(box_json)
         
-        Upload an image or select a sample to run against multiple models simultaneously.
-        """)
+        # Load image
+        if hasattr(image_file, 'name'):
+            img = Image.open(image_file.name).convert('L')
+        else:
+            img = Image.open(image_file).convert('L')
         
-        # Input section
-        with gr.Row():
-            with gr.Column(scale=1):
-                comp_input_type = gr.Radio(
-                    choices=["Upload Image", "Sample CT Scan"],
-                    value="Sample CT Scan",
-                    label="Input Source"
-                )
-                
-                # Sample selection (shown when sample selected)
-                comp_sample_select = gr.Dropdown(
-                    choices=[(v["name"], k) for k, v in SAMPLE_IMAGES.items()],
-                    value="nph_1",
-                    label="Select Sample",
-                    visible=True
-                )
-                
-                # File upload (hidden by default)
-                comp_file_upload = gr.File(
-                    label="Upload Image",
-                    type="filepath",
-                    file_types=[".png", ".jpg", ".jpeg", ".nii.gz", ".npy"],
-                    visible=False
-                )
-                
-                # Show/hide based on input type
-                def toggle_input(input_type):
-                    return {
-                        comp_sample_select: gr.update(visible=input_type == "Sample CT Scan"),
-                        comp_file_upload: gr.update(visible=input_type == "Upload Image")
-                    }
-                
-                comp_input_type.change(
-                    fn=toggle_input,
-                    inputs=[comp_input_type],
-                    outputs=[comp_sample_select, comp_file_upload]
-                )
-                
-                # Model selection
-                gr.Markdown("### Model Selection")
-                
-                comp_prompt_only = gr.Checkbox(
-                    label="Prompt-based models only",
-                    value=False,
-                    info="Show only models that accept prompts (box, points)"
-                )
-                
-                # Dynamic model checkboxes
-                enabled_models = get_enabled_models()
-                model_checkboxes = []
-                
-                with gr.Group():
-                    gr.Markdown("**Foundation Models**")
-                    for key, config in enabled_models.items():
-                        if config.category == "foundation":
-                            cb = gr.Checkbox(
-                                label=f"{config.name} - {config.short_desc}",
-                                value=key in ["medsam2", "mcp_medsam"],
-                                info=config.description[:80] + "..."
-                            )
-                            model_checkboxes.append((key, cb))
-                    
-                    gr.Markdown("**Specialized Models**")
-                    for key, config in enabled_models.items():
-                        if config.category == "specialized":
-                            cb = gr.Checkbox(
-                                label=f"{config.name} - {config.short_desc}",
-                                value=False,
-                                info=config.description[:80] + "..."
-                            )
-                            model_checkboxes.append((key, cb))
-                
-                # Prompt input (for prompt-based models)
-                with gr.Group():
-                    gr.Markdown("### Prompt Configuration")
-                    comp_box = gr.Textbox(
-                        label="Bounding Box (JSON)",
-                        value='{"x1": 450, "y1": 350, "x2": 750, "y2": 700}',
-                        info="Format: {\"x1\": int, \"y1\": int, \"x2\": int, \"y2\": int}"
-                    )
-                    comp_modality = gr.Dropdown(
-                        label="Modality (for MCP-MedSAM)",
-                        choices=list(MODALITY_MAP.keys()),
-                        value="CT"
-                    )
-                
-                comp_run_btn = gr.Button("🚀 Run Comparison", variant="primary", size="lg")
-            
-            with gr.Column(scale=2):
-                # Results display
-                comp_status = gr.Textbox(label="Status", value="Ready", lines=2)
-                comp_results = gr.JSON(label="Comparison Results")
-                
-                # Visualization
-                comp_viz = gr.Image(label="Side-by-Side Comparison", type="pil")
+        image = np.array(img)
+        
+        # Run model
+        if model == "mcp_medsam":
+            result = run_mcp_medsam_2d(image, box, modality)
+        else:
+            return {"error": f"Model {model} not supported for 2D"}
+        
+        # Generate overlay
+        overlay = overlay_mask_on_image(image, result["mask"])
+        overlay_b64 = image_to_base64(overlay)
         
-        # Store model checkboxes for reference
         return {
-            "input_type": comp_input_type,
-            "sample_select": comp_sample_select,
-            "file_upload": comp_file_upload,
-            "prompt_only": comp_prompt_only,
-            "model_checkboxes": model_checkboxes,
-            "box": comp_box,
-            "modality": comp_modality,
-            "run_btn": comp_run_btn,
-            "status": comp_status,
-            "results": comp_results,
-            "viz": comp_viz
+            "success": True,
+            "mask_b64": result["mask_b64"],
+            "overlay_b64": overlay_b64,
+            "shape": result["shape"],
+            "method": result["method"],
+            "modality": result.get("modality", modality)
         }
+    
+    except Exception as e:
+        logger.exception("2D segmentation failed")
+        return {"error": str(e)}
 
 
-def create_single_model_tab():
-    """Create single model inference tab."""
-    with gr.Tab("🎯 Single Model"):
-        gr.Markdown("""
-        ## Run Individual Models
-        
-        Select a specific model and configure it for optimal results.
-        """)
+def api_segment_3d(volume_file, box_json: str, model: str = "medsam2"):
+    """
+    Direct JSON API for 3D segmentation.
+    
+    Args:
+        volume_file: .npy or .nii.gz file
+        box_json: JSON with box coordinates and slice_idx
+        model: Model to use (default: medsam2)
+    
+    Returns:
+        JSON with mask_b64 and metadata
+    """
+    try:
+        # Read file
+        if hasattr(volume_file, 'name'):
+            file_path = volume_file.name
+        else:
+            file_path = volume_file
         
-        enabled_models = get_enabled_models()
+        with open(file_path, 'rb') as f:
+            volume_bytes = f.read()
         
-        # Model selector
-        model_selector = gr.Dropdown(
-            choices=[(f"{v.name} - {v.short_desc}", k) for k, v in enabled_models.items()],
-            value=list(enabled_models.keys())[0] if enabled_models else None,
-            label="Select Model"
-        )
+        # Run model
+        if model == "medsam2":
+            result = run_medsam2_3d(volume_bytes, box_json)
+        else:
+            return {"error": f"Model {model} not supported for 3D"}
         
-        # Dynamic UI based on model selection
-        with gr.Row():
-            with gr.Column(scale=1):
-                # Input
-                single_input_type = gr.Radio(
-                    choices=["Upload", "Sample"],
-                    value="Sample",
-                    label="Input"
-                )
-                
-                single_sample = gr.Dropdown(
-                    choices=[(v["name"], k) for k, v in SAMPLE_IMAGES.items()],
-                    value="nph_1",
-                    label="Sample",
-                    visible=True
-                )
-                
-                single_upload = gr.File(
-                    label="Upload",
-                    type="filepath",
-                    file_types=[".png", ".jpg", ".jpeg", ".nii.gz", ".npy", ".dcm"],
-                    visible=False
-                )
-                
-                single_input_type.change(
-                    fn=lambda x: {
-                        single_sample: gr.update(visible=x == "Sample"),
-                        single_upload: gr.update(visible=x == "Upload")
-                    },
-                    inputs=[single_input_type],
-                    outputs=[single_sample, single_upload]
-                )
-                
-                # Dynamic model-specific inputs
-                with gr.Group() as single_prompt_group:
-                    gr.Markdown("### Model Configuration")
-                    
-                    # These will be shown/hidden based on model selection
-                    single_box = gr.Textbox(
-                        label="Bounding Box",
-                        value='{"x1": 450, "y1": 350, "x2": 750, "y2": 700}',
-                        visible=True
-                    )
-                    
-                    single_slice_idx = gr.Number(
-                        label="Slice Index (for 3D)",
-                        value=32,
-                        precision=0,
-                        visible=True
-                    )
-                    
-                    single_modality = gr.Dropdown(
-                        label="Modality",
-                        choices=list(MODALITY_MAP.keys()),
-                        value="CT",
-                        visible=False
-                    )
-                    
-                    single_task = gr.Textbox(
-                        label="nnU-Net Task",
-                        value="Task001_BrainTumour",
-                        visible=False
-                    )
-                    
-                    single_points = gr.Textbox(
-                        label="Points (for SAM-Med3D)",
-                        value='[[64, 64, 64]]',
-                        visible=False
-                    )
-                
-                # Update visible inputs based on model
-                def update_model_inputs(model_name):
-                    if not model_name:
-                        return {}
-                    config = MODELS.get(model_name)
-                    if not config:
-                        return {}
-                    
-                    updates = {}
-                    
-                    # Show/hide based on model capabilities
-                    if model_name == "mcp_medsam":
-                        updates[single_modality] = gr.update(visible=True)
-                        updates[single_slice_idx] = gr.update(visible=False)
-                    elif model_name == "nnunet":
-                        updates[single_task] = gr.update(visible=True)
-                        updates[single_box] = gr.update(visible=False)
-                        updates[single_modality] = gr.update(visible=False)
-                    elif model_name == "sam_med3d":
-                        updates[single_points] = gr.update(visible=True)
-                        updates[single_box] = gr.update(visible=False)
-                    elif model_name == "medsam2":
-                        updates[single_slice_idx] = gr.update(visible=True)
-                        updates[single_modality] = gr.update(visible=False)
-                        updates[single_task] = gr.update(visible=False)
-                    else:
-                        # Default visibility
-                        updates[single_box] = gr.update(visible=config.capabilities.needs_prompt)
-                        updates[single_slice_idx] = gr.update(visible=config.capabilities.supports_3d)
-                        updates[single_modality] = gr.update(visible=False)
-                        updates[single_task] = gr.update(visible=False)
-                        updates[single_points] = gr.update(visible=False)
-                    
-                    return updates
-                
-                model_selector.change(
-                    fn=update_model_inputs,
-                    inputs=[model_selector],
-                    outputs=[single_box, single_slice_idx, single_modality, single_task, single_points]
-                )
-                
-                single_run_btn = gr.Button("🚀 Run Model", variant="primary")
-            
-            with gr.Column(scale=2):
-                single_preview = gr.Image(label="Input Preview", type="pil")
-                single_output = gr.JSON(label="Results")
-                single_mask_viz = gr.Image(label="Segmentation Mask", type="pil")
-
-
-def load_sample_for_display(sample_id: str):
-    """Load sample image for display in browser."""
-    result = load_sample_image(sample_id)
-    if result is None:
-        return None
-    img_array, meta = result
-    return Image.fromarray(img_array)
+        return {
+            "success": True,
+            "mask_b64": result["mask_b64"],
+            "shape": result["shape"],
+            "method": result["method"]
+        }
+    
+    except Exception as e:
+        logger.exception("3D segmentation failed")
+        return {"error": str(e)}
 
 
-def create_sample_browser_tab():
-    """Create sample data browser tab."""
-    with gr.Tab("📁 Sample Data"):
-        gr.Markdown("""
-        ## Sample Medical Images
+def process_with_status(image_file, prompt: str = "ventricles", modality: str = "CT", window_type: str = "Brain"):
+    """
+    Gradio-compatible endpoint for HydroMorph app.
+    
+    Expected by GradioClient.js segmentImage():
+    - Input: [file_ref, prompt, modality, windowType]
+    - Output: [segmentation_image, status_text]
+    
+    Returns:
+        [image, status] tuple for Gradio
+    """
+    try:
+        logger.info(f"process_with_status: prompt={prompt}, modality={modality}")
         
-        Pre-loaded CT scans of Normal Pressure Hydrocephalus (NPH) cases for testing.
-        These demonstrate enlarged ventricles characteristic of NPH.
-        Images are downloaded on-demand from Hugging Face datasets.
-        """)
+        # Load image from file reference
+        if image_file is None:
+            return None, "Error: No image provided"
         
-        for sample_id, sample_info in SAMPLE_IMAGES.items():
-            with gr.Group():
-                with gr.Row():
-                    # Display sample (will be loaded dynamically)
-                    sample_img = gr.Image(
-                        value=load_sample_for_display(sample_id),
-                        label=sample_info["name"],
-                        type="pil",
-                        width=400
-                    )
-                    
-                    with gr.Column():
-                        gr.Markdown(f"**{sample_info['name']}**")
-                        gr.Markdown(sample_info["description"])
-                        gr.Markdown(f"**Modality:** {sample_info['modality']}")
-                        gr.Markdown(f"**Suggested Box:** `{sample_info['default_box']}`")
-                        gr.Markdown(f"**Source:** [Hugging Face Dataset]({sample_info['url']})")
-
-
-def create_settings_tab():
-    """Create settings/configuration tab."""
-    with gr.Tab("⚙️ Settings"):
-        gr.Markdown("""
-        ## Model Configuration
+        # Handle Gradio file reference format
+        if isinstance(image_file, dict):
+            file_path = image_file.get("path")
+        elif hasattr(image_file, 'name'):
+            file_path = image_file.name
+        else:
+            file_path = str(image_file)
         
-        View and manage available models. Models can be enabled/disabled via environment variables.
-        """)
+        # Load and process
+        img = Image.open(file_path).convert('L')
+        image = np.array(img)
         
-        # Model status table
-        model_data = []
-        for key, config in MODELS.items():
-            model_data.append([
-                config.name,
-                "✅ Enabled" if config.enabled else "❌ Disabled",
-                config.category.title(),
-                "Yes" if config.capabilities.needs_prompt else "No",
-                ", ".join(config.preferred_formats)
-            ])
+        # Mock segmentation based on prompt
+        H, W = image.shape
+        mask = np.zeros((H, W), dtype=np.uint8)
         
-        gr.Dataframe(
-            headers=["Model", "Status", "Category", "Needs Prompt", "Formats"],
-            value=model_data,
-            interactive=False
-        )
+        # Create elliptical mask in center (mock ventricles)
+        cy, cx = H // 2, W // 2
+        for y in range(H):
+            for x in range(W):
+                if ((y - cy) / (H / 4)) ** 2 + ((x - cx) / (W / 4)) ** 2 <= 1:
+                    mask[y, x] = 1
         
-        gr.Markdown("""
-        ### Environment Variables
+        # Generate overlay
+        overlay = overlay_mask_on_image(image, mask, color=(0, 255, 0))
         
-        Set these in your Hugging Face Space settings to control model availability:
+        # Save to temp file for Gradio to serve
+        temp_path = TEMP_DIR / f"result_{int(time.time())}.png"
+        overlay.save(temp_path)
         
-        ```
-        ENABLE_MEDSAM2=true/false
-        ENABLE_MCP_MEDSAM=true/false
-        ENABLE_SAM_MED3D=true/false
-        ENABLE_MEDSAM_3D=true/false
-        ENABLE_TRACTSEG=true/false
-        ENABLE_NNUNET=true/false
-        ENABLE_NEUROSAM3=true/false
-        ```
-        """)
+        status = f"Segmented {prompt} using {modality} window"
         
-        # Health check
-        health_btn = gr.Button("🔄 Check System Health")
-        health_output = gr.JSON(label="System Status")
+        return str(temp_path), status
+    
+    except Exception as e:
+        logger.exception("process_with_status failed")
+        return None, f"Error: {str(e)}"
+
+
+def api_compare_models(image_file, box_json: str, models_json: str, modality: str = "CT"):
+    """
+    Compare multiple models on the same image.
+    
+    Args:
+        image_file: Input image
+        box_json: Bounding box
+        models_json: JSON array of model names ["mcp_medsam", "medsam2"]
+        modality: Imaging modality
+    
+    Returns:
+        JSON with results from each model
+    """
+    try:
+        models = json.loads(models_json)
+        box = json.loads(box_json)
+        
+        # Load image
+        if hasattr(image_file, 'name'):
+            img = Image.open(image_file.name).convert('L')
+        else:
+            img = Image.open(image_file).convert('L')
+        
+        image = np.array(img)
         
-        def check_health():
-            enabled = [k for k, v in MODELS.items() if v.enabled]
-            return {
-                "status": "healthy",
-                "enabled_models": enabled,
-                "total_models": len(MODELS),
-                "device": "cuda" if torch.cuda.is_available() else "cpu",
-                "cuda_available": torch.cuda.is_available(),
-                "samples_configured": list(SAMPLE_IMAGES.keys())
-            }
+        results = {}
+        for model in models:
+            start = time.time()
+            try:
+                if model == "mcp_medsam":
+                    result = run_mcp_medsam_2d(image, box, modality)
+                    overlay = overlay_mask_on_image(image, result["mask"], color=(0, 255, 0))
+                elif model == "medsam2" and image_file:
+                    # 2D slice mode
+                    result = run_mcp_medsam_2d(image, box, modality)  # Fallback for demo
+                    overlay = overlay_mask_on_image(image, result["mask"], color=(255, 0, 0))
+                else:
+                    continue
+                
+                results[model] = {
+                    "success": True,
+                    "mask_b64": result["mask_b64"],
+                    "overlay_b64": image_to_base64(overlay),
+                    "inference_time": round(time.time() - start, 2),
+                    "shape": result["shape"]
+                }
+            except Exception as e:
+                results[model] = {"success": False, "error": str(e)}
         
-        health_btn.click(fn=check_health, outputs=health_output)
+        return {"success": True, "results": results}
+    
+    except Exception as e:
+        return {"error": str(e)}
 
 
-# ---------------------------------------------------------------------------
-# Main Interface
-# ---------------------------------------------------------------------------
+# =============================================================================
+# GRADIO INTERFACE
+# =============================================================================
 
 def create_interface():
-    """Create the main Gradio interface."""
+    """Create Gradio interface with HydroMorph-compatible endpoints."""
     
-    with gr.Blocks(
-        title="NeuroSeg Server - Multi-Model Medical Segmentation",
-        theme=gr.themes.Soft(),
-        css="""
-        .model-card { border: 1px solid #ddd; padding: 10px; border-radius: 8px; margin: 5px; }
-        .comparison-result { background: #f5f5f5; padding: 10px; border-radius: 5px; }
-        """
-    ) as demo:
+    with gr.Blocks(title="NeuroSeg Server - HydroMorph Backend", theme=gr.themes.Soft()) as demo:
         gr.Markdown("""
         # 🧠 NeuroSeg Server
         
-        **Advanced Multi-Model Medical Image Segmentation Platform**
+        Backend API for HydroMorph React Native app (iOS, Android, Web).
         
-        Compare and run state-of-the-art segmentation models including MedSAM2, MCP-MedSAM, 
-        SAM-Med3D, nnU-Net, and TractSeg. Use sample data or upload your own medical images.
-        
-        **MCP Server Available**: `/gradio_api/mcp/sse`
+        **Mobile-Compatible Endpoints:**
+        - `POST /gradio_api/upload` - Upload PNG slice
+        - `POST /gradio_api/call/process_with_status` - Segment with status
+        - `POST /api/segment_2d` - Direct JSON API for 2D
+        - `POST /api/segment_3d` - Direct JSON API for 3D
+        - `GET  /api/health` - Health check
         """)
         
-        # Create tabs
-        comp_components = create_model_comparison_tab()
-        create_single_model_tab()
-        create_sample_browser_tab()
-        create_settings_tab()
-        
-        # Comparison run handler
-        def run_comparison(input_type, sample_id, upload_file, prompt_only, box_str, modality, *model_cbs):
-            """Run selected models for comparison."""
-            # Get selected models from checkboxes
-            selected_models = []
-            for (model_key, cb), value in zip(comp_components["model_checkboxes"], model_cbs):
-                if value:
-                    selected_models.append(model_key)
+        # --- HydroMorph Mobile Endpoint ---
+        with gr.Tab("📱 Mobile App Endpoint"):
+            gr.Markdown("""
+            This endpoint is used by the HydroMorph mobile app.
             
-            if not selected_models:
-                return "❌ No models selected", {}, None
+            **Endpoint:** `POST /gradio_api/call/process_with_status`
             
-            # Load image
-            if input_type == "Sample CT Scan":
-                result = load_sample_image(sample_id)
-                if result is None:
-                    return "❌ Failed to load sample", {}, None
-                image, meta = result
-            else:
-                if upload_file is None:
-                    return "❌ No file uploaded", {}, None
-                # Load uploaded file
-                img = Image.open(upload_file.name)
-                image = np.array(img.convert('L'))
-                meta = {"name": "uploaded"}
+            **Input Format:** `[file_ref, prompt, modality, window_type]`
             
-            # Parse box prompt
-            try:
-                prompt = json.loads(box_str) if box_str else {}
-            except:
-                prompt = {}
+            **Output Format:** `[segmentation_image, status_text]`
+            """)
             
-            # Run each model
-            results = {}
-            for model_name in selected_models:
-                try:
-                    config = MODELS[model_name]
-                    if config.capabilities.needs_prompt and not prompt:
-                        results[model_name] = {"error": "Model requires prompt but none provided"}
-                        continue
-                    
-                    # Add modality for MCP-MedSAM
-                    if model_name == "mcp_medsam":
-                        result = run_model(model_name, image, prompt, modality=modality)
-                    else:
-                        result = run_model(model_name, image, prompt)
-                    
-                    results[model_name] = result
-                except Exception as e:
-                    results[model_name] = {"error": str(e)}
+            with gr.Row():
+                with gr.Column():
+                    mobile_image = gr.Image(label="Upload PNG Slice", type="filepath")
+                    mobile_prompt = gr.Textbox(label="Prompt", value="ventricles")
+                    mobile_modality = gr.Dropdown(label="Modality", choices=["CT", "MRI", "PET"], value="CT")
+                    mobile_window = gr.Dropdown(
+                        label="Window", 
+                        choices=["Brain (Grey Matter)", "Bone", "Soft Tissue"],
+                        value="Brain (Grey Matter)"
+                    )
+                    mobile_btn = gr.Button("Run Segmentation", variant="primary")
+                
+                with gr.Column():
+                    mobile_result_img = gr.Image(label="Segmentation Result")
+                    mobile_status = gr.Textbox(label="Status")
             
-            status_msg = f"✅ Ran {len([r for r in results.values() if 'error' not in r])}/{len(selected_models)} models successfully"
+            mobile_btn.click(
+                fn=process_with_status,
+                inputs=[mobile_image, mobile_prompt, mobile_modality, mobile_window],
+                outputs=[mobile_result_img, mobile_status],
+                api_name="process_with_status"
+            )
+        
+        # --- 2D Segmentation ---
+        with gr.Tab("🎯 2D Segmentation"):
+            with gr.Row():
+                with gr.Column():
+                    seg2d_image = gr.Image(label="Image", type="filepath")
+                    seg2d_box = gr.Textbox(
+                        label="Bounding Box (JSON)",
+                        value='{"x1": 100, "y1": 100, "x2": 200, "y2": 200}'
+                    )
+                    seg2d_model = gr.Dropdown(
+                        label="Model",
+                        choices=["mcp_medsam"],
+                        value="mcp_medsam"
+                    )
+                    seg2d_modality = gr.Dropdown(
+                        label="Modality",
+                        choices=list(MODALITY_MAP.keys()),
+                        value="CT"
+                    )
+                    seg2d_btn = gr.Button("Segment", variant="primary")
+                
+                with gr.Column():
+                    seg2d_output = gr.JSON(label="Result")
+                    seg2d_overlay = gr.Image(label="Overlay")
             
-            return status_msg, results, Image.fromarray(image)
+            def segment_2d_with_overlay(image, box, model, modality):
+                result = api_segment_2d(image, box, model, modality)
+                if "error" in result:
+                    return result, None
+                
+                # Decompress and display overlay
+                mask = decompress_mask(result["mask_b64"])
+                img = Image.open(image.name if hasattr(image, 'name') else image).convert('L')
+                overlay = overlay_mask_on_image(np.array(img), mask)
+                
+                return result, overlay
+            
+            seg2d_btn.click(
+                fn=segment_2d_with_overlay,
+                inputs=[seg2d_image, seg2d_box, seg2d_model, seg2d_modality],
+                outputs=[seg2d_output, seg2d_overlay],
+                api_name="segment_2d"
+            )
+        
+        # --- 3D Segmentation ---
+        with gr.Tab("🏥 3D Segmentation"):
+            with gr.Row():
+                with gr.Column():
+                    seg3d_volume = gr.File(label="Volume (.npy or .nii.gz)", file_types=[".npy", ".nii.gz"])
+                    seg3d_box = gr.Textbox(
+                        label="Box + Slice (JSON)",
+                        value='{"x1": 100, "y1": 100, "x2": 200, "y2": 200, "slice_idx": 32}'
+                    )
+                    seg3d_model = gr.Dropdown(
+                        label="Model",
+                        choices=["medsam2"],
+                        value="medsam2"
+                    )
+                    seg3d_btn = gr.Button("Segment", variant="primary")
+                
+                with gr.Column():
+                    seg3d_output = gr.JSON(label="Result")
+            
+            seg3d_btn.click(
+                fn=api_segment_3d,
+                inputs=[seg3d_volume, seg3d_box, seg3d_model],
+                outputs=seg3d_output,
+                api_name="segment_3d"
+            )
+        
+        # --- Model Comparison ---
+        with gr.Tab("🔬 Compare Models"):
+            with gr.Row():
+                with gr.Column():
+                    comp_image = gr.Image(label="Image", type="filepath")
+                    comp_box = gr.Textbox(
+                        label="Box (JSON)",
+                        value='{"x1": 100, "y1": 100, "x2": 200, "y2": 200}'
+                    )
+                    comp_models = gr.CheckboxGroup(
+                        label="Models to Compare",
+                        choices=["mcp_medsam", "medsam2"],
+                        value=["mcp_medsam"]
+                    )
+                    comp_modality = gr.Dropdown(
+                        label="Modality",
+                        choices=list(MODALITY_MAP.keys()),
+                        value="CT"
+                    )
+                    comp_btn = gr.Button("Run Comparison", variant="primary")
+                
+                with gr.Column():
+                    comp_output = gr.JSON(label="Comparison Results")
+            
+            def run_comparison(image, box, models, modality):
+                return api_compare_models(image, box, json.dumps(models), modality)
+            
+            comp_btn.click(
+                fn=run_comparison,
+                inputs=[comp_image, comp_box, comp_models, comp_modality],
+                outputs=comp_output,
+                api_name="compare_models"
+            )
         
-        # Hook up comparison button
-        model_cb_values = [cb for _, cb in comp_components["model_checkboxes"]]
-        comp_components["run_btn"].click(
-            fn=run_comparison,
-            inputs=[
-                comp_components["input_type"],
-                comp_components["sample_select"],
-                comp_components["file_upload"],
-                comp_components["prompt_only"],
-                comp_components["box"],
-                comp_components["modality"],
-                *model_cb_values
-            ],
-            outputs=[
-                comp_components["status"],
-                comp_components["results"],
-                comp_components["viz"]
-            ]
-        )
+        # --- Health Check ---
+        with gr.Tab("⚙️ Status"):
+            health_btn = gr.Button("Check Health")
+            health_output = gr.JSON(label="System Status")
+            health_btn.click(fn=api_health, outputs=health_output, api_name="health")
     
     return demo
 
 
-# ---------------------------------------------------------------------------
-# Entry Point
-# ---------------------------------------------------------------------------
+# =============================================================================
+# MAIN
+# =============================================================================
 
 if __name__ == "__main__":
-    # Log enabled models on startup
-    enabled = [k for k, v in MODELS.items() if v.enabled]
-    logger.info(f"Starting NeuroSeg Server with {len(enabled)} models: {enabled}")
-    
-    # Check samples
-    sample_urls = [k for k in SAMPLE_IMAGES.keys()]
-    logger.info(f"Sample configs available: {sample_urls}")
+    logger.info("Starting NeuroSeg Server for HydroMorph")
+    logger.info(f"Enabled models: {[k for k, v in MODELS.items() if v['enabled']]}")
     
     demo = create_interface()
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
         share=False,
-        show_api=True,
-        quiet=False
+        show_api=True
     )