Add Phase 11: FastAPI backend for slicer-grade frontend support

seg_app/backend/__init__.py

@@ -0,0 +1,10 @@
+"""
+Backend API module for seg_app.
+
+Provides a FastAPI-based REST API for slicer-grade frontend support.
+Wraps the existing orchestrator without modification.
+"""
+
+from seg_app.backend.api import create_api_app
+
+__all__ = ["create_api_app"]

seg_app/backend/api.py

@@ -0,0 +1,586 @@
+"""
+FastAPI backend for medical image segmentation.
+
+This module provides a REST API that wraps the existing orchestrator
+to support slicer-grade frontends (VTK.js, MONAI-style viewers).
+
+Design principles:
+- No modification to existing orchestrator or model code
+- Importable without side effects (no model loading on import)
+- Simple in-memory state for single-user research use
+- Meaningful HTTP errors without exposing stack traces
+
+Usage:
+    from seg_app.backend.api import create_api_app
+    app = create_api_app()
+    # Run with: uvicorn seg_app.backend.api:app --reload
+"""
+
+import io
+import logging
+import tempfile
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+from fastapi import FastAPI, File, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import Response
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Pydantic Models for API Request/Response
+# =============================================================================
+
+class PromptTypeEnum(str, Enum):
+    """Prompt types matching orchestrator.PromptType."""
+    point_positive = "point_positive"
+    point_negative = "point_negative"
+    bounding_box = "bounding_box"
+
+
+class PromptModel(BaseModel):
+    """A single prompt for interactive segmentation."""
+    prompt_type: PromptTypeEnum
+    coordinates: List[int] = Field(
+        ...,
+        description="For points: [d, h, w]; for boxes: [d1, h1, w1, d2, h2, w2]"
+    )
+
+
+class PromptsModel(BaseModel):
+    """Collection of prompts for segmentation refinement."""
+    items: List[PromptModel] = Field(default_factory=list)
+
+
+class VolumeUploadResponse(BaseModel):
+    """Response after successful volume upload."""
+    volume_id: str
+    shape: Tuple[int, int, int]
+    spacing: Tuple[float, float, float]
+    message: str = "Volume uploaded successfully"
+
+
+class SegmentationRequest(BaseModel):
+    """Request to run segmentation on an uploaded volume."""
+    volume_id: str
+    model_id: str = Field(
+        default="medical-sam-3d",
+        description="Model ID: 'medical-sam-3d' or 'unet3d-brain-tumor'"
+    )
+    prompts: Optional[PromptsModel] = None
+
+
+class RefineRequest(BaseModel):
+    """Request to refine an existing segmentation with prompts."""
+    volume_id: str
+    prompts: PromptsModel
+
+
+class SegmentationResponse(BaseModel):
+    """Response after segmentation completes."""
+    volume_id: str
+    model_id: str
+    task_name: str
+    mask_shape: Tuple[int, int, int]
+    status: str = "completed"
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class ErrorResponse(BaseModel):
+    """Standard error response."""
+    detail: str
+    error_code: str
+
+
+class AvailableModelsResponse(BaseModel):
+    """List of available segmentation models."""
+    models: List[Dict[str, str]]
+
+
+# =============================================================================
+# In-Memory State Storage
+# =============================================================================
+
+@dataclass
+class VolumeState:
+    """State for a single uploaded volume."""
+    volume_id: str
+    array: np.ndarray
+    spacing: Tuple[float, float, float]
+    affine: Optional[np.ndarray] = None
+    mask: Optional[np.ndarray] = None
+    last_model_id: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+class StateManager:
+    """Simple in-memory state manager for volumes and masks.
+    
+    Single-user assumption: no concurrency handling.
+    """
+    
+    def __init__(self):
+        self._volumes: Dict[str, VolumeState] = {}
+    
+    def create_volume(
+        self,
+        array: np.ndarray,
+        spacing: Tuple[float, float, float],
+        affine: Optional[np.ndarray] = None,
+    ) -> str:
+        """Store a new volume and return its ID."""
+        volume_id = str(uuid.uuid4())[:8]  # Short ID for convenience
+        self._volumes[volume_id] = VolumeState(
+            volume_id=volume_id,
+            array=array,
+            spacing=spacing,
+            affine=affine,
+        )
+        logger.info(f"Created volume {volume_id} with shape {array.shape}")
+        return volume_id
+    
+    def get_volume(self, volume_id: str) -> Optional[VolumeState]:
+        """Retrieve a volume by ID."""
+        return self._volumes.get(volume_id)
+    
+    def update_mask(
+        self,
+        volume_id: str,
+        mask: np.ndarray,
+        model_id: str,
+    ) -> None:
+        """Update the segmentation mask for a volume."""
+        if volume_id not in self._volumes:
+            raise KeyError(f"Volume not found: {volume_id}")
+        self._volumes[volume_id].mask = mask
+        self._volumes[volume_id].last_model_id = model_id
+        logger.info(f"Updated mask for volume {volume_id} using model {model_id}")
+    
+    def delete_volume(self, volume_id: str) -> bool:
+        """Delete a volume and its associated mask."""
+        if volume_id in self._volumes:
+            del self._volumes[volume_id]
+            logger.info(f"Deleted volume {volume_id}")
+            return True
+        return False
+    
+    def list_volumes(self) -> List[str]:
+        """List all stored volume IDs."""
+        return list(self._volumes.keys())
+
+
+# Global state manager instance
+_state_manager = StateManager()
+
+
+# =============================================================================
+# Helper Functions
+# =============================================================================
+
+def _convert_prompts_to_orchestrator(prompts_model: Optional[PromptsModel]):
+    """Convert Pydantic PromptsModel to orchestrator Prompts object."""
+    if prompts_model is None or len(prompts_model.items) == 0:
+        return None
+    
+    # Import here to avoid circular imports and side effects
+    from seg_app.inference.orchestrator import Prompts, PromptType, Prompt
+    
+    prompts = Prompts()
+    for p in prompts_model.items:
+        if p.prompt_type == PromptTypeEnum.point_positive:
+            prompt_type = PromptType.POINT_POSITIVE
+        elif p.prompt_type == PromptTypeEnum.point_negative:
+            prompt_type = PromptType.POINT_NEGATIVE
+        else:
+            prompt_type = PromptType.BOUNDING_BOX
+        
+        prompts.items.append(Prompt(prompt_type, tuple(p.coordinates)))
+    
+    return prompts
+
+
+def _get_available_model_ids() -> List[str]:
+    """Get list of valid model IDs."""
+    # Import here to avoid side effects
+    from seg_app.inference.orchestrator import get_available_models
+    return [m["id"] for m in get_available_models()]
+
+
+# =============================================================================
+# FastAPI Application Factory
+# =============================================================================
+
+def create_api_app() -> FastAPI:
+    """Create and configure the FastAPI application.
+    
+    Returns:
+        Configured FastAPI app instance
+    """
+    app = FastAPI(
+        title="Brain Lesion Segmentation API",
+        description=(
+            "REST API for 3D medical image segmentation. "
+            "Supports volume upload, segmentation, and interactive refinement."
+        ),
+        version="1.0.0",
+        docs_url="/docs",
+        redoc_url="/redoc",
+    )
+    
+    # Configure CORS for local frontend development
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=[
+            "http://localhost:3000",      # React dev server
+            "http://localhost:5173",      # Vite dev server
+            "http://localhost:8080",      # Generic dev server
+            "http://127.0.0.1:3000",
+            "http://127.0.0.1:5173",
+            "http://127.0.0.1:8080",
+        ],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+    
+    # -------------------------------------------------------------------------
+    # Health Check
+    # -------------------------------------------------------------------------
+    
+    @app.get("/health")
+    async def health_check():
+        """Health check endpoint."""
+        return {"status": "healthy"}
+    
+    # -------------------------------------------------------------------------
+    # Model Information
+    # -------------------------------------------------------------------------
+    
+    @app.get("/models", response_model=AvailableModelsResponse)
+    async def list_models():
+        """Get list of available segmentation models."""
+        from seg_app.inference.orchestrator import get_available_models
+        return AvailableModelsResponse(models=get_available_models())
+    
+    # -------------------------------------------------------------------------
+    # Volume Upload
+    # -------------------------------------------------------------------------
+    
+    @app.post("/volume/upload", response_model=VolumeUploadResponse)
+    async def upload_volume(file: UploadFile = File(...)):
+        """Upload a 3D medical volume (NIfTI format).
+        
+        Accepts .nii or .nii.gz files.
+        Returns volume metadata and a volume_id for subsequent operations.
+        """
+        # Validate file extension
+        filename = file.filename or ""
+        if not (filename.endswith(".nii") or filename.endswith(".nii.gz")):
+            raise HTTPException(
+                status_code=400,
+                detail="Invalid file format. Please upload a NIfTI file (.nii or .nii.gz)"
+            )
+        
+        try:
+            # Save uploaded file to temporary location
+            suffix = ".nii.gz" if filename.endswith(".nii.gz") else ".nii"
+            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                content = await file.read()
+                tmp.write(content)
+                tmp_path = tmp.name
+            
+            # Load volume using existing I/O utilities
+            from seg_app.data.io import load_nifti
+            volume_data = load_nifti(tmp_path)
+            
+            # Clean up temp file
+            Path(tmp_path).unlink(missing_ok=True)
+            
+            # Store in state manager
+            volume_id = _state_manager.create_volume(
+                array=volume_data.array,
+                spacing=volume_data.spacing,
+                affine=volume_data.affine,
+            )
+            
+            return VolumeUploadResponse(
+                volume_id=volume_id,
+                shape=volume_data.array.shape,
+                spacing=volume_data.spacing,
+            )
+            
+        except Exception as e:
+            logger.error(f"Failed to load volume: {e}")
+            raise HTTPException(
+                status_code=400,
+                detail=f"Failed to load volume: {str(e)}"
+            )
+    
+    # -------------------------------------------------------------------------
+    # Segmentation
+    # -------------------------------------------------------------------------
+    
+    @app.post("/segment", response_model=SegmentationResponse)
+    async def run_segmentation(request: SegmentationRequest):
+        """Run segmentation on an uploaded volume.
+        
+        Requires volume_id from a previous upload.
+        Optionally accepts prompts for SAM-based models.
+        """
+        # Validate volume exists
+        volume_state = _state_manager.get_volume(request.volume_id)
+        if volume_state is None:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Volume not found: {request.volume_id}. Please upload a volume first."
+            )
+        
+        # Validate model ID
+        valid_models = _get_available_model_ids()
+        if request.model_id not in valid_models:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Invalid model_id: {request.model_id}. Valid options: {valid_models}"
+            )
+        
+        try:
+            # Import orchestrator here to avoid side effects
+            from seg_app.inference.orchestrator import run_segmentation as orchestrator_run
+            
+            # Convert prompts
+            prompts = _convert_prompts_to_orchestrator(request.prompts)
+            
+            # Determine if this is refinement mode
+            existing_mask = None
+            if prompts is not None and volume_state.mask is not None:
+                existing_mask = volume_state.mask
+            
+            # Run segmentation
+            result = orchestrator_run(
+                volume=volume_state.array,
+                task_name="brain_lesion",
+                prompts=prompts,
+                existing_mask=existing_mask,
+                spacing=volume_state.spacing,
+                model_id=request.model_id,
+            )
+            
+            # Store mask
+            _state_manager.update_mask(
+                request.volume_id,
+                result.mask,
+                result.model_id,
+            )
+            
+            return SegmentationResponse(
+                volume_id=request.volume_id,
+                model_id=result.model_id,
+                task_name=result.task_name,
+                mask_shape=result.mask.shape,
+                metadata=result.metadata,
+            )
+            
+        except ValueError as e:
+            # User-recoverable errors (e.g., SAM requires prompts)
+            raise HTTPException(status_code=400, detail=str(e))
+        except Exception as e:
+            logger.error(f"Segmentation failed: {e}")
+            raise HTTPException(
+                status_code=500,
+                detail=f"Segmentation failed: {str(e)}"
+            )
+    
+    # -------------------------------------------------------------------------
+    # Refinement
+    # -------------------------------------------------------------------------
+    
+    @app.post("/refine", response_model=SegmentationResponse)
+    async def refine_segmentation(request: RefineRequest):
+        """Refine an existing segmentation using interactive prompts.
+        
+        Requires a previous segmentation to exist for the volume.
+        Uses SAM-Med3D for refinement.
+        """
+        # Validate volume exists
+        volume_state = _state_manager.get_volume(request.volume_id)
+        if volume_state is None:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Volume not found: {request.volume_id}. Please upload a volume first."
+            )
+        
+        # Validate existing mask
+        if volume_state.mask is None:
+            raise HTTPException(
+                status_code=400,
+                detail="No existing segmentation found. Run /segment first before refinement."
+            )
+        
+        # Validate prompts
+        if request.prompts is None or len(request.prompts.items) == 0:
+            raise HTTPException(
+                status_code=400,
+                detail="Prompts are required for refinement."
+            )
+        
+        try:
+            from seg_app.inference.orchestrator import run_segmentation as orchestrator_run
+            
+            prompts = _convert_prompts_to_orchestrator(request.prompts)
+            
+            # Run refinement (SAM mode)
+            result = orchestrator_run(
+                volume=volume_state.array,
+                task_name="brain_lesion",
+                prompts=prompts,
+                existing_mask=volume_state.mask,
+                spacing=volume_state.spacing,
+                model_id="medical-sam-3d",
+                full_reinference=False,
+            )
+            
+            # Update mask
+            _state_manager.update_mask(
+                request.volume_id,
+                result.mask,
+                result.model_id,
+            )
+            
+            return SegmentationResponse(
+                volume_id=request.volume_id,
+                model_id=result.model_id,
+                task_name=result.task_name,
+                mask_shape=result.mask.shape,
+                status="refined",
+                metadata=result.metadata,
+            )
+            
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+        except Exception as e:
+            logger.error(f"Refinement failed: {e}")
+            raise HTTPException(
+                status_code=500,
+                detail=f"Refinement failed: {str(e)}"
+            )
+    
+    # -------------------------------------------------------------------------
+    # Mask Retrieval
+    # -------------------------------------------------------------------------
+    
+    @app.get("/mask/{volume_id}")
+    async def get_mask(volume_id: str, format: str = "npy"):
+        """Retrieve the segmentation mask for a volume.
+        
+        Args:
+            volume_id: ID of the volume
+            format: Output format - 'npy' (compressed numpy) or 'nifti'
+        
+        Returns:
+            Binary mask data in requested format
+        """
+        # Validate volume exists
+        volume_state = _state_manager.get_volume(volume_id)
+        if volume_state is None:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Volume not found: {volume_id}"
+            )
+        
+        # Validate mask exists
+        if volume_state.mask is None:
+            raise HTTPException(
+                status_code=404,
+                detail=f"No segmentation mask found for volume {volume_id}. Run /segment first."
+            )
+        
+        if format == "npy":
+            # Return compressed numpy array
+            buffer = io.BytesIO()
+            np.savez_compressed(buffer, mask=volume_state.mask)
+            buffer.seek(0)
+            return Response(
+                content=buffer.getvalue(),
+                media_type="application/octet-stream",
+                headers={
+                    "Content-Disposition": f"attachment; filename=mask_{volume_id}.npz"
+                }
+            )
+        
+        elif format == "nifti":
+            # Return as NIfTI file
+            try:
+                import nibabel as nib
+                
+                # Create NIfTI image with affine
+                affine = volume_state.affine
+                if affine is None:
+                    # Create identity affine with spacing
+                    affine = np.diag([*volume_state.spacing, 1.0])
+                
+                nifti_img = nib.Nifti1Image(
+                    volume_state.mask.astype(np.uint8),
+                    affine=affine
+                )
+                
+                buffer = io.BytesIO()
+                nib.save(nifti_img, buffer)
+                buffer.seek(0)
+                
+                return Response(
+                    content=buffer.getvalue(),
+                    media_type="application/octet-stream",
+                    headers={
+                        "Content-Disposition": f"attachment; filename=mask_{volume_id}.nii.gz"
+                    }
+                )
+            except Exception as e:
+                logger.error(f"Failed to create NIfTI: {e}")
+                raise HTTPException(
+                    status_code=500,
+                    detail=f"Failed to create NIfTI file: {str(e)}"
+                )
+        
+        else:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Invalid format: {format}. Supported: 'npy', 'nifti'"
+            )
+    
+    # -------------------------------------------------------------------------
+    # Volume Management
+    # -------------------------------------------------------------------------
+    
+    @app.get("/volumes")
+    async def list_volumes():
+        """List all uploaded volume IDs."""
+        return {"volume_ids": _state_manager.list_volumes()}
+    
+    @app.delete("/volume/{volume_id}")
+    async def delete_volume(volume_id: str):
+        """Delete an uploaded volume and its mask."""
+        if _state_manager.delete_volume(volume_id):
+            return {"message": f"Volume {volume_id} deleted"}
+        else:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Volume not found: {volume_id}"
+            )
+    
+    return app
+
+
+# Create app instance for uvicorn
+app = create_api_app()
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="127.0.0.1", port=8000)