app.py

"""
OncoSeg API - 3D Tumor Segmentation for Medical Imaging
HuggingFace Space: tp53/oncoseg-api

This provides a Gradio API for 3D tumor segmentation using the OncoSeg model.
Accepts NIfTI files and returns segmentation masks.

FIX: Returns JSON as string to avoid Gradio schema issues with nested dicts.
"""

import gradio as gr
import numpy as np
import json
import base64
import gzip
from typing import Optional
import os

# Model loading (lazy)
model = None
model_name = "default"


def load_model(checkpoint: str = "default"):
    """Load the segmentation model (stub - replace with actual model loading)"""
    global model, model_name

    # In production, load actual OncoSeg model here
    # For now, return a stub that generates demo segmentation
    model_name = checkpoint
    model = "loaded"

    return f"Model '{checkpoint}' loaded"


def segment_nifti(
    nifti_base64: str, checkpoint: str = "default", threshold: float = 0.5
) -> str:
    """
    Segment a 3D NIfTI volume.

    Args:
        nifti_base64: Base64-encoded gzipped NIfTI file
        checkpoint: Model checkpoint to use
        threshold: Segmentation threshold (0-1)

    Returns:
        JSON string with segmentation results (NOT dict - Gradio compatibility)
    """
    global model, model_name

    try:
        # Load model if needed
        if model is None or model_name != checkpoint:
            load_model(checkpoint)

        # Decode the NIfTI file
        nifti_bytes = base64.b64decode(nifti_base64)

        # Check if gzipped
        if nifti_bytes[:2] == b"\x1f\x8b":
            nifti_bytes = gzip.decompress(nifti_bytes)

        # Parse NIfTI header to get dimensions
        # NIfTI-1 header: dims at offset 40, 8 int16 values
        # dim[0] = number of dimensions, dim[1-7] = sizes
        import struct

        # Check for NIfTI magic number
        magic = nifti_bytes[344:348]
        if magic not in [b"n+1\x00", b"ni1\x00"]:
            # Try as raw volume (assume 64x64x64 for demo)
            shape = (64, 64, 64)
        else:
            dims = struct.unpack("<8h", nifti_bytes[40:56])
            ndim = dims[0]
            shape = tuple(dims[1 : ndim + 1])

        # Generate segmentation mask
        # In production, run actual model inference here
        mask = generate_demo_mask(shape, threshold)

        # Encode mask as base64 for transmission
        mask_bytes = mask.astype(np.uint8).tobytes()
        mask_compressed = gzip.compress(mask_bytes)
        mask_base64 = base64.b64encode(mask_compressed).decode("utf-8")

        # Calculate tumor statistics
        tumor_voxels = int(np.sum(mask > 0))
        total_voxels = int(np.prod(shape))
        tumor_percentage = (
            (tumor_voxels / total_voxels) * 100 if total_voxels > 0 else 0
        )

        # Find bounding box of tumor
        if tumor_voxels > 0:
            coords = np.where(mask > 0)
            bbox = {
                "min": [int(c.min()) for c in coords],
                "max": [int(c.max()) for c in coords],
            }
        else:
            bbox = None

        result = {
            "success": True,
            "shape": list(shape),
            "mask_base64": mask_base64,
            "mask_dtype": "uint8",
            "statistics": {
                "tumor_voxels": tumor_voxels,
                "total_voxels": total_voxels,
                "tumor_percentage": round(tumor_percentage, 2),
                "bounding_box": bbox,
            },
            "model": checkpoint,
            "threshold": threshold,
        }

        # IMPORTANT: Return as JSON string, not dict
        # This fixes Gradio's JSON schema validation issues
        return json.dumps(result)

    except Exception as e:
        error_result = {
            "success": False,
            "error": str(e),
            "error_type": type(e).__name__,
        }
        return json.dumps(error_result)


def generate_demo_mask(shape: tuple, threshold: float = 0.5) -> np.ndarray:
    """
    Generate a demo segmentation mask (ellipsoid tumor).
    Replace this with actual model inference in production.
    """
    # Create coordinate grids
    z, y, x = np.ogrid[: shape[0], : shape[1], : shape[2]]

    # Center of the volume
    cz, cy, cx = shape[0] // 2, shape[1] // 2, shape[2] // 2

    # Ellipsoid radii (tumor size ~15-25% of volume)
    rz = shape[0] * 0.15
    ry = shape[1] * 0.18
    rx = shape[2] * 0.20

    # Create ellipsoid mask
    distance = ((z - cz) / rz) ** 2 + ((y - cy) / ry) ** 2 + ((x - cx) / rx) ** 2
    mask = (distance <= 1.0).astype(np.uint8)

    # Add some irregularity
    np.random.seed(42)  # Reproducible
    noise = np.random.rand(*shape) * 0.3
    mask = ((distance <= 1.0 + noise * 0.5) & (distance <= 1.3)).astype(np.uint8)

    return mask


def get_available_checkpoints() -> str:
    """
    Get list of available model checkpoints.

    Returns:
        JSON string with checkpoint list
    """
    checkpoints = {
        "checkpoints": [
            {
                "id": "default",
                "name": "OncoSeg Default",
                "description": "General purpose tumor segmentation",
                "modalities": ["CT", "MRI"],
            },
            {
                "id": "liver",
                "name": "OncoSeg Liver",
                "description": "Optimized for liver tumors",
                "modalities": ["CT"],
            },
            {
                "id": "brain",
                "name": "OncoSeg Brain",
                "description": "Optimized for brain tumors",
                "modalities": ["MRI"],
            },
            {
                "id": "lung",
                "name": "OncoSeg Lung",
                "description": "Optimized for lung nodules",
                "modalities": ["CT"],
            },
        ]
    }
    return json.dumps(checkpoints)


def health_check() -> str:
    """Health check endpoint"""
    return json.dumps(
        {
            "status": "healthy",
            "model_loaded": model is not None,
            "model_name": model_name,
        }
    )


# Create Gradio interface
with gr.Blocks(title="OncoSeg API") as demo:
    gr.Markdown("""
    # OncoSeg API - 3D Tumor Segmentation
    
    This API provides 3D tumor segmentation for medical imaging (NIfTI format).
    
    ## API Endpoints
    
    Use the Gradio API client to call these functions:
    
    - `segment_nifti(nifti_base64, checkpoint, threshold)` - Segment a NIfTI volume
    - `get_available_checkpoints()` - List available model checkpoints
    - `health_check()` - Check API health
    
    ## Usage Example (Python)
    
    ```python
    from gradio_client import Client
    import base64
    import json
    
    client = Client("tp53/oncoseg-api")
    
    # Load and encode NIfTI file
    with open("scan.nii.gz", "rb") as f:
        nifti_b64 = base64.b64encode(f.read()).decode()
    
    # Call segmentation
    result_str = client.predict(
        nifti_base64=nifti_b64,
        checkpoint="default",
        threshold=0.5,
        api_name="/segment_nifti"
    )
    
    # Parse result (returns JSON string)
    result = json.loads(result_str)
    print(f"Tumor: {result['statistics']['tumor_percentage']:.1f}%")
    ```
    """)

    with gr.Tab("Segment"):
        with gr.Row():
            with gr.Column():
                nifti_input = gr.Textbox(
                    label="NIfTI Base64",
                    placeholder="Base64-encoded NIfTI file",
                    lines=3,
                )
                checkpoint_input = gr.Dropdown(
                    choices=["default", "liver", "brain", "lung"],
                    value="default",
                    label="Model Checkpoint",
                )
                threshold_input = gr.Slider(
                    minimum=0.1, maximum=0.9, value=0.5, step=0.1, label="Threshold"
                )
                segment_btn = gr.Button("Segment", variant="primary")

            with gr.Column():
                output = gr.Textbox(label="Result (JSON)", lines=10)

        segment_btn.click(
            fn=segment_nifti,
            inputs=[nifti_input, checkpoint_input, threshold_input],
            outputs=output,
        )

    with gr.Tab("Checkpoints"):
        checkpoints_btn = gr.Button("Get Checkpoints")
        checkpoints_output = gr.Textbox(label="Available Checkpoints", lines=10)
        checkpoints_btn.click(fn=get_available_checkpoints, outputs=checkpoints_output)

    with gr.Tab("Health"):
        health_btn = gr.Button("Health Check")
        health_output = gr.Textbox(label="Status", lines=3)
        health_btn.click(fn=health_check, outputs=health_output)


# Launch
if __name__ == "__main__":
    demo.launch()