Rf33d/DR-CKD
0
1"""2DR-CKD Risk Detection — Backend3Run: python app.py model.pth4"""5 6import os, io, sys, base647import numpy as np8import cv29import torch10import torch.nn as nn11import timm12from PIL import Image13from torchvision import transforms14from flask import Flask, request, jsonify, send_file, render_template15from reportlab.lib.pagesizes import A416from reportlab.pdfgen import canvas17from reportlab.lib.utils import ImageReader18import datetime19 20app = Flask(__name__, template_folder='templates', static_folder='static')21app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16 MB max upload22 23# ─────────────────────────────────────────────24# 1. MODEL DEFINITION (exact copy from notebook)25# ─────────────────────────────────────────────26class ConvNeXtDR(nn.Module):27 def __init__(self, num_classes=5, pretrained=False):28 super().__init__()29 self.backbone = timm.create_model(30 'convnext_large_in22ft1k',31 pretrained=pretrained,32 num_classes=033 )34 in_features = self.backbone.num_features # 1536 for convnext_large35 self.head = nn.Sequential(36 nn.LayerNorm(in_features),37 nn.Dropout(0.3),38 nn.Linear(in_features, 512),39 nn.GELU(),40 nn.Dropout(0.2),41 nn.Linear(512, num_classes)42 )43 44 def forward(self, x):45 features = self.backbone(x)46 return self.head(features)47 48 49# ─────────────────────────────────────────────50# 2. GRAD-CAM (exact copy from notebook block 71)51# ─────────────────────────────────────────────52class GradCAM:53 def __init__(self, model, target_layer):54 self.model = model55 self.target_layer = target_layer56 self.gradients = None57 self.activations = None58 self._register_hooks()59 60 def _register_hooks(self):61 def fwd(module, inp, out):62 self.activations = out.detach()63 def bwd(module, grad_in, grad_out):64 self.gradients = grad_out[0].detach()65 self.target_layer.register_forward_hook(fwd)66 self.target_layer.register_full_backward_hook(bwd)67 68 def generate(self, input_tensor, class_idx):69 self.model.zero_grad()70 # Need grad for backward71 inp = input_tensor.clone().detach().requires_grad_(True)72 output = self.model(inp)73 score = output[0, class_idx]74 score.backward()75 weights = self.gradients[0].mean(dim=(1, 2))76 cam = (weights[:, None, None] * self.activations[0]).sum(0)77 cam = torch.relu(cam).cpu().numpy()78 cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)79 return cam80 81 82# ─────────────────────────────────────────────83# 3. PREPROCESSING (exact copy from notebook block 3)84# ─────────────────────────────────────────────85TARGET_SIZE = (224, 224)86 87def crop_black_border(image, tolerance=7):88 gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)89 _, thresh = cv2.threshold(gray, tolerance, 255, cv2.THRESH_BINARY)90 coords = cv2.findNonZero(thresh)91 if coords is None:92 return image93 x, y, w, h = cv2.boundingRect(coords)94 return image[y:y+h, x:x+w]95 96def ben_graham_preprocessing(image, sigmaX=10):97 image = cv2.resize(image, TARGET_SIZE)98 blurred = cv2.GaussianBlur(image, (0, 0), sigmaX)99 result = cv2.addWeighted(image, 4, blurred, -4, 128)100 return result101 102def apply_clahe(image):103 lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)104 l, a, b = cv2.split(lab)105 clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))106 l_eq = clahe.apply(l)107 lab_eq = cv2.merge([l_eq, a, b])108 return cv2.cvtColor(lab_eq, cv2.COLOR_LAB2BGR)109 110val_transform = transforms.Compose([111 transforms.Resize(TARGET_SIZE),112 transforms.ToTensor(),113 transforms.Normalize(mean=[0.485, 0.456, 0.406],114 std=[0.229, 0.224, 0.225])115])116 117def preprocess_image(image_bytes):118 """Run full pipeline. Returns (tensor, display_rgb_array)."""119 nparr = np.frombuffer(image_bytes, np.uint8)120 img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)121 if img_bgr is None:122 return None, None123 img_bgr = crop_black_border(img_bgr)124 img_bgr = ben_graham_preprocessing(img_bgr)125 img_bgr = apply_clahe(img_bgr)126 img_bgr = cv2.resize(img_bgr, TARGET_SIZE)127 img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)128 pil_img = Image.fromarray(img_rgb)129 tensor = val_transform(pil_img).unsqueeze(0)130 return tensor, img_rgb131 132 133# ─────────────────────────────────────────────134# 4. CKD RISK MAP (from notebook block 78)135# ─────────────────────────────────────────────136CKD_RISK_MAP = {137 0: {138 'dr_grade': 'No DR',139 'risk_category': 'Very Low Risk',140 'egfr_range': '>90',141 'risk_percentage': 2.5,142 'ckd_stage': 'G1 — Normal',143 'clinical_action': 'Annual screening. No immediate nephrology referral.'144 },145 1: {146 'dr_grade': 'Mild DR',147 'risk_category': 'Low Risk',148 'egfr_range': '60–89',149 'risk_percentage': 10.0,150 'ckd_stage': 'G2 — Mildly decreased',151 'clinical_action': 'Monitor renal function every 6 months.'152 },153 2: {154 'dr_grade': 'Moderate DR',155 'risk_category': 'Moderate Risk',156 'egfr_range': '45–89',157 'risk_percentage': 35.0,158 'ckd_stage': 'G2–G3a',159 'clinical_action': 'Nephrology review recommended. HbA1c + BP control.'160 },161 3: {162 'dr_grade': 'Severe DR',163 'risk_category': 'High Risk',164 'egfr_range': '20–59',165 'risk_percentage': 60.0,166 'ckd_stage': 'G3b–G4',167 'clinical_action': 'Urgent nephrology referral. Renal function monitoring monthly.'168 },169 4: {170 'dr_grade': 'Proliferative DR',171 'risk_category': 'Very High Risk',172 'egfr_range': '<30',173 'risk_percentage': 85.0,174 'ckd_stage': 'G4–G5',175 'clinical_action': 'Immediate nephrology referral. Consider dialysis planning.'176 }177}178 179RISK_COLORS = {180 'Very Low Risk': '#22c55e',181 'Low Risk': '#84cc16',182 'Moderate Risk': '#f59e0b',183 'High Risk': '#ef4444',184 'Very High Risk': '#7f1d1d'185}186 187 188# ─────────────────────────────────────────────189# 5. FUNDUS VALIDATION (simple heuristic)190# ─────────────────────────────────────────────191def is_fundus_image(image_bytes):192 """193 Simple check: fundus images have dark corners and a194 reddish/orange circular bright centre.195 """196 nparr = np.frombuffer(image_bytes, np.uint8)197 img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)198 if img is None:199 return False200 h, w = img.shape[:2]201 gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)202 203 cx, cy = w // 2, h // 2204 r6h, r6w = max(1, h // 6), max(1, w // 6)205 r8h, r8w = max(1, h // 8), max(1, w // 8)206 207 center_brightness = float(gray[cy-r6h:cy+r6h, cx-r6w:cx+r6w].mean())208 corner_brightness = float((209 gray[:r8h, :r8w].mean() + gray[:r8h, -r8w:].mean() +210 gray[-r8h:, :r8w].mean() + gray[-r8h:, -r8w:].mean()211 ) / 4)212 213 b_ch, g_ch, r_ch = cv2.split(img)214 is_reddish = float(r_ch.mean()) > float(g_ch.mean()) * 0.75215 has_dark_border = center_brightness > corner_brightness * 1.15216 217 return has_dark_border and is_reddish218 219 220# ─────────────────────────────────────────────221# 6. MODEL LOADING222# ─────────────────────────────────────────────223device = torch.device('cpu')224model = None225gradcam_engine = None226 227def load_model(pth_path):228 global model, gradcam_engine229 print(f"Loading model from {pth_path} ...")230 m = ConvNeXtDR(num_classes=5, pretrained=False)231 m.load_state_dict(torch.load(pth_path, map_location=device))232 m.eval()233 model = m234 gradcam_engine = GradCAM(model, model.backbone.stages[-1])235 print("Model ready.")236 237 238# ─────────────────────────────────────────────239# 7. HELPERS240# ─────────────────────────────────────────────241def array_to_b64(arr):242 pil = Image.fromarray(arr.astype(np.uint8))243 buf = io.BytesIO()244 pil.save(buf, format='PNG')245 return base64.b64encode(buf.getvalue()).decode()246 247 248# ─────────────────────────────────────────────249# 8. ROUTES250# ─────────────────────────────────────────────251@app.route('/')252def index():253 return render_template('index.html')254 255 256@app.route('/predict', methods=['POST'])257def predict():258 if 'image' not in request.files:259 return jsonify({'error': 'No image uploaded.'}), 400260 261 img_file = request.files['image']262 image_bytes = img_file.read()263 264 # Validate265 if not is_fundus_image(image_bytes):266 return jsonify({267 'error': 'This does not look like a retinal fundus image. '268 'Please upload a valid fundus photograph.'269 }), 400270 271 if model is None:272 return jsonify({'error': 'Model not loaded. Place Phase 2 .pth as model.pth.'}), 500273 274 # Preprocess275 tensor, img_rgb = preprocess_image(image_bytes)276 if tensor is None:277 return jsonify({'error': 'Could not read image file.'}), 400278 279 tensor = tensor.to(device)280 281 # Inference282 with torch.no_grad():283 output = model(tensor)284 probs = torch.softmax(output, dim=1)[0]285 pred_class = int(probs.argmax().item())286 confidence = float(probs[pred_class].item())287 288 # GradCAM289 cam = gradcam_engine.generate(tensor, pred_class)290 cam_resized = cv2.resize(cam, (img_rgb.shape[1], img_rgb.shape[0]))291 heatmap_bgr = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)292 heatmap_rgb = cv2.cvtColor(heatmap_bgr, cv2.COLOR_BGR2RGB)293 overlay = cv2.addWeighted(img_rgb, 0.6, heatmap_rgb, 0.4, 0)294 295 ckd_info = CKD_RISK_MAP[pred_class]296 297 return jsonify({298 'dr_grade': pred_class,299 'dr_label': ckd_info['dr_grade'],300 'confidence': round(confidence * 100, 1),301 'ckd': ckd_info,302 'risk_color': RISK_COLORS[ckd_info['risk_category']],303 'original_img': array_to_b64(img_rgb),304 'gradcam_img': array_to_b64(overlay),305 'all_probs': [round(float(p) * 100, 1) for p in probs],306 'class_names': ['No DR', 'Mild', 'Moderate', 'Severe', 'Proliferative']307 })308 309 310@app.route('/generate_pdf', methods=['POST'])311def generate_pdf():312 data = request.json313 buf = io.BytesIO()314 c = canvas.Canvas(buf, pagesize=A4)315 pw, ph = A4316 317 # ── Header bar ──318 c.setFillColorRGB(0.04, 0.16, 0.28)319 c.rect(0, ph - 90, pw, 90, fill=1, stroke=0)320 c.setFillColorRGB(1, 1, 1)321 c.setFont("Helvetica-Bold", 17)322 c.drawString(36, ph - 38, "Diabetic Retinopathy — CKD Risk Report")323 c.setFont("Helvetica", 10)324 c.drawString(36, ph - 58, "AI-Assisted Early Detection | Deep Learning — ConvNeXt-Large")325 c.setFont("Helvetica", 9)326 c.drawString(36, ph - 74, f"Generated: {datetime.datetime.now().strftime('%B %d, %Y %H:%M')}")327 328 y = ph - 115329 330 def draw_section(title, rows, y):331 # Section title332 c.setFillColorRGB(0.04, 0.16, 0.28)333 c.setFont("Helvetica-Bold", 12)334 c.drawString(36, y, title)335 y -= 5336 c.setStrokeColorRGB(0.04, 0.16, 0.28)337 c.setLineWidth(0.8)338 c.line(36, y, pw - 36, y)339 y -= 16340 c.setFont("Helvetica", 10)341 for label, value in rows:342 c.setFillColorRGB(0.35, 0.35, 0.35)343 c.setFont("Helvetica-Bold", 10)344 c.drawString(36, y, f"{label}:")345 c.setFillColorRGB(0.1, 0.1, 0.1)346 c.setFont("Helvetica", 10)347 c.drawString(185, y, str(value))348 y -= 18349 return y - 8350 351 ckd = data.get('ckd', {})352 all_probs = data.get('all_probs', [])353 class_names = ['No DR', 'Mild', 'Moderate', 'Severe', 'Proliferative']354 prob_str = ' | '.join([f"{class_names[i]}: {p}%" for i, p in enumerate(all_probs)])355 356 y = draw_section("Prediction Results", [357 ("DR Grade", data.get('dr_label', 'N/A')),358 ("Confidence", f"{data.get('confidence', 'N/A')}%"),359 ("All Probabilities", prob_str),360 ], y)361 362 y = draw_section("CKD Risk Assessment", [363 ("Risk Category", ckd.get('risk_category', 'N/A')),364 ("CKD Risk", f"{ckd.get('risk_percentage', 'N/A')}%"),365 ("Estimated eGFR", ckd.get('egfr_range', 'N/A') + " mL/min/1.73m²"),366 ("CKD Stage", ckd.get('ckd_stage', 'N/A')),367 ("Clinical Action", ckd.get('clinical_action', 'N/A')),368 ], y)369 370 # ── Images ──371 c.setFillColorRGB(0.04, 0.16, 0.28)372 c.setFont("Helvetica-Bold", 12)373 c.drawString(36, y, "GradCAM Analysis")374 y -= 5375 c.line(36, y, pw - 36, y)376 y -= 170377 378 def draw_img(b64_str, x, y, label):379 try:380 img_bytes = base64.b64decode(b64_str)381 img_buf = io.BytesIO(img_bytes)382 reader = ImageReader(img_buf)383 c.drawImage(reader, x, y, width=155, height=155)384 c.setFillColorRGB(0.4, 0.4, 0.4)385 c.setFont("Helvetica", 8)386 c.drawCentredString(x + 77, y - 12, label)387 except Exception:388 pass389 390 draw_img(data.get('original_img', ''), 36, y, "Preprocessed Original")391 draw_img(data.get('gradcam_img', ''), 210, y, "GradCAM Overlay")392 393 # ── Disclaimer ──394 c.setFillColorRGB(0.6, 0.6, 0.6)395 c.setFont("Helvetica-Oblique", 7.5)396 disclaimer = ("This report is generated by an AI model and is intended for research purposes only. "397 "It does not constitute medical advice and must not replace professional clinical diagnosis.")398 c.drawString(36, 28, disclaimer)399 c.setStrokeColorRGB(0.8, 0.8, 0.8)400 c.line(36, 40, pw - 36, 40)401 402 c.save()403 buf.seek(0)404 return send_file(405 buf,406 mimetype='application/pdf',407 as_attachment=True,408 download_name='CKD_Risk_Report.pdf'409 )410 411 412# ─────────────────────────────────────────────413# 9. ENTRY POINT414# ─────────────────────────────────────────────415if __name__ == '__main__':416 pth_path = sys.argv[1] if len(sys.argv) > 1 else 'model.pth'417 if os.path.exists(pth_path):418 load_model(pth_path)419 else:420 print(f"\n⚠ WARNING: '{pth_path}' not found.")421 print(" Place your Phase 2 .pth file as 'model.pth' in this folder, then restart.\n")422 app.run(debug=False, host='0.0.0.0', port=7860)423