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()