app.py

"""
app.py - Gradio Demo for ResNet-50 Baseline (Diabetic Retinopathy Classification)
"""

import os
import json
import torch
import torch.nn as nn
import numpy as np
import cv2
from PIL import Image
from torchvision import models, transforms
from huggingface_hub import hf_hub_download
import gradio as gr

# ========================== MODEL DEFINITION ==========================

class ResNet50_DR(nn.Module):
    def __init__(self, num_classes=5, drop_rate=0.3):
        super().__init__()
        self.model = models.resnet50(weights=None)
        in_features = self.model.fc.in_features
        self.model.fc = nn.Sequential(
            nn.Dropout(p=drop_rate),
            nn.Linear(in_features, num_classes),
        )

    def forward(self, x):
        return self.model(x)

    def load_state_dict(self, state_dict, strict=True):
        if any(k.startswith("model.") for k in state_dict.keys()):
            return super().load_state_dict(state_dict, strict=strict)
        else:
            return self.model.load_state_dict(state_dict, strict=strict)



# ========================== PREPROCESSING ==========================

CROP_TOLERANCE = 12
BEN_SIGMA = 10
TARGET_SIZE = (224, 224)
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def crop_fundus_circle(img, tolerance=CROP_TOLERANCE):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    _, mask = cv2.threshold(gray, tolerance, 255, cv2.THRESH_BINARY)
    coords = cv2.findNonZero(mask)
    if coords is None:
        return img
    x, y, w, h = cv2.boundingRect(coords)
    return img[y: y + h, x: x + w]


def auto_detect_border(img, thresh=0.05):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    return float((gray < CROP_TOLERANCE).mean()) > thresh


def letterbox_resize(img, target_size=TARGET_SIZE):
    h, w = img.shape[:2]
    th, tw = target_size
    scale = min(tw / w, th / h)
    nw, nh = int(w * scale), int(h * scale)
    resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_CUBIC)
    canvas = np.zeros((th, tw, 3), dtype=np.uint8)
    pad_y = (th - nh) // 2
    pad_x = (tw - nw) // 2
    canvas[pad_y: pad_y + nh, pad_x: pad_x + nw] = resized
    return canvas


def ben_graham_transform(img, sigma_x=BEN_SIGMA):
    blur = cv2.GaussianBlur(img, (0, 0), sigmaX=sigma_x)
    enhanced = cv2.addWeighted(img, 4, blur, -4, 128)
    return np.clip(enhanced, 0, 255).astype(np.uint8)


def preprocess_image(pil_image):
    """Full pipeline: PIL Image -> preprocessed BGR numpy array."""
    img_rgb = np.array(pil_image.convert("RGB"))
    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)

    if auto_detect_border(img_bgr):
        img_bgr = crop_fundus_circle(img_bgr)

    img_bgr = letterbox_resize(img_bgr, TARGET_SIZE)
    img_bgr = ben_graham_transform(img_bgr)
    return img_bgr


def to_tensor(img_bgr):
    """BGR numpy -> normalized PyTorch tensor."""
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(img_rgb)
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    ])
    return transform(pil_img)


# ========================== LOAD MODEL ==========================

REPO_ID = "chrisnguyenx/ResNet50-P3"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

CLASS_NAMES = {
    0: "No DR",
    1: "Mild",
    2: "Moderate",
    3: "Severe",
    4: "Proliferative DR",
}

CLINICAL_ADVICE = {
    0: {
        "title": "Mat Binh Thuong (No DR)",
        "emoji": "✅",
        "color": "#10b981",
        "advice": "Chua phat hien ton thuong vong mac tieu duong. Khuyen nghi kham mat dinh ky 12 thang/lan.",
        "urgency": "Binh thuong",
    },
    1: {
        "title": "Benh Nhe (Mild DR)",
        "emoji": "🔵",
        "color": "#3b82f6",
        "advice": "Xuat hien vi phinh mach nho. Khuyen nghi tai kham sau 6-12 thang va kiem soat duong huyet.",
        "urgency": "Theo doi dinh ky",
    },
    2: {
        "title": "Benh Trung Binh (Moderate DR)",
        "emoji": "🟡",
        "color": "#f59e0b",
        "advice": "Ton thuong xuat huyet/xuat tiet muc do vua. Can kham bac si nhan khoa trong 3-6 thang.",
        "urgency": "Kham chuyen khoa",
    },
    3: {
        "title": "Benh Nang (Severe DR)",
        "emoji": "🔴",
        "color": "#ef4444",
        "advice": "Ton thuong nghiem trong o nhieu goc phan tu vong mac. CAN chuyen kham gap trong 2-4 tuan.",
        "urgency": "Can can thiep som",
    },
    4: {
        "title": "Tang Sinh Nguy Hiem (Proliferative DR)",
        "emoji": "🟣",
        "color": "#8b5cf6",
        "advice": "Tang sinh tan mach nguy co gay mo mat vinh vien! CAN DIEU TRI KHAN CAP.",
        "urgency": "KHAN CAP",
    },
}


def load_model():
    """Load model from local file or download weights from HF Hub."""
    local_weights = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "resnet50_baseline_fold1.pth")
    if os.path.exists(local_weights):
        weights_path = local_weights
        print(f"Loading local weights from {weights_path}")
    elif os.path.exists("resnet50_baseline_fold1.pth"):
        weights_path = "resnet50_baseline_fold1.pth"
    else:
        weights_path = hf_hub_download(
            repo_id=REPO_ID,
            filename="resnet50_baseline_fold1.pth",
        )

    model = ResNet50_DR(num_classes=5, drop_rate=0.3)

    checkpoint = torch.load(weights_path, map_location=DEVICE, weights_only=False)
    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
        state_dict = checkpoint["model_state_dict"]
    else:
        state_dict = checkpoint

    model.load_state_dict(state_dict)
    model.to(DEVICE)
    model.eval()
    return model


print("Loading model from HuggingFace Hub...")
model = load_model()
print(f"Model loaded on {DEVICE}")


# ========================== INFERENCE ==========================

def predict(image):
    """Main prediction function for Gradio."""
    if image is None:
        return None, "Please upload a retinal fundus image.", ""

    # Preprocess
    img_bgr = preprocess_image(image)
    tensor = to_tensor(img_bgr).unsqueeze(0).to(DEVICE)

    # Inference
    with torch.no_grad():
        outputs = model(tensor)
        probs = torch.softmax(outputs, dim=1)[0].cpu().numpy()

    pred_class = int(np.argmax(probs))
    confidence = float(probs[pred_class]) * 100.0

    # Build results
    label_results = {CLASS_NAMES[i]: float(probs[i]) for i in range(5)}

    clinical = CLINICAL_ADVICE[pred_class]
    clinical_text = f"""
### {clinical['emoji']} {clinical['title']}

**Confidence:** {confidence:.1f}%
**Urgency:** {clinical['urgency']}

**Clinical Advice:**
{clinical['advice']}
"""

    # Return preprocessed image for visualization
    preprocessed_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    preprocessed_pil = Image.fromarray(preprocessed_rgb)

    return label_results, clinical_text, preprocessed_pil


# ========================== GRADIO UI ==========================

with gr.Blocks(
    title="DR Classification - ResNet-50 Baseline",
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="purple",
    ),
) as demo:
    gr.Markdown(
        """
        # 🔬 Diabetic Retinopathy Classification
        ### ResNet-50 Baseline Model | 5-Class ICDR Standard

        Upload a **retinal fundus image** to classify the severity of Diabetic Retinopathy.

        | Class | Description |
        |-------|-------------|
        | **0 - No DR** | No visible retinopathy |
        | **1 - Mild** | Microaneurysms only |
        | **2 - Moderate** | More than just microaneurysms |
        | **3 - Severe** | Extensive hemorrhages |
        | **4 - Proliferative DR** | Neovascularization / vitreous hemorrhage |
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                type="pil",
                label="Upload Fundus Image",
                height=350,
            )
            predict_btn = gr.Button(
                "🔍 Analyze Image",
                variant="primary",
                size="lg",
            )

        with gr.Column(scale=1):
            output_label = gr.Label(
                label="Classification Probabilities",
                num_top_classes=5,
            )
            output_clinical = gr.Markdown(label="Clinical Guidance")

    with gr.Row():
        output_preprocessed = gr.Image(
            label="Preprocessed Image (Ben Graham Enhanced)",
            height=250,
        )

    predict_btn.click(
        fn=predict,
        inputs=input_image,
        outputs=[output_label, output_clinical, output_preprocessed],
    )

    gr.Examples(
        examples=[],
        inputs=input_image,
        label="Example Images (upload your own fundus images)",
    )

    gr.Markdown(
        """
        ---
        **Model:** ResNet-50 Baseline
        | **Architecture:** ResNet-50
        """
    )

demo.launch()