File size: 3,948 Bytes
d1a4251 4aa2eee 361243a d1a4251 4aa2eee d1a4251 4aa2eee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """
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 # softmax temperature for calibrated confidence; tune on a val split if time allows
# --- Model ---
# Pretrained on APTOS/EyePACS/IDRID/DDR fundus data (timm hub format).
# If this exact repo id doesn't load cleanly, swap for another fundus-pretrained
# timm/HF checkpoint (e.g. search "diabetic retinopathy" on huggingface.co/models)
# and keep everything else (quality gate, Grad-CAM, calibration) unchanged.
MODEL_ID = "hf_hub:ClementP/FundusDRGrading-efficientnet_b0"
model = timm.create_model(MODEL_ID, pretrained=True)
model.eval()
# last conv block for Grad-CAM (EfficientNet: conv_head is the final conv layer)
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]
# Grad-CAM heatmap
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()
|