| """ |
| Eye Screen backend — deploy on Hugging Face Spaces (Gradio SDK). |
| Pipeline: image quality gate -> pretrained DR classifier -> Grad-CAM -> calibrated confidence. |
| This is the "real model" half of the app. Google AI Studio frontend calls this via its API. |
| """ |
|
|
| import spaces |
| import cv2 |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| import timm |
| import gradio as gr |
| from PIL import Image |
| from pytorch_grad_cam import GradCAM |
| from pytorch_grad_cam.utils.image import show_cam_on_image |
|
|
| LABELS = ["No DR", "Mild", "Moderate", "Severe", "Proliferative DR"] |
| TEMPERATURE = 1.8 |
|
|
| |
| |
| |
| |
| |
| MODEL_ID = "hf_hub:ClementP/FundusDRGrading-efficientnet_b0" |
| model = timm.create_model(MODEL_ID, pretrained=True) |
| model.eval() |
|
|
| |
| target_layers = [model.conv_head] if hasattr(model, "conv_head") else [list(model.children())[-3]] |
| cam = GradCAM(model=model, target_layers=target_layers) |
|
|
| IMG_SIZE = 224 |
|
|
|
|
| def check_quality(img_bgr): |
| gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY) |
| blur_score = cv2.Laplacian(gray, cv2.CV_64F).var() |
| brightness = gray.mean() |
|
|
| issues = [] |
| if blur_score < 100: |
| issues.append("Image appears blurry/out of focus") |
| if brightness < 40: |
| issues.append("Image is underexposed (too dark)") |
| if brightness > 210: |
| issues.append("Image is overexposed (too bright)") |
|
|
| return issues, {"blur_score": round(float(blur_score), 1), "brightness": round(float(brightness), 1)} |
|
|
|
|
| def preprocess(img_pil): |
| img = img_pil.resize((IMG_SIZE, IMG_SIZE)) |
| arr = np.array(img).astype(np.float32) / 255.0 |
| mean = np.array([0.485, 0.456, 0.406]) |
| std = np.array([0.229, 0.224, 0.225]) |
| arr = (arr - mean) / std |
| tensor = torch.tensor(arr.transpose(2, 0, 1), dtype=torch.float32).unsqueeze(0) |
| return tensor, arr |
|
|
|
|
| @spaces.GPU(duration=30) |
| def predict(image): |
| if image is None: |
| return "No image uploaded", "", None, {} |
|
|
| img_pil = Image.fromarray(image).convert("RGB") |
| img_bgr = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) |
|
|
| issues, metrics = check_quality(img_bgr) |
| if issues: |
| msg = "RETAKE REQUESTED: " + "; ".join(issues) |
| return msg, "N/A — quality check failed before classification", None, metrics |
|
|
| tensor, norm_img = preprocess(img_pil) |
|
|
| with torch.no_grad(): |
| logits = model(tensor) |
| calibrated_probs = F.softmax(logits / TEMPERATURE, dim=1)[0] |
|
|
| pred_idx = int(torch.argmax(calibrated_probs)) |
| confidence = float(calibrated_probs[pred_idx]) * 100 |
| severity = LABELS[pred_idx] |
|
|
| |
| grayscale_cam = cam(input_tensor=tensor, targets=None)[0] |
| display_img = np.array(img_pil.resize((IMG_SIZE, IMG_SIZE))).astype(np.float32) / 255.0 |
| overlay = show_cam_on_image(display_img, grayscale_cam, use_rgb=True) |
|
|
| result_text = f"{severity} ({confidence:.1f}% confidence)" |
| disclaimer = "Screening aid only — not a diagnosis. Consult an ophthalmologist for confirmation." |
| return result_text, disclaimer, overlay, metrics |
|
|
|
|
| demo = gr.Interface( |
| fn=predict, |
| inputs=gr.Image(label="Fundus photo"), |
| outputs=[ |
| gr.Textbox(label="Result"), |
| gr.Textbox(label="Note"), |
| gr.Image(label="Grad-CAM heatmap"), |
| gr.JSON(label="Quality metrics (debug)"), |
| ], |
| title="Eye Screen — DR Grading API", |
| description="Upload a retinal fundus image. Rejects low-quality photos before grading.", |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
|
|
