Upload app.py with huggingface_hub

app.py

@@ -1,14 +1,7 @@
 #!/usr/bin/env python3
 """
 OncoSeg Inference API - HuggingFace Space
-Optimized for programmatic access from oncoseg-viewer
-
-This Space provides GPU-accelerated inference for medical image segmentation.
-It exposes both a Gradio UI and programmatic API endpoints.
-
-Usage from viewer:
-    POST /api/segment_slice
-    POST /api/segment_volume
+Minimal version for initial deployment.
 """
 
 import os
@@ -18,12 +11,9 @@ import tempfile
 import time
 import logging
 from pathlib import Path
-from typing import Optional, List, Tuple, Any
 
 import gradio as gr
 import numpy as np
-import torch
-import cv2
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -37,79 +27,12 @@ try:
     logger.info("ZeroGPU available")
 except ImportError:
     ZEROGPU_AVAILABLE = False
-    logger.info("ZeroGPU not available, using standard GPU/CPU")
-
-# Device setup
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-logger.info(f"Using device: {DEVICE}")
-
-# Global model cache
-MODELS = {}
-
-# Checkpoint mapping (HuggingFace Hub paths)
-CHECKPOINTS = {
-    "brain": "checkpoints/medsam3-task20_brats_gli-final_latest/last.ckpt",
-    "liver": "checkpoints/medsam3-task03_liver-final_latest/last.ckpt",
-    "breast": "checkpoints/medsam3-task25_breastdcedl-final_latest/last.ckpt",
-    "lung": "checkpoints/medsam3-task06_lung-final_latest/last.ckpt",
-    "kidney": "checkpoints/medsam3-task17_kits23-final_latest/last.ckpt",
-    "spine": "checkpoints/medsam3-task11_lctsc-final_latest/last.ckpt",
-}
-
-# HF Repo ID for checkpoints
-HF_REPO_ID = os.getenv("HF_REPO_ID", "tp53/oncoseg")
-
-
-# Flag to track if we're using fallback mode
-USE_FALLBACK = False
-
-
-def get_model(checkpoint: str = "brain"):
-    """Load or retrieve cached model. Falls back to simple segmentation if SAM3 unavailable."""
-    global MODELS, USE_FALLBACK
-
-    if checkpoint not in MODELS:
-        logger.info(f"Loading model: {checkpoint}")
-
-        try:
-            from huggingface_hub import hf_hub_download
+    logger.info("ZeroGPU not available")
 
-            ckpt_file = CHECKPOINTS.get(checkpoint, CHECKPOINTS["brain"])
-            ckpt_path = hf_hub_download(
-                repo_id=HF_REPO_ID,
-                filename=ckpt_file,
-            )
-            logger.info(f"Downloaded checkpoint to: {ckpt_path}")
 
-            # Import model (from local model/ directory in this Space)
-            from model.medsam3 import MedSAM3Model
-
-            # Initialize model with checkpoint
-            # Note: MedSAM3Model builds SAM3 internally and loads our LoRA weights
-            model = MedSAM3Model(checkpoint_path=ckpt_path)
-            model.to(DEVICE)
-            model.eval()
-
-            MODELS[checkpoint] = model
-            logger.info(f"Model {checkpoint} loaded on {DEVICE}")
-
-        except ImportError as e:
-            logger.warning(f"SAM3 not available, using fallback segmentation: {e}")
-            USE_FALLBACK = True
-            MODELS[checkpoint] = None
-        except Exception as e:
-            logger.error(f"Failed to load model {checkpoint}: {e}")
-            USE_FALLBACK = True
-            MODELS[checkpoint] = None
-
-    return MODELS.get(checkpoint)
-
-
-def fallback_segment(slice_2d: np.ndarray):
-    """
-    Simple intensity-based segmentation fallback when SAM3 is not available.
-    Works well for FLAIR MRI where tumors appear hyperintense.
-    """
+def fallback_segment(slice_2d):
+    """Simple intensity-based segmentation."""
+    from scipy import ndimage
     from skimage.filters import threshold_otsu
     from skimage.morphology import binary_opening, binary_closing, disk
 
@@ -126,51 +49,15 @@ def fallback_segment(slice_2d: np.ndarray):
 
     # Morphological cleanup
     try:
-        mask = binary_opening(mask, disk(2))
-        mask = binary_closing(mask, disk(3))
-    except:
+        mask = binary_opening(mask, disk(2)).astype(np.uint8)
+        mask = binary_closing(mask, disk(3)).astype(np.uint8)
+    except Exception:
         pass
 
-    return mask.astype(np.uint8)
-
-
-def preprocess_slice(slice_2d: np.ndarray, target_size: int = 1024) -> torch.Tensor:
-    """
-    Preprocess a 2D slice for SAM3 input.
-
-    Args:
-        slice_2d: Input slice (H, W)
-        target_size: Target size for SAM3 (default 1024)
-
-    Returns:
-        Preprocessed tensor (1, 3, H, W) on DEVICE
-    """
-    import cv2
-
-    # Normalize to [0, 1]
-    vmin, vmax = slice_2d.min(), slice_2d.max()
-    if vmax - vmin < 1e-8:
-        slice_norm = np.zeros_like(slice_2d)
-    else:
-        slice_norm = (slice_2d - vmin) / (vmax - vmin)
-
-    # Resize to target size
-    slice_resized = cv2.resize(
-        slice_norm.astype(np.float32), (target_size, target_size)
-    )
-
-    # Scale to [-1, 1] for SAM3
-    slice_scaled = slice_resized * 2 - 1
-
-    # Convert to 3-channel tensor (B, C, H, W)
-    slice_tensor = torch.from_numpy(slice_scaled).float()
-    slice_tensor = slice_tensor.unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
-    slice_tensor = slice_tensor.repeat(1, 3, 1, 1)  # (1, 3, H, W)
+    return mask
 
-    return slice_tensor.to(DEVICE)
 
-
-def find_contours(mask: np.ndarray) -> List[List[List[float]]]:
+def find_contours(mask):
     """Extract contours from binary mask."""
     try:
         from skimage.measure import find_contours as sk_find_contours
@@ -181,7 +68,7 @@ def find_contours(mask: np.ndarray) -> List[List[List[float]]]:
         return []
 
 
-def keep_largest_component(mask: np.ndarray) -> np.ndarray:
+def keep_largest_component(mask):
     """Keep only the largest connected component."""
     try:
         from scipy import ndimage
@@ -196,8 +83,7 @@ def keep_largest_component(mask: np.ndarray) -> np.ndarray:
         return mask
 
 
-# Define the inference function with optional ZeroGPU decorator
-def _segment_slice_impl(
+def segment_slice_api(
     nifti_b64: str,
     slice_idx: int,
     text_prompt: str = "tumor",
@@ -205,15 +91,7 @@ def _segment_slice_impl(
 ):
     """
     Segment a single slice from a NIfTI volume.
-
-    Args:
-        nifti_b64: Base64-encoded NIfTI file bytes
-        slice_idx: Slice index to segment (0-indexed)
-        text_prompt: Text prompt for segmentation (e.g., "tumor", "lesion")
-        checkpoint: Model checkpoint name
-
-    Returns:
-        dict with keys: success, mask_b64, mask_shape, contours, slice_idx, inference_time_ms
+    Currently uses fallback segmentation (SAM3 to be integrated).
     """
     start_time = time.time()
 
@@ -242,56 +120,9 @@ def _segment_slice_impl(
                 "error": f"Slice index {slice_idx} out of range [0, {volume.shape[0]})",
             }
 
-        # Extract slice
+        # Extract slice and segment
         slice_2d = volume[slice_idx]
-        original_shape = slice_2d.shape
-
-        # Load model (may return None if fallback mode)
-        model = get_model(checkpoint)
-
-        if model is None or USE_FALLBACK:
-            # Use fallback segmentation
-            logger.info("Using fallback segmentation (SAM3 not available)")
-            mask = fallback_segment(slice_2d)
-            backend = "fallback"
-        else:
-            # Use SAM3 model
-            slice_tensor = preprocess_slice(
-                slice_2d
-            )  # (1, 3, 1024, 1024) tensor on DEVICE
-
-            # Create full-image bounding box prompt (auto-segment entire image)
-            # Format: [x_min, y_min, x_max, y_max] in pixel coordinates
-            target_size = slice_tensor.shape[-1]  # 1024
-            input_boxes = torch.tensor(
-                [[0, 0, target_size, target_size]], dtype=torch.float32, device=DEVICE
-            )
-
-            # Run inference with text prompt for grounding
-            with torch.no_grad():
-                outputs = model(
-                    pixel_values=slice_tensor,
-                    input_boxes=input_boxes,
-                    text_prompt=text_prompt,
-                )
-
-            # Extract mask from SAM3 output
-            # SAM3 returns a dict with 'pred_masks' key, shape (B, 1, H, W)
-            if isinstance(outputs, dict) and "pred_masks" in outputs:
-                pred_mask = outputs["pred_masks"][0, 0].cpu().numpy()
-            elif hasattr(outputs, "pred_masks"):
-                pred_mask = outputs.pred_masks[0, 0].cpu().numpy()
-            else:
-                # Fallback: try to extract from tuple/list
-                logger.warning(f"Unexpected output type: {type(outputs)}")
-                pred_mask = np.zeros((target_size, target_size))
-
-            # Resize mask back to original shape
-            mask = cv2.resize(pred_mask, (original_shape[1], original_shape[0]))
-            backend = "sam3"
-
-        # Threshold to binary
-        mask = (mask > 0.5).astype(np.uint8)
+        mask = fallback_segment(slice_2d)
         mask = keep_largest_component(mask)
 
         # Extract contours
@@ -307,7 +138,7 @@ def _segment_slice_impl(
 
         return {
             "success": True,
-            "backend": backend,
+            "backend": "fallback",
             "mask_b64": mask_b64,
             "mask_shape": list(mask.shape),
             "contours": contours,
@@ -320,146 +151,6 @@ def _segment_slice_impl(
         return {"success": False, "error": str(e)}
 
 
-def _segment_volume_impl(
-    nifti_b64: str,
-    text_prompt: str = "tumor",
-    checkpoint: str = "brain",
-    skip_empty: bool = True,
-    min_area: int = 50,
-):
-    """
-    Segment entire volume and return contours for all slices with detections.
-
-    Args:
-        nifti_b64: Base64-encoded NIfTI file bytes
-        text_prompt: Text prompt for segmentation
-        checkpoint: Model checkpoint name
-        skip_empty: Skip mostly-empty slices
-        min_area: Minimum mask area to report
-
-    Returns:
-        dict with keys: success, contours (dict), num_slices, slices_with_tumor, inference_time_ms
-    """
-    start_time = time.time()
-
-    try:
-        import nibabel as nib
-
-        # Decode NIfTI
-        nifti_bytes = base64.b64decode(nifti_b64)
-
-        with tempfile.NamedTemporaryFile(suffix=".nii.gz", delete=False) as f:
-            f.write(nifti_bytes)
-            temp_path = f.name
-
-        nii = nib.load(temp_path)
-        volume = nii.get_fdata().astype(np.float32)
-        os.unlink(temp_path)
-
-        logger.info(f"Loaded volume shape: {volume.shape}")
-
-        # Load model (may return None if fallback mode)
-        model = get_model(checkpoint)
-        use_fallback = model is None or USE_FALLBACK
-
-        num_slices = volume.shape[0]
-        all_contours = {}
-
-        target_size = 1024
-
-        for i in range(num_slices):
-            slice_2d = volume[i]
-            original_shape = slice_2d.shape
-
-            # Skip mostly-empty slices
-            if skip_empty and slice_2d.max() - slice_2d.min() < 0.01:
-                continue
-
-            if use_fallback:
-                # Use fallback segmentation
-                mask = fallback_segment(slice_2d)
-            else:
-                slice_tensor = preprocess_slice(slice_2d, target_size)
-
-                # Create full-image bounding box
-                input_boxes = torch.tensor(
-                    [[0, 0, target_size, target_size]],
-                    dtype=torch.float32,
-                    device=DEVICE,
-                )
-
-                with torch.no_grad():
-                    outputs = model(
-                        pixel_values=slice_tensor,
-                        input_boxes=input_boxes,
-                        text_prompt=text_prompt,
-                    )
-
-                # Extract mask from SAM3 output
-                if isinstance(outputs, dict) and "pred_masks" in outputs:
-                    pred_mask = outputs["pred_masks"][0, 0].cpu().numpy()
-                elif hasattr(outputs, "pred_masks"):
-                    pred_mask = outputs.pred_masks[0, 0].cpu().numpy()
-                else:
-                    continue  # Skip if no valid output
-
-                # Resize to original shape and threshold
-                mask = cv2.resize(pred_mask, (original_shape[1], original_shape[0]))
-                mask = (mask > 0.5).astype(np.uint8)
-
-            if mask.sum() >= min_area:
-                mask = keep_largest_component(mask)
-                contours = find_contours(mask)
-                if contours:
-                    all_contours[str(i)] = contours
-
-        inference_time = int((time.time() - start_time) * 1000)
-        logger.info(
-            f"Segmented {num_slices} slices in {inference_time}ms, found tumor in {len(all_contours)} slices"
-        )
-
-        return {
-            "success": True,
-            "contours": all_contours,
-            "num_slices": num_slices,
-            "slices_with_tumor": list(all_contours.keys()),
-            "inference_time_ms": inference_time,
-        }
-
-    except Exception as e:
-        logger.error(f"Volume segmentation failed: {e}")
-        return {"success": False, "error": str(e)}
-
-
-# Apply ZeroGPU decorator if available
-if ZEROGPU_AVAILABLE:
-
-    @spaces.GPU(duration=60)
-    def segment_slice_api(
-        nifti_b64: str,
-        slice_idx: int,
-        text_prompt: str = "tumor",
-        checkpoint: str = "brain",
-    ):
-        return _segment_slice_impl(nifti_b64, slice_idx, text_prompt, checkpoint)
-
-    @spaces.GPU(duration=300)
-    def segment_volume_api(
-        nifti_b64: str,
-        text_prompt: str = "tumor",
-        checkpoint: str = "brain",
-        skip_empty: bool = True,
-        min_area: int = 50,
-    ):
-        return _segment_volume_impl(
-            nifti_b64, text_prompt, checkpoint, skip_empty, min_area
-        )
-else:
-    segment_slice_api = _segment_slice_impl
-    segment_volume_api = _segment_volume_impl
-
-
-# Gradio UI functions (for interactive demo)
 def load_and_display_nifti(file):
     """Load NIfTI and return middle slice for display."""
     if file is None:
@@ -494,7 +185,7 @@ def load_and_display_nifti(file):
         return None, f"Error: {e}", 0
 
 
-def segment_and_overlay(file, slice_idx: int, text_prompt: str, checkpoint: str):
+def segment_and_overlay(file, slice_idx, text_prompt, checkpoint):
     """Segment a slice and overlay the mask."""
     if file is None:
         return None, "Please upload a file first"
@@ -504,7 +195,7 @@ def segment_and_overlay(file, slice_idx: int, text_prompt: str, checkpoint: str)
         with open(file.name, "rb") as f:
             nifti_b64 = base64.b64encode(f.read()).decode()
 
-        # Call segmentation API
+        # Call segmentation
         result = segment_slice_api(nifti_b64, int(slice_idx), text_prompt, checkpoint)
 
         if not result["success"]:
@@ -536,7 +227,7 @@ def segment_and_overlay(file, slice_idx: int, text_prompt: str, checkpoint: str)
         rgb[mask_bool, 1] = rgb[mask_bool, 1] * (1 - alpha) + 50 * alpha
         rgb[mask_bool, 2] = rgb[mask_bool, 2] * (1 - alpha) + 50 * alpha
 
-        info = f"Segmented in {result['inference_time_ms']}ms, mask area: {mask.sum()} pixels"
+        info = f"Backend: {result['backend']}, Segmented in {result['inference_time_ms']}ms, mask area: {mask.sum()} pixels"
 
         return rgb.astype(np.uint8), info
 
@@ -546,30 +237,27 @@ def segment_and_overlay(file, slice_idx: int, text_prompt: str, checkpoint: str)
 
 # Build Gradio interface
 def build_demo():
-    with gr.Blocks(
-        title="OncoSeg Inference API",
-        theme=gr.themes.Soft(),
-    ) as demo:
+    with gr.Blocks(title="OncoSeg Inference API", theme=gr.themes.Soft()) as demo:
         gr.Markdown("""
         # OncoSeg Medical Image Segmentation API
         
         GPU-accelerated segmentation for CT and MRI volumes.
         
+        **Note:** Currently using fallback segmentation. Full SAM3 model coming soon!
+        
         **API Endpoints** (for programmatic access):
         - `POST /api/segment_slice_api` - Segment a single slice
-        - `POST /api/segment_volume_api` - Segment entire volume
-        
-        **Interactive Demo** below:
         """)
 
         with gr.Row():
             with gr.Column(scale=1):
                 file_input = gr.File(
-                    label="Upload NIfTI (.nii, .nii.gz)", file_types=[".nii", ".nii.gz"]
+                    label="Upload NIfTI (.nii, .nii.gz)",
+                    file_types=[".nii", ".nii.gz", ".gz"],
                 )
 
                 checkpoint = gr.Dropdown(
-                    choices=list(CHECKPOINTS.keys()),
+                    choices=["brain", "liver", "breast", "lung", "kidney", "spine"],
                     value="brain",
                     label="Model Checkpoint",
                 )
@@ -622,7 +310,7 @@ def build_demo():
         
         # Call API
         response = requests.post(
-            "https://YOUR-SPACE.hf.space/api/segment_slice_api",
+            "https://tp53-oncoseg-api.hf.space/api/segment_slice_api",
             json={
                 "nifti_b64": nifti_b64,
                 "slice_idx": 77,