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