Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import timm | |
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| from PIL import Image | |
| import requests | |
| import json | |
| import fitz | |
| import pytesseract | |
| import re | |
| import os | |
| import pickle | |
| import pandas as pd | |
| from sklearn.preprocessing import LabelEncoder | |
| from sklearn.impute import SimpleImputer | |
| from huggingface_hub import hf_hub_download | |
| import joblib | |
| import shap | |
| # ── Device ──────────────────────────────────────────────────────────────────── | |
| device = torch.device("cpu") | |
| REPO_ID = "Dina-Raslan/ckd-retinal-model" | |
| # ── UCI Column Definitions ──────────────────────────────────────────────────── | |
| NUMERIC_COLS = ["age", "bp", "bgr", "bu", "sc", "sod", "pot", "hemo", "pcv", "wbcc", "rbcc"] | |
| CATEGORICAL_COLS = ["sg", "al", "su", "rbc", "pc", "pcc", "ba", "htn", "dm", "cad", "appet", "pe", "ane"] | |
| ALL_COLS = NUMERIC_COLS + CATEGORICAL_COLS | |
| # القيم الممكنة لكل عمود categorical - لازم تكون ثابتة علشان LabelEncoder يشتغل صح | |
| LABEL_ENCODER_CLASSES = { | |
| "sg": ["1.005", "1.010", "1.015", "1.020", "1.025"], | |
| "al": ["0", "1", "2", "3", "4", "5"], | |
| "su": ["0", "1", "2", "3", "4", "5"], | |
| "rbc": ["abnormal", "normal"], | |
| "pc": ["abnormal", "normal"], | |
| "pcc": ["notpresent", "present"], | |
| "ba": ["notpresent", "present"], | |
| "htn": ["no", "yes"], | |
| "dm": ["no", "yes"], | |
| "cad": ["no", "yes"], | |
| "appet": ["good", "poor"], | |
| "pe": ["no", "yes"], | |
| "ane": ["no", "yes"], | |
| } | |
| def encode_categoricals(df): | |
| """Encode categorical columns consistently using fixed classes""" | |
| for col in CATEGORICAL_COLS: | |
| le = LabelEncoder() | |
| le.classes_ = np.array(LABEL_ENCODER_CLASSES[col]) | |
| val = str(df[col].iloc[0]).strip() | |
| # لو القيمة مش موجودة في الـ classes خد أقرب قيمة | |
| if val not in LABEL_ENCODER_CLASSES[col]: | |
| val = LABEL_ENCODER_CLASSES[col][0] | |
| df[col] = le.transform([val]) | |
| return df | |
| # ── Retinal Model ───────────────────────────────────────────────────────────── | |
| class RetinalModelV2(nn.Module): | |
| def __init__(self, num_classes=2): | |
| super().__init__() | |
| self.backbone = timm.create_model( | |
| "efficientnet_b3", pretrained=False, | |
| num_classes=0, global_pool="avg", drop_rate=0.3 | |
| ) | |
| self.feature_dim = self.backbone.num_features | |
| self.classifier = nn.Sequential( | |
| nn.Linear(self.feature_dim, 512), nn.BatchNorm1d(512), | |
| nn.ReLU(), nn.Dropout(0.5), | |
| nn.Linear(512, 256), nn.BatchNorm1d(256), | |
| nn.ReLU(), nn.Dropout(0.4), | |
| nn.Linear(256, num_classes) | |
| ) | |
| def forward(self, x): | |
| features = self.backbone(x) | |
| return self.classifier(features), features | |
| # ── Fusion Model ────────────────────────────────────────────────────────────── | |
| class CrossModalAttentionFusion(nn.Module): | |
| def __init__(self, retinal_dim=64, clinical_dim=24, hidden_dim=128, num_classes=2): | |
| super().__init__() | |
| self.retinal_proj = nn.Sequential( | |
| nn.Linear(retinal_dim, hidden_dim), nn.LayerNorm(hidden_dim), | |
| nn.ReLU(), nn.Dropout(0.3) | |
| ) | |
| self.clinical_proj = nn.Sequential( | |
| nn.Linear(clinical_dim, hidden_dim), nn.LayerNorm(hidden_dim), | |
| nn.ReLU(), nn.Dropout(0.3) | |
| ) | |
| self.cross_attention = nn.MultiheadAttention( | |
| embed_dim=hidden_dim, num_heads=4, dropout=0.1, batch_first=True | |
| ) | |
| self.self_attention = nn.MultiheadAttention( | |
| embed_dim=hidden_dim, num_heads=4, dropout=0.1, batch_first=True | |
| ) | |
| self.classifier = nn.Sequential( | |
| nn.Linear(hidden_dim * 2, 128), nn.BatchNorm1d(128), | |
| nn.ReLU(), nn.Dropout(0.4), | |
| nn.Linear(128, 64), nn.ReLU(), | |
| nn.Dropout(0.3), nn.Linear(64, num_classes) | |
| ) | |
| def forward(self, retinal, clinical): | |
| r = self.retinal_proj(retinal).unsqueeze(1) | |
| c = self.clinical_proj(clinical).unsqueeze(1) | |
| r_attended, _ = self.cross_attention(query=r, key=c, value=c) | |
| c_attended, _ = self.self_attention(query=c, key=c, value=c) | |
| fused = torch.cat([r_attended.squeeze(1), c_attended.squeeze(1)], dim=1) | |
| return self.classifier(fused) | |
| # ── Load All Models ─────────────────────────────────────────────────────────── | |
| print("Loading models...") | |
| retinal_model_path = hf_hub_download(repo_id=REPO_ID, filename="best_model.pth") | |
| retinal_model = RetinalModelV2(num_classes=2) | |
| retinal_model.load_state_dict(torch.load(retinal_model_path, map_location="cpu"), strict=False) | |
| retinal_model.eval() | |
| print("Retinal model loaded") | |
| fusion_model_path = hf_hub_download(repo_id=REPO_ID, filename="best_fusion_model.pth") | |
| fusion_model = CrossModalAttentionFusion(retinal_dim=64, clinical_dim=24, hidden_dim=128, num_classes=2) | |
| fusion_model.load_state_dict(torch.load(fusion_model_path, map_location="cpu")) | |
| fusion_model.eval() | |
| print("Fusion model loaded") | |
| pca_path = hf_hub_download(repo_id=REPO_ID, filename="pca_model.pkl") | |
| with open(pca_path, "rb") as f: | |
| pca_model = pickle.load(f) | |
| print("PCA model loaded") | |
| scaler_path = hf_hub_download(repo_id=REPO_ID, filename="clinical_scaler.pkl") | |
| with open(scaler_path, "rb") as f: | |
| clinical_scaler = pickle.load(f) | |
| print("Clinical scaler loaded") | |
| gb_path = hf_hub_download(repo_id=REPO_ID, filename="gb_clinical_model_joblib.pkl") | |
| gb_model = joblib.load(gb_path) | |
| print("GB clinical model loaded") | |
| # SHAP explainer for the clinical model (created once, reused for every request) | |
| shap_explainer = shap.TreeExplainer(gb_model) | |
| print("SHAP explainer ready") | |
| print("All models ready.") | |
| # تحقق من اسماء الأعمدة في الـ scaler | |
| if hasattr(clinical_scaler, 'feature_names_in_'): | |
| print("Scaler trained on:", list(clinical_scaler.feature_names_in_)) | |
| else: | |
| print("Scaler has no feature_names_in_ - will use numpy array directly") | |
| OPTIMAL_THRESHOLD = 0.50 | |
| # ── Preprocessing ───────────────────────────────────────────────────────────── | |
| def remove_black_border(img): | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| _, thresh = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY) | |
| contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| if contours: | |
| x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea)) | |
| img = img[y:y+h, x:x+w] | |
| return img | |
| def lab_ace_enhancement(img): | |
| lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB) | |
| l, a, b = cv2.split(lab) | |
| clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) | |
| l = clahe.apply(l) | |
| return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2BGR) | |
| def graham_preprocessing(img): | |
| return cv2.addWeighted(img, 4, cv2.GaussianBlur(img, (0, 0), 10), -4, 128) | |
| def green_channel_enhancement(img): | |
| b, g, r = cv2.split(img) | |
| g = cv2.normalize(g, None, 0, 255, cv2.NORM_MINMAX) | |
| return cv2.merge([b, g, r]) | |
| def preprocess_retinal(img_pil, target_size=512): | |
| img = cv2.cvtColor(np.array(img_pil.convert("RGB")), cv2.COLOR_RGB2BGR) | |
| img = remove_black_border(img) | |
| img = cv2.resize(img, (target_size, target_size)) | |
| img = lab_ace_enhancement(img) | |
| img = graham_preprocessing(img) | |
| img = green_channel_enhancement(img) | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| return img.astype(np.float32) / 255.0 | |
| # ── Clinical Feature Preparation ────────────────────────────────────────────── | |
| def prepare_clinical_features( | |
| # من التحاليل | |
| sc, bu, hemo, bgr, sod, pot, pcv, wbcc, rbcc, | |
| sg, al, su, rbc, pc, pcc, ba, | |
| # مانيوال | |
| age, bp, htn, dm, cad, appet, pe, ane | |
| ): | |
| """ | |
| تحضير الـ 24 feature بالترتيب الصح للـ scaler | |
| """ | |
| row = { | |
| # Numeric | |
| "age": float(age), | |
| "bp": float(bp), | |
| "bgr": float(bgr), | |
| "bu": float(bu), | |
| "sc": float(sc), | |
| "sod": float(sod), | |
| "pot": float(pot), | |
| "hemo": float(hemo), | |
| "pcv": float(pcv), | |
| "wbcc": float(wbcc), | |
| "rbcc": float(rbcc), | |
| # Categorical | |
| "sg": str(sg), | |
| "al": str(al), | |
| "su": str(su), | |
| "rbc": str(rbc), | |
| "pc": str(pc), | |
| "pcc": str(pcc), | |
| "ba": str(ba), | |
| "htn": str(htn), | |
| "dm": str(dm), | |
| "cad": str(cad), | |
| "appet": str(appet), | |
| "pe": str(pe), | |
| "ane": str(ane), | |
| } | |
| df = pd.DataFrame([row]) | |
| # Encode categoricals بطريقة صح (fixed classes) | |
| df = encode_categoricals(df) | |
| # رتب الأعمدة بالترتيب الصح | |
| df = df[ALL_COLS] | |
| # اعمل numpy array مباشرة علشان نتجنب مشكلة feature names | |
| X = df.values.astype(np.float64) | |
| # Scale | |
| scaled = clinical_scaler.transform(X) | |
| return scaled | |
| def get_top_shap_features(clinical_feat, top_n=6): | |
| """ | |
| Compute SHAP values for the clinical (GB) model's prediction and | |
| return the top_n features ranked by absolute impact, together with | |
| each feature's percentage share of total |SHAP| impact and the | |
| combined percentage share of all remaining features. | |
| clinical_feat must be the scaled numpy array in ALL_COLS order. | |
| Returns: | |
| top_features: list of (feature_name, shap_value, pct_impact) tuples | |
| other_pct: combined percentage impact of the remaining features | |
| """ | |
| shap_values = shap_explainer.shap_values(clinical_feat) | |
| # Some SHAP/model versions return a list per class, others a single array | |
| if isinstance(shap_values, list): | |
| shap_row = shap_values[1][0] # class "CKD" (positive class) | |
| else: | |
| shap_row = shap_values[0] | |
| abs_vals = np.abs(shap_row) | |
| total_abs = abs_vals.sum() | |
| # Avoid division by zero in the (unlikely) case all SHAP values are 0 | |
| if total_abs == 0: | |
| pct_vals = np.zeros_like(abs_vals) | |
| else: | |
| pct_vals = (abs_vals / total_abs) * 100.0 | |
| sorted_idx = np.argsort(abs_vals)[::-1] | |
| top_idx = sorted_idx[:top_n] | |
| other_idx = sorted_idx[top_n:] | |
| top_features = [ | |
| (ALL_COLS[i], float(shap_row[i]), float(pct_vals[i])) | |
| for i in top_idx | |
| ] | |
| other_pct = float(pct_vals[other_idx].sum()) if len(other_idx) else 0.0 | |
| return top_features, other_pct | |
| # ── Grad-CAM++ ──────────────────────────────────────────────────────────────── | |
| class GradCAMPlusPlus: | |
| def __init__(self, model, target_layer): | |
| self.model = model | |
| self.gradients = None | |
| self.activations = None | |
| target_layer.register_forward_hook(self._save_activation) | |
| target_layer.register_full_backward_hook(self._save_gradient) | |
| def _save_activation(self, module, input, output): | |
| self.activations = output.detach() | |
| def _save_gradient(self, module, grad_input, grad_output): | |
| self.gradients = grad_output[0].detach() | |
| def generate(self, img_tensor, class_idx=None): | |
| self.model.eval() | |
| img_tensor = img_tensor.unsqueeze(0) | |
| logits, _ = self.model(img_tensor) | |
| if class_idx is None: | |
| class_idx = logits.argmax(dim=1).item() | |
| self.model.zero_grad() | |
| logits[0, class_idx].backward() | |
| grads = self.gradients | |
| acts = self.activations | |
| grads_sq = grads ** 2 | |
| grads_cub = grads ** 3 | |
| denom = 2 * grads_sq + acts * grads_cub | |
| denom = torch.where(denom != 0, denom, torch.ones_like(denom)) | |
| alpha = grads_sq / denom | |
| weights = (alpha * F.relu(grads)).sum(dim=[2, 3], keepdim=True) | |
| cam = (weights * acts).sum(dim=1, keepdim=True) | |
| cam = F.relu(cam) | |
| cam = F.interpolate(cam, size=(512, 512), mode="bilinear", align_corners=False) | |
| cam = cam.squeeze().cpu().numpy() | |
| cam -= cam.min() | |
| if cam.max() > 0: | |
| cam /= cam.max() | |
| return cam | |
| target_layer = retinal_model.backbone.blocks[6][1].conv_pwl | |
| gradcam = GradCAMPlusPlus(retinal_model, target_layer) | |
| def make_retinal_mask(size=512): | |
| mask = np.zeros((size, size), dtype=np.float32) | |
| cv2.circle(mask, (size//2, size//2), int(size * 0.47), 1.0, -1) | |
| return mask | |
| # ── OCR ─────────────────────────────────────────────────────────────────────── | |
| LAB_PATTERNS = { | |
| # Blood Tests | |
| "sc": [r"creatinine[\s\S]{0,40}?(\d+\.?\d*)\s*mg", r"s\.?\s*creat[\s\S]{0,30}?(\d+\.?\d*)"], | |
| "bu": [r"(?:blood\s*urea|urea|bun)[\s\S]{0,30}?(\d+\.?\d*)\s*mg", r"b\.?u\.?n[\s\S]{0,20}?(\d+\.?\d*)"], | |
| "hemo": [r"h[ae]moglobin[\s\S]{0,30}?(\d+\.?\d*)\s*g", r"hgb[\s\S]{0,20}?(\d+\.?\d*)"], | |
| "bgr": [r"(?:blood\s*glucose|glucose|random\s*blood\s*sugar|rbs)[\s\S]{0,30}?(\d+\.?\d*)\s*mg", | |
| r"fbs[\s\S]{0,20}?(\d+\.?\d*)"], | |
| "sod": [r"sodium[\s\S]{0,30}?(\d+\.?\d*)\s*m?eq", r"na\+?[\s\S]{0,20}?(\d+\.?\d*)"], | |
| "pot": [r"potassium[\s\S]{0,30}?(\d+\.?\d*)\s*m?eq", r"k\+?[\s\S]{0,20}?(\d+\.?\d*)"], | |
| "pcv": [r"(?:packed\s*cell\s*volume|hematocrit|pcv|hct)[\s\S]{0,30}?(\d+\.?\d*)\s*%?"], | |
| "wbcc": [r"(?:white\s*blood\s*cell|wbc|leukocyte)[\s\S]{0,30}?(\d+[\d,]*\.?\d*)", | |
| r"wbcc[\s\S]{0,20}?(\d+[\d,]*\.?\d*)"], | |
| "rbcc": [r"(?:red\s*blood\s*cell|rbc\s*count|erythrocyte)[\s\S]{0,30}?(\d+\.?\d*)", | |
| r"rbcc[\s\S]{0,20}?(\d+\.?\d*)"], | |
| # Urine Tests | |
| "sg": [r"(?:specific\s*gravity|sp\.?\s*gr\.?)[\s\S]{0,30}?(1\.\d{3})"], | |
| "al": [r"(?:albumin|protein)[\s\S]{0,30}?(\d)\s*\+?", r"albumin[\s\S]{0,20}?(\d)"], | |
| "su": [r"(?:sugar|glucose)\s*(?:in\s*urine|urine)[\s\S]{0,30}?(\d)", r"glycosuria[\s\S]{0,20}?(\d)"], | |
| "bp": [r"(?:blood\s*pressure|bp|diastolic)[\s\S]{0,30}?(\d{2,3})\s*mm", | |
| r"\d{2,3}\s*/\s*(\d{2,3})\s*mm"], | |
| } | |
| VALUE_RANGES = { | |
| "sc": (0.1, 20), "bu": (1, 300), | |
| "hemo": (1, 25), "bgr": (30, 600), | |
| "sod": (100, 170), "pot": (1, 10), | |
| "pcv": (10, 60), "wbcc": (2000, 30000), | |
| "rbcc": (1, 8), "sg": (1.000, 1.030), | |
| "al": (0, 5), "su": (0, 5), | |
| "bp": (50, 250), | |
| } | |
| def parse_lab_values(text): | |
| text_lower = text.lower() | |
| values = {} | |
| for key, patterns in LAB_PATTERNS.items(): | |
| for pattern in patterns: | |
| match = re.search(pattern, text_lower) | |
| if match: | |
| try: | |
| val_str = match.group(1).replace(",", "") | |
| val = float(val_str) | |
| lo, hi = VALUE_RANGES.get(key, (0, 9999)) | |
| if lo <= val <= hi: | |
| values[key] = val | |
| break | |
| except (ValueError, IndexError): | |
| pass | |
| # استنتاج categorical من النص | |
| # RBC in urine | |
| if re.search(r"rbc\s*(?:in\s*urine|urine)[\s\S]{0,30}?abnormal", text_lower): | |
| values["rbc"] = "abnormal" | |
| elif re.search(r"rbc\s*(?:in\s*urine|urine)[\s\S]{0,30}?normal", text_lower): | |
| values["rbc"] = "normal" | |
| # Pus Cells | |
| if re.search(r"pus\s*cells?[\s\S]{0,30}?abnormal", text_lower): | |
| values["pc"] = "abnormal" | |
| elif re.search(r"pus\s*cells?[\s\S]{0,30}?normal", text_lower): | |
| values["pc"] = "normal" | |
| # Pus Cell Clumps | |
| if re.search(r"pus\s*cell\s*clumps?[\s\S]{0,30}?present", text_lower): | |
| values["pcc"] = "present" | |
| elif re.search(r"pus\s*cell\s*clumps?[\s\S]{0,30}?not\s*present", text_lower): | |
| values["pcc"] = "notpresent" | |
| # Bacteria | |
| if re.search(r"bacteri[\s\S]{0,30}?present", text_lower): | |
| values["ba"] = "present" | |
| elif re.search(r"bacteri[\s\S]{0,30}?not\s*present", text_lower): | |
| values["ba"] = "notpresent" | |
| # Specific Gravity - snap to nearest valid value | |
| if "sg" in values: | |
| valid_sg = [1.005, 1.010, 1.015, 1.020, 1.025] | |
| sg_val = values["sg"] | |
| nearest = min(valid_sg, key=lambda x: abs(x - sg_val)) | |
| values["sg"] = f"{nearest:.3f}" | |
| # Albumin و Sugar -> snap to integer string | |
| for col in ["al", "su"]: | |
| if col in values: | |
| values[col] = str(int(round(values[col]))) | |
| return values | |
| def process_lab_file(lab_file): | |
| if lab_file is None: | |
| return ("No file uploaded.", | |
| gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), | |
| gr.update(), gr.update(), gr.update(), gr.update(), | |
| gr.update(), gr.update(), gr.update()) | |
| path = lab_file.name if hasattr(lab_file, "name") else lab_file | |
| if path.lower().endswith(".pdf"): | |
| doc = fitz.open(path) | |
| text = "".join(page.get_text() for page in doc) | |
| else: | |
| img = Image.open(path) | |
| img_np = np.array(img.convert("RGB")) | |
| img_cv = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) | |
| img_cv = cv2.resize(img_cv, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC) | |
| gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY) | |
| gray = cv2.fastNlMeansDenoising(gray, h=10) | |
| _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| text = pytesseract.image_to_string(thresh, config="--oem 3 --psm 6") | |
| values = parse_lab_values(text) | |
| if values: | |
| found = [k for k in values] | |
| status = f" Extracted {len(found)} values: {', '.join(found)}\nPlease verify and fill missing fields manually." | |
| else: | |
| status = "️ Could not extract values automatically.\nPlease enter all values manually." | |
| def upd(key, default): | |
| return gr.update(value=values[key]) if key in values else gr.update(value=default) | |
| def upd_cat(key, default): | |
| return gr.update(value=values[key]) if key in values else gr.update(value=default) | |
| return ( | |
| status, | |
| upd("sc", 1.2), | |
| upd("bu", 30), | |
| upd("hemo", 13), | |
| upd("bgr", 100), | |
| upd("sod", 135), | |
| upd("pot", 4.0), | |
| upd("pcv", 39), | |
| upd("wbcc", 7500), | |
| upd("rbcc", 4.5), | |
| upd("bp", 80), | |
| upd_cat("sg", "1.020"), | |
| upd_cat("al", "0"), | |
| upd_cat("su", "0"), | |
| upd_cat("rbc", "normal"), | |
| upd_cat("pc", "normal"), | |
| upd_cat("pcc", "notpresent"), | |
| upd_cat("ba", "notpresent"), | |
| ) | |
| # ── Main Prediction ─────────────────────────────────────────────────────────── | |
| def predict_ckd_risk( | |
| retinal_image, | |
| # Blood tests | |
| sc, bu, hemo, bgr, sod, pot, pcv, wbcc, rbcc, | |
| # Blood pressure | |
| bp, | |
| # Urine tests (categorical) | |
| sg, al, su, rbc, pc, pcc, ba, | |
| # Manual entry | |
| age, htn, dm, cad, appet, pe, ane | |
| ): | |
| if retinal_image is None: | |
| return None, "️ Please upload a retinal image.", "", "" | |
| # 1. Preprocess retinal image | |
| img_processed = preprocess_retinal(retinal_image) | |
| img_tensor = torch.tensor(img_processed).permute(2, 0, 1) | |
| # 2. Extract retinal features | |
| with torch.no_grad(): | |
| logits, retinal_features = retinal_model(img_tensor.unsqueeze(0)) | |
| retinal_prob = torch.softmax(logits, dim=1)[0, 1].item() | |
| # 3. PCA on retinal features | |
| retinal_feat_np = retinal_features.cpu().numpy() | |
| retinal_feat_pca = pca_model.transform(retinal_feat_np) | |
| # 4. Prepare clinical features (24 columns, correct order) | |
| clinical_feat = prepare_clinical_features( | |
| sc=sc, bu=bu, hemo=hemo, bgr=bgr, sod=sod, pot=pot, | |
| pcv=pcv, wbcc=wbcc, rbcc=rbcc, | |
| sg=sg, al=al, su=su, rbc=rbc, pc=pc, pcc=pcc, ba=ba, | |
| age=age, bp=bp, | |
| htn=htn, dm=dm, cad=cad, appet=appet, pe=pe, ane=ane | |
| ) | |
| # 5. GB clinical prediction | |
| clinical_prob = gb_model.predict_proba(clinical_feat)[0, 1] | |
| # 5.b Top-6 SHAP features driving the clinical prediction (with % impact) | |
| top_shap_features, other_shap_pct = get_top_shap_features(clinical_feat, top_n=6) | |
| # 6. Fusion model prediction | |
| retinal_tensor = torch.tensor(retinal_feat_pca, dtype=torch.float32) | |
| clinical_tensor = torch.tensor(clinical_feat, dtype=torch.float32) | |
| with torch.no_grad(): | |
| fusion_logits = fusion_model(retinal_tensor, clinical_tensor) | |
| fusion_prob = torch.softmax(fusion_logits, dim=1)[0, 1].item() | |
| # 7. Final risk | |
| if fusion_prob >= 0.6: | |
| risk_label = " HIGH RISK" | |
| risk_color = "red" | |
| elif fusion_prob >= 0.35: | |
| risk_label = " MEDIUM RISK" | |
| risk_color = "orange" | |
| else: | |
| risk_label = " LOW RISK" | |
| risk_color = "green" | |
| # 8. Grad-CAM++ | |
| cam = gradcam.generate(img_tensor) | |
| retinal_mask = make_retinal_mask(512) | |
| cam = cam * retinal_mask | |
| cam -= cam.min() | |
| if cam.max() > 0: | |
| cam /= cam.max() | |
| img_rgb = (img_processed * 255).astype(np.uint8) | |
| heatmap = cv2.applyColorMap((cam * 255).astype(np.uint8), cv2.COLORMAP_JET) | |
| heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) | |
| heat_mask = (cam > 0.3).astype(np.float32) | |
| heat_mask = np.stack([heat_mask] * 3, axis=-1) | |
| overlay = (img_rgb * (1 - 0.6 * heat_mask) + heatmap * 0.6 * heat_mask).astype(np.uint8) | |
| circ_mask = np.stack([retinal_mask] * 3, axis=-1) | |
| overlay = (overlay * circ_mask).astype(np.uint8) | |
| fig, axes = plt.subplots(1, 3, figsize=(15, 5)) | |
| axes[0].imshow(img_rgb); axes[0].set_title("Preprocessed Retinal Image"); axes[0].axis("off") | |
| axes[1].imshow(cam, cmap="jet"); axes[1].set_title("Grad-CAM++ Heatmap"); axes[1].axis("off") | |
| axes[2].imshow(overlay); axes[2].set_title("Overlay"); axes[2].axis("off") | |
| plt.suptitle( | |
| f"CKD Risk: {risk_label} | Fusion Score: {fusion_prob:.3f}", | |
| fontsize=13, color=risk_color, fontweight="bold" | |
| ) | |
| plt.tight_layout() | |
| gradcam_path = "/tmp/gradcam_result.png" | |
| plt.savefig(gradcam_path, dpi=120) | |
| plt.close() | |
| clinical_summary = f""" | |
| ╔══════════════════════════════════════╗ | |
| ║ CLINICAL VALUES ║ | |
| ╠══════════════════════════════════════╣ | |
| BLOOD TESTS | |
| Creatinine (sc): {sc} mg/dL | |
| Urea/BUN (bu): {bu} mg/dL | |
| Hemoglobin (hemo): {hemo} g/dL | |
| Blood Glucose (bgr): {bgr} mg/dL | |
| Sodium (sod): {sod} mEq/L | |
| Potassium (pot): {pot} mEq/L | |
| Packed Cell Vol (pcv):{pcv} % | |
| WBC Count (wbcc): {wbcc} cells/cumm | |
| RBC Count (rbcc): {rbcc} millions/cumm | |
| Blood Pressure (bp): {bp} mmHg | |
| URINE TESTS | |
| Specific Gravity (sg):{sg} | |
| Albumin (al): {al} | |
| Sugar (su): {su} | |
| RBC in Urine (rbc): {rbc} | |
| Pus Cells (pc): {pc} | |
| Pus Cell Clumps (pcc):{pcc} | |
| Bacteria (ba): {ba} | |
| PATIENT INFO | |
| Age: {age} years | |
| Hypertension (htn): {htn} | |
| Diabetes (dm): {dm} | |
| Coronary Art. (cad): {cad} | |
| Appetite (appet): {appet} | |
| Pedal Edema (pe): {pe} | |
| Anemia (ane): {ane} | |
| ╠══════════════════════════════════════╣ | |
| AI PREDICTION BREAKDOWN | |
| Retinal Score: {retinal_prob:.4f} | |
| Clinical Score: {clinical_prob:.4f} | |
| Fusion Score: {fusion_prob:.4f} | |
| Risk Level: {risk_label} | |
| ╠══════════════════════════════════════╣ | |
| TOP 6 CLINICAL FEATURES (SHAP) | |
| {chr(10).join(f" {i+1}. {name:<8s} {'+' if val >= 0 else ''}{val:.4f} ({pct:.1f}% impact)" for i, (name, val, pct) in enumerate(top_shap_features))} | |
| Other 18 factors combined: {other_shap_pct:.1f}% impact | |
| ╚══════════════════════════════════════╝ | |
| """ | |
| risk_display = f"## {risk_label}\n\n**Fusion Score:** {fusion_prob:.4f}\n\n**Retinal:** {retinal_prob:.4f} | **Clinical:** {clinical_prob:.4f}" | |
| return gradcam_path, risk_display, clinical_summary, f"Done. Risk: {risk_label}" | |
| # ── Chatbot ─────────────────────────────────────────────────────────────────── | |
| GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "") | |
| SYSTEM_PROMPT = """You are a CKD Health Coach, a specialized AI assistant for patients | |
| with Chronic Kidney Disease (CKD) or those at risk. | |
| - Provide lifestyle and dietary recommendations specific to CKD | |
| - Offer fluid intake and exercise advice | |
| - Answer basic kidney health questions | |
| - NEVER diagnose or replace a doctor | |
| - For severe symptoms, recommend emergency care immediately""" | |
| EMERGENCY_KEYWORDS = ["chest pain", "difficulty breathing", "severe swelling", | |
| "cannot breathe", "unconscious", "bleeding", "fainted", "heart attack"] | |
| def check_emergency(message): | |
| return any(kw in message.lower() for kw in EMERGENCY_KEYWORDS) | |
| def chat_with_groq(message, history): | |
| if check_emergency(message): | |
| return " This sounds like a medical emergency. Please call emergency services immediately." | |
| messages = [{"role": "system", "content": SYSTEM_PROMPT}] | |
| for item in history: | |
| if isinstance(item, dict): | |
| messages.append({"role": item["role"], "content": item["content"]}) | |
| elif isinstance(item, (list, tuple)) and len(item) == 2: | |
| messages.append({"role": "user", "content": item[0]}) | |
| messages.append({"role": "assistant", "content": item[1]}) | |
| messages.append({"role": "user", "content": message}) | |
| if not GROQ_API_KEY: | |
| return "Chatbot unavailable: GROQ_API_KEY not set in Space secrets." | |
| try: | |
| response = requests.post( | |
| "https://api.groq.com/openai/v1/chat/completions", | |
| headers={"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}, | |
| json={"model": "llama-3.3-70b-versatile", "messages": messages, | |
| "max_tokens": 512, "temperature": 0.7}, | |
| timeout=30 | |
| ) | |
| return response.json()["choices"][0]["message"]["content"] | |
| except Exception as e: | |
| return f"Error: {str(e)}" | |
| # ── Gradio UI ───────────────────────────────────────────────────────────────── | |
| with gr.Blocks(title="CKD-MultiVision") as demo: | |
| gr.Markdown(""" | |
| # CKD-MultiVision | |
| ### Multimodal AI System for Early Chronic Kidney Disease Detection | |
| *Retinal Image + Clinical Lab Values → AI Risk Assessment* | |
| """) | |
| with gr.Tabs(): | |
| # ── Tab 1: Risk Prediction ───────────────────────────────────────── | |
| with gr.TabItem(" Risk Prediction"): | |
| # Lab file upload | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("### Upload Lab Report (Optional)") | |
| lab_file_input = gr.File( | |
| label="Upload Lab Report (PDF or Image)", | |
| file_types=[".pdf", ".png", ".jpg", ".jpeg"] | |
| ) | |
| extract_btn = gr.Button(" Extract Values from File", variant="secondary") | |
| extract_status = gr.Textbox(label="Extraction Status", lines=3, interactive=False) | |
| gr.Markdown("---") | |
| with gr.Row(): | |
| # Left: Retinal Image | |
| with gr.Column(scale=1): | |
| gr.Markdown("### ️ Retinal Fundus Image") | |
| retinal_input = gr.Image( | |
| label="Upload Retinal Image (JPG/PNG)", | |
| type="pil", height=300 | |
| ) | |
| # Middle: Blood Tests | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Blood Tests") | |
| sc_input = gr.Number(label="Serum Creatinine - sc (mg/dL)", value=1.2, minimum=0.1, maximum=20) | |
| bu_input = gr.Number(label="Blood Urea/BUN - bu (mg/dL)", value=30, minimum=1, maximum=300) | |
| hemo_input = gr.Number(label="Hemoglobin - hemo (g/dL)", value=13, minimum=1, maximum=25) | |
| bgr_input = gr.Number(label="Blood Glucose Random - bgr (mg/dL)", value=100, minimum=30, maximum=600) | |
| sod_input = gr.Number(label="Sodium - sod (mEq/L)", value=135, minimum=100, maximum=170) | |
| pot_input = gr.Number(label="Potassium - pot (mEq/L)", value=4.0, minimum=1, maximum=10) | |
| pcv_input = gr.Number(label="Packed Cell Volume/Hematocrit - pcv (%)", value=39, minimum=10, maximum=60) | |
| wbcc_input = gr.Number(label="WBC Count - wbcc (cells/cumm)", value=7500, minimum=2000, maximum=30000) | |
| rbcc_input = gr.Number(label="RBC Count - rbcc (millions/cumm)", value=4.5, minimum=1, maximum=8) | |
| bp_input = gr.Number(label="Blood Pressure Diastolic - bp (mmHg)", value=80, minimum=50, maximum=200) | |
| # Right: Urine Tests + Manual | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Urine Tests") | |
| sg_input = gr.Dropdown(label="Specific Gravity - sg", | |
| choices=["1.005","1.010","1.015","1.020","1.025"], | |
| value="1.020") | |
| al_input = gr.Dropdown(label="Albumin in Urine - al (0=None, 5=Heavy)", | |
| choices=["0","1","2","3","4","5"], value="0") | |
| su_input = gr.Dropdown(label="Sugar in Urine - su (0=None, 5=Heavy)", | |
| choices=["0","1","2","3","4","5"], value="0") | |
| rbc_input = gr.Radio(label="RBC in Urine - rbc", | |
| choices=["normal","abnormal"], value="normal") | |
| pc_input = gr.Radio(label="Pus Cells - pc", | |
| choices=["normal","abnormal"], value="normal") | |
| pcc_input = gr.Radio(label="Pus Cell Clumps - pcc", | |
| choices=["notpresent","present"], value="notpresent") | |
| ba_input = gr.Radio(label="Bacteria in Urine - ba", | |
| choices=["notpresent","present"], value="notpresent") | |
| gr.Markdown("### Patient Info (Manual)") | |
| age_input = gr.Number(label="Age (years)", value=45, minimum=2, maximum=90) | |
| htn_input = gr.Radio(label="Hypertension - htn", choices=["no","yes"], value="no") | |
| dm_input = gr.Radio(label="Diabetes Mellitus - dm", choices=["no","yes"], value="no") | |
| cad_input = gr.Radio(label="Coronary Artery Disease - cad", choices=["no","yes"], value="no") | |
| appet_input = gr.Radio(label="Appetite - appet", choices=["good","poor"], value="good") | |
| pe_input = gr.Radio(label="Pedal Edema (Swelling) - pe", choices=["no","yes"], value="no") | |
| ane_input = gr.Radio(label="Anemia - ane", choices=["no","yes"], value="no") | |
| predict_btn = gr.Button(" Run CKD Risk Assessment", variant="primary", size="lg") | |
| with gr.Row(): | |
| with gr.Column(): | |
| gradcam_output = gr.Image(label="Retinal Analysis + Grad-CAM++") | |
| with gr.Column(): | |
| risk_output = gr.Markdown(label="Risk Result") | |
| clinical_output = gr.Textbox(label="Full Report", lines=20, interactive=False) | |
| status_output = gr.Textbox(label="Status", interactive=False) | |
| # ── Tab 2: CKD Health Coach ──────────────────────────────────────── | |
| with gr.TabItem(" CKD Health Coach"): | |
| gr.Markdown("### Ask about diet, lifestyle, exercise, and kidney health.") | |
| chatbot = gr.Chatbot(height=450) | |
| chat_input = gr.Textbox(placeholder="Ask a question about CKD...", label="Your Message") | |
| send_btn = gr.Button("Send", variant="primary") | |
| gr.Examples( | |
| examples=["What foods should I avoid with CKD?", | |
| "How much water should I drink daily?", | |
| "Can I exercise with kidney disease?", | |
| "What are warning signs I should watch for?"], | |
| inputs=chat_input | |
| ) | |
| # ── Tab 3: About ─────────────────────────────────────────────────── | |
| with gr.TabItem("About"): | |
| gr.Markdown(""" | |
| ## About CKD-MultiVision | |
| **Input Fields (24 total — matching UCI CKD Dataset):** | |
| | # | Field | Source | Description | | |
| |---|-------|--------|-------------| | |
| | 1 | age | Manual | Patient age in years | | |
| | 2 | bp | Lab | Blood pressure diastolic (mmHg) | | |
| | 3 | bgr | Lab | Blood glucose random (mg/dL) | | |
| | 4 | bu | Lab | Blood urea (mg/dL) | | |
| | 5 | sc | Lab | Serum creatinine (mg/dL) | | |
| | 6 | sod | Lab | Sodium (mEq/L) | | |
| | 7 | pot | Lab | Potassium (mEq/L) | | |
| | 8 | hemo | Lab | Hemoglobin (g/dL) | | |
| | 9 | pcv | Lab | Packed cell volume (%) | | |
| | 10 | wbcc | Lab | WBC count (cells/cumm) | | |
| | 11 | rbcc | Lab | RBC count (millions/cumm) | | |
| | 12 | sg | Lab | Specific gravity | | |
| | 13 | al | Lab | Albumin in urine (0-5) | | |
| | 14 | su | Lab | Sugar in urine (0-5) | | |
| | 15 | rbc | Lab | RBC in urine | | |
| | 16 | pc | Lab | Pus cells | | |
| | 17 | pcc | Lab | Pus cell clumps | | |
| | 18 | ba | Lab | Bacteria | | |
| | 19 | htn | Manual | Hypertension | | |
| | 20 | dm | Manual | Diabetes mellitus | | |
| | 21 | cad | Manual | Coronary artery disease | | |
| | 22 | appet | Manual | Appetite | | |
| | 23 | pe | Manual | Pedal edema | | |
| | 24 | ane | Manual | Anemia | | |
| **Pipeline:** | |
| - Retinal Image → EfficientNet-B3 → Feature Vector (1536) → PCA (64 dim) | |
| - Clinical Values (24) → StandardScaler → Gradient Boosting + Fusion Model | |
| - Cross-Modal Attention Fusion → Final Risk Level | |
| *Research prototype — not for clinical use.* | |
| """) | |
| # ── Event Handlers ───────────────────────────────────────────────────── | |
| extract_btn.click( | |
| fn=process_lab_file, | |
| inputs=[lab_file_input], | |
| outputs=[ | |
| extract_status, | |
| sc_input, bu_input, hemo_input, bgr_input, sod_input, | |
| pot_input, pcv_input, wbcc_input, rbcc_input, bp_input, | |
| sg_input, al_input, su_input, rbc_input, pc_input, pcc_input, ba_input, | |
| ] | |
| ) | |
| predict_btn.click( | |
| fn=predict_ckd_risk, | |
| inputs=[ | |
| retinal_input, | |
| # Blood tests | |
| sc_input, bu_input, hemo_input, bgr_input, | |
| sod_input, pot_input, pcv_input, wbcc_input, rbcc_input, | |
| bp_input, | |
| # Urine tests | |
| sg_input, al_input, su_input, rbc_input, pc_input, pcc_input, ba_input, | |
| # Manual | |
| age_input, htn_input, dm_input, cad_input, appet_input, pe_input, ane_input, | |
| ], | |
| outputs=[gradcam_output, risk_output, clinical_output, status_output] | |
| ) | |
| def respond(message, history): | |
| if not message.strip(): | |
| return history, "" | |
| bot_response = chat_with_groq(message, history) | |
| history.append({"role": "user", "content": message}) | |
| history.append({"role": "assistant", "content": bot_response}) | |
| return history, "" | |
| send_btn.click(fn=respond, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input]) | |
| chat_input.submit(fn=respond, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input]) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Soft()) |