app.py

import gradio as gr                         # Gradio for the web interface
import numpy as np                          # NumPy for array operations
import onnxruntime as ort                   # ONNX Runtime to run both models
from PIL import Image                       # Pillow for image processing
from huggingface_hub import hf_hub_download # download models from HuggingFace
from transformers import SiglipProcessor   # handles MedSigLIP preprocessing
from torchvision import transforms          # handles gatekeeper preprocessing

# ======================== CONFIG ========================

# HuggingFace repo where all models are stored
REPO_ID  = "KhanyiTapiwa00/medsiglip-diagnosis"

# HuggingFace model ID for loading the SiglipProcessor
MODEL_ID = "KhanyiTapiwa00/medsiglip-diagnosis"

# Diagnostic classes for the main MedSigLIP classifier
classes = ["Normal", "CIN1", "CIN2", "CIN3", "Cancer"]

# ======================== DOWNLOAD MODELS ========================

# Download the gatekeeper ONNX model — ResNet50 binary classifier (91 MB)
print("Downloading gatekeeper model...")
GATEKEEPER_PATH = hf_hub_download(repo_id=REPO_ID, filename="gatekeeper_direct.onnx")
print(f"Gatekeeper saved to: {GATEKEEPER_PATH}")

# Download the main MedSigLIP ONNX model — full diagnosis model (1.64 GB)
print("Downloading MedSigLIP model...")
MEDSIGLIP_PATH = hf_hub_download(repo_id=REPO_ID, filename="medsiglip_direct.onnx")
print(f"MedSigLIP saved to: {MEDSIGLIP_PATH}")

# ======================== LOAD MODELS ========================

# Load gatekeeper into ONNX inference session
print("Loading gatekeeper model...")
gatekeeper_session = ort.InferenceSession(GATEKEEPER_PATH)
gatekeeper_input   = gatekeeper_session.get_inputs()[0].name   # input tensor name
print(f"Gatekeeper input: {gatekeeper_input} | shape: {gatekeeper_session.get_inputs()[0].shape}")

# Load MedSigLIP into ONNX inference session
print("Loading MedSigLIP model...")
medsiglip_session = ort.InferenceSession(MEDSIGLIP_PATH)
medsiglip_input   = medsiglip_session.get_inputs()[0].name     # input tensor name
print(f"MedSigLIP input: {medsiglip_input} | shape: {medsiglip_session.get_inputs()[0].shape}")

# Load SiglipProcessor — handles MedSigLIP preprocessing exactly as during training
print("Loading SiglipProcessor...")
processor = SiglipProcessor.from_pretrained(MODEL_ID)
print("All models loaded successfully!")

# ======================== GATEKEEPER PREPROCESSING ========================

# Gatekeeper uses ImageNet normalization — ResNet50 expects 224x224
# Must match exactly what was used during gatekeeper training
gatekeeper_transforms = transforms.Compose([
    transforms.Resize((224, 224)),           # resize to ResNet50 expected size
    transforms.ToTensor(),                   # convert PIL image to tensor (C,H,W)
    transforms.Normalize(                    # ImageNet normalization
        mean=[0.485, 0.456, 0.406],          # ImageNet RGB mean per channel
        std=[0.229, 0.224, 0.225]            # ImageNet RGB std per channel
    )
])

# ======================== HELPER FUNCTIONS ========================

def softmax(x):
    """Convert raw logits to probabilities that sum to 1."""
    e_x = np.exp(x - np.max(x))   # subtract max for numerical stability
    return e_x / e_x.sum(axis=0)  # normalize

def calculate_entropy(probs):
    """High entropy = uncertain. Low entropy = confident."""
    return -np.sum(probs * np.log(probs + 1e-10))  # 1e-10 avoids log(0)

# ======================== PREDICTION LOGIC ========================

def predict(image):
    """
    Two-stage inference pipeline:
    Stage 1 — Gatekeeper: is this a colposcopy image?
    Stage 2 — MedSigLIP: what is the cervical diagnosis?
    """

    # Guard: user uploaded nothing
    if image is None:
        return "Please upload a colposcopy image"

    try:
        # Convert Gradio numpy array to PIL Image in RGB
        img_pil = Image.fromarray(image).convert("RGB")

        # ================================================
        # STAGE 1 — GATEKEEPER
        # Checks if this is a valid colposcopy image
        # before passing to the main classifier
        # ================================================

        # Preprocess for gatekeeper — 224x224 with ImageNet normalization
        gatekeeper_tensor = gatekeeper_transforms(img_pil)         # apply transforms
        gatekeeper_tensor = gatekeeper_tensor.unsqueeze(0)         # add batch dim → (1,3,224,224)
        gatekeeper_array  = gatekeeper_tensor.numpy()              # convert to numpy for ONNX

        # Run gatekeeper inference
        gate_outputs = gatekeeper_session.run(None, {gatekeeper_input: gatekeeper_array})
        gate_logits  = gate_outputs[0][0]          # shape (2,) — [cervix, not_cervix]
        gate_probs   = softmax(gate_logits)         # convert to probabilities

        # Log gatekeeper result for debugging
        print(f"Gatekeeper — cervix: {gate_probs[0]:.2%} | not_cervix: {gate_probs[1]:.2%}")

        # Gatekeeper decision — 0=cervix, 1=not_cervix
        gate_pred = int(np.argmax(gate_probs))

        # Reject image if gatekeeper says not_cervix
        if gate_pred == 1:
            return (
                "### ❌ Image Rejected\n\n"
                "This does not appear to be a colposcopy image.\n\n"
                f"**Gatekeeper Confidence:**\n"
                f"- Colposcopy: {gate_probs[0]:.1%}\n"
                f"- Not Colposcopy: {gate_probs[1]:.1%}\n\n"
                "Please upload a valid cervical colposcopy image."
            )

        # ================================================
        # STAGE 2 — MEDSIGLIP DIAGNOSIS
        # Only reached if gatekeeper approves the image
        # ================================================

        # Preprocess for MedSigLIP using SiglipProcessor
        # Matches exactly how images were preprocessed during training
        inputs       = processor(images=img_pil, return_tensors="np")
        pixel_values = inputs["pixel_values"].astype(np.float32)   # shape (1,3,448,448)

        # Log for debugging
        print(f"MedSigLIP input shape: {pixel_values.shape} | min: {pixel_values.min():.3f} | max: {pixel_values.max():.3f}")

        # Run MedSigLIP inference
        med_outputs = medsiglip_session.run(None, {medsiglip_input: pixel_values})
        logits      = med_outputs[0][0]    # shape (5,) — one score per class
        print(f"MedSigLIP logits: {logits}")

        # Convert logits to probabilities
        probs      = softmax(logits)
        confidence = float(np.max(probs))      # highest class probability
        entropy    = calculate_entropy(probs)  # uncertainty measure
        pred_idx   = np.argmax(probs)          # predicted class index
        pred_label = classes[pred_idx]         # predicted class name

        print(f"Probs: {probs} | Confidence: {confidence:.2%} | Entropy: {entropy:.3f}")

        # ---- Hallucination check ----
        # High confidence AND high entropy = contradictory = unreliable
        if confidence > 0.75 and entropy > 1.6:
            return (
                "### ⚠️ Uncertain Prediction\n\n"
                "The model produced a contradictory result.\n"
                "This may indicate poor image quality. Please retake the image.\n"
            )

        # ---- Low confidence check ----
        # Model is unsure — show all probabilities without a single prediction
        if confidence < 0.45:
            formatted_probs = "\n".join([f"- {c}: {p:.1%}" for c, p in zip(classes, probs)])
            return "### ⚠️ Uncertain Result — Please Review With Clinician\n\n" + formatted_probs

        # ---- Final result ----
        result  = f"### Prediction: {pred_label}\n"
        result += f"### Confidence: {confidence:.1%}\n\n"
        result += f"**Image Quality Check:** ✅ Valid colposcopy image ({gate_probs[0]:.1%} confidence)\n\n"
        result += "**Detailed Probabilities:**\n"
        for c, p in zip(classes, probs):
            result += f"- {c}: {p:.1%}\n"
        return result

    except Exception as e:
        # Surface any unexpected errors in the UI
        return f"Error during inference: {str(e)}"

# ======================== GRADIO INTERFACE ========================

with gr.Blocks() as demo:
    gr.Markdown("# MedSigLIP Colposcopy Diagnosis")
    gr.Markdown(
        "**Two-stage pipeline:** Gatekeeper (ResNet50) validates the image first, "
        "then MedSigLIP makes the cervical diagnosis."
    )

    # Image input — Gradio returns numpy array when type='numpy'
    inp = gr.Image(label="Upload colposcopy image", type="numpy")

    # Output rendered as Markdown
    out = gr.Markdown()

    # Trigger prediction whenever image changes
    inp.change(predict, inp, out)

# Launch — HuggingFace Spaces handles the server automatically
demo.launch()