"""
Local Gradio demo for kavin-aravindhan/vit-oct-wamd.

There is no hosted Space for this model (Gradio/Docker Spaces require a
paid HF plan; this repo is free). Run this script locally instead:

    pip install -r requirements-demo.txt
    python demo_app.py

Then open the printed local URL in a browser and upload an OCT B-scan.

NOTE: requirements-demo.txt pins a specific gradio + transformers +
huggingface_hub + pydantic/starlette/fastapi combination that was verified
to actually work together end-to-end. This isn't cosmetic -- gradio's
latest releases require huggingface_hub>=1.0, which conflicts with the
transformers==4.53.0 needed to load best_model.pt's state dict, and an
unpinned newer pydantic/starlette paired with older gradio hits real
upstream bugs (schema generation and template caching both broke in
testing). Deviating from requirements-demo.txt may silently break this.

Loads only the classifier weights (image_encoder + cls_head) from the
checkpoint -- the auxiliary T5 text-alignment branch in the full checkpoint
is unused at inference time, so it's skipped here (via strict=False) to
keep memory/startup time down.
"""
import gc

import cv2
import gradio as gr
import numpy as np
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from transformers import SiglipVisionModel

MODEL_REPO = "kavin-aravindhan/vit-oct-wamd"
IMAGE_SIZE = 384
LABELS = ["normal", "wet_amd"]


class LeanSigLIPClassifier(nn.Module):
    """image_encoder + cls_head only -- matches the classification path of
    modeling.py's forward(), without the unused alignment branch."""

    def __init__(self, dropout_rate=0.057129660535791494):
        super().__init__()
        self.image_encoder = SiglipVisionModel.from_pretrained("google/siglip-so400m-patch14-384")
        self.dropout = nn.Dropout(dropout_rate)
        self.cls_head = nn.Linear(1152, 2)

    def forward(self, images):
        img_features = self.image_encoder(pixel_values=images).last_hidden_state
        cls_features = self.dropout(img_features[:, 0])
        return self.cls_head(cls_features)


def load_model():
    ckpt_path = hf_hub_download(repo_id=MODEL_REPO, filename="best_model.pt")
    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    model = LeanSigLIPClassifier()
    # strict=False: checkpoint also has siglip_loss.* (unused T5 branch) keys we don't have.
    missing, unexpected = model.load_state_dict(checkpoint["model_state_dict"], strict=False)
    assert not missing, f"Missing expected classifier weights: {missing}"
    del checkpoint
    gc.collect()
    model.eval()
    return model


print("Loading model (first launch may take a minute)...")
MODEL = load_model()
print("Model ready.")


def preprocess(pil_image):
    img_array = np.array(pil_image.convert("RGB"), dtype=np.float32) / 255.0
    img_resized = cv2.resize(img_array, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_LINEAR)
    img_tensor = torch.from_numpy(img_resized).permute(2, 0, 1)
    img_tensor = (img_tensor - 0.5) / 0.5
    return img_tensor.unsqueeze(0)


@torch.no_grad()
def classify(pil_image):
    if pil_image is None:
        return {}
    pixel_values = preprocess(pil_image)
    logits = MODEL(pixel_values)
    probs = torch.softmax(logits, dim=-1)[0]
    return {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}


DESCRIPTION = """
# vit-oct-wamd

A SigLIP vision transformer fine-tuned to classify OCT B-scans as **normal**
or **wet AMD**. Upload an OCT B-scan image to see the model's prediction.

Model: [kavin-aravindhan/vit-oct-wamd](https://huggingface.co/kavin-aravindhan/vit-oct-wamd)

⚠️ **Research use only.** This model is released as supporting material for
a paper and has not been clinically validated. Do not use it for diagnosis
or clinical decision-making.
"""

demo = gr.Interface(
    fn=classify,
    inputs=gr.Image(type="pil", label="OCT B-scan"),
    outputs=gr.Label(num_top_classes=2, label="Prediction"),
    title="vit-oct-wamd — OCT Wet AMD Classifier",
    description=DESCRIPTION,
    examples=None,
    allow_flagging="never",
)

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