Update app.py

app.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 import timm
 from torchvision import transforms
@@ -6,14 +7,15 @@ from PIL import Image
 import numpy as np
 import cv2
 import gradio as gr
+import os
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 MODEL_PATH = "best_timm_pneumonia.pt"
+SEG_PATHS = ["lung_unet_lite.pt", "lung_unet_lite.ts"]  # accepted filenames
 IMG_SIZE = 384
 
-
 # ---------------------------------------------------------
-# ✅ Preprocessing
+# Preprocessing
 # ---------------------------------------------------------
 val_tf = transforms.Compose([
     transforms.Grayscale(num_output_channels=3),
@@ -23,9 +25,8 @@ val_tf = transforms.Compose([
     transforms.Normalize([0.485]*3, [0.229]*3),
 ])
 
-
 # ---------------------------------------------------------
-# ✅ Load EfficientNet Model
+# Classifier (EfficientNet-B0)
 # ---------------------------------------------------------
 def build_model():
     return timm.create_model("tf_efficientnet_b0_ns", pretrained=False, num_classes=1)
@@ -34,34 +35,106 @@ model = build_model().to(DEVICE)
 model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
 model.eval()
 
+# ---------------------------------------------------------
+# Lightweight Lung U-Net (architecture only; ~5MB weights expected)
+# ---------------------------------------------------------
+class DSConv(nn.Module):
+    def __init__(self, c_in, c_out):
+        super().__init__()
+        self.dw = nn.Conv2d(c_in, c_in, 3, padding=1, groups=c_in, bias=False)
+        self.pw = nn.Conv2d(c_in, c_out, 1, bias=False)
+        self.bn = nn.BatchNorm2d(c_out)
+        self.act = nn.ReLU(inplace=True)
+    def forward(self, x):
+        x = self.dw(x); x = self.pw(x); x = self.bn(x); return self.act(x)
+
+class Block(nn.Module):
+    def __init__(self, c_in, c_out):
+        super().__init__()
+        self.c1 = DSConv(c_in, c_out)
+        self.c2 = DSConv(c_out, c_out)
+    def forward(self, x): return self.c2(self.c1(x))
+
+class Up(nn.Module):
+    def __init__(self, c_in, c_out):
+        super().__init__()
+        self.up = nn.ConvTranspose2d(c_in, c_in//2, 2, stride=2)
+        self.conv = Block(c_in, c_out)
+    def forward(self, x, skip):
+        x = self.up(x)
+        # pad if sizes mismatch
+        dh, dw = skip.shape[2]-x.shape[2], skip.shape[3]-x.shape[3]
+        x = F.pad(x, [dw//2, dw-dw//2, dh//2, dh-dh//2])
+        x = torch.cat([skip, x], dim=1)
+        return self.conv(x)
+
+class LungUNetLite(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.enc1 = Block(1, 32)
+        self.enc2 = Block(32, 64)
+        self.enc3 = Block(64, 128)
+        self.enc4 = Block(128, 256)
+        self.pool = nn.MaxPool2d(2)
+        self.bott = Block(256, 256)
+        self.up3 = Up(256, 128)
+        self.up2 = Up(128, 64)
+        self.up1 = Up(64, 32)
+        self.outc = nn.Conv2d(32, 1, 1)
+
+    def forward(self, x):
+        e1 = self.enc1(x)          # 1/1
+        e2 = self.enc2(self.pool(e1))  # 1/2
+        e3 = self.enc3(self.pool(e2))  # 1/4
+        e4 = self.enc4(self.pool(e3))  # 1/8
+        b  = self.bott(self.pool(e4))  # 1/16
+        x  = self.up3(b, e4)
+        x  = self.up2(x, e3)
+        x  = self.up1(x, e2)
+        x  = self.outc(x)
+        return x
+
+# Try to load a tiny lung segmentation model if the user uploaded it
+lung_net = None
+seg_loaded_msg = ""
+for p in SEG_PATHS:
+    if os.path.exists(p):
+        try:
+            if p.endswith(".ts"):  # TorchScript
+                lung_net = torch.jit.load(p, map_location=DEVICE)
+            else:                  # state_dict
+                lung_net = LungUNetLite().to(DEVICE)
+                lung_net.load_state_dict(torch.load(p, map_location=DEVICE))
+            lung_net.eval()
+            seg_loaded_msg = f"✅ Lung model loaded: {p}"
+            break
+        except Exception as e:
+            seg_loaded_msg = f"⚠️ Failed to load {p}: {e}"
+
+if lung_net is None and seg_loaded_msg == "":
+    seg_loaded_msg = "ℹ️ Using classical (non-ML) lung mask; upload 'lung_unet_lite.pt' to upgrade."
 
 # ---------------------------------------------------------
-# ✅ Auto-Crop Dark Borders
+# Auto-crop borders
 # ---------------------------------------------------------
 def autocrop_chest(pil_img, th=10):
     gray = np.array(pil_img.convert("L"))
     mask = gray > th
-    if mask.sum() < 10:
-        return pil_img
+    if mask.sum() < 10: return pil_img
     ys, xs = np.where(mask)
     return Image.fromarray(gray[ys.min():ys.max(), xs.min():xs.max()]).convert("RGB")
 
-
 # ---------------------------------------------------------
-# ✅ Classical Lung Mask (no models required)
+# Classical lung mask (fallback, no ML)
 # ---------------------------------------------------------
-def get_lung_mask(pil_img):
+def classical_lung_mask(pil_img):
     img = np.array(pil_img.convert("L"))
-
     clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
     g = clahe.apply(img)
-
     g = cv2.GaussianBlur(g, (5, 5), 0)
     _, th = cv2.threshold(g, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
-
     th = cv2.morphologyEx(th, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8), iterations=1)
     th = cv2.morphologyEx(th, cv2.MORPH_CLOSE, np.ones((15, 15), np.uint8), iterations=1)
-
     num, labels, stats, _ = cv2.connectedComponentsWithStats(th, connectivity=8)
     if num > 1:
         areas = [(i, stats[i, cv2.CC_STAT_AREA]) for i in range(1, num)]
@@ -70,79 +143,76 @@ def get_lung_mask(pil_img):
         mask = np.where(np.isin(labels, list(keep)), 1.0, 0.0).astype(np.float32)
     else:
         mask = (th > 0).astype(np.float32)
-
     mask = cv2.dilate(mask, np.ones((11, 11), np.uint8), iterations=1)
     mask = cv2.GaussianBlur(mask, (21, 21), 0)
-
     h, w = mask.shape
     b = int(0.06 * w)
     mask[:b, :] = mask[-b:, :] = mask[:, :b] = mask[:, -b:] = 0
+    mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
+    return mask
 
+# ---------------------------------------------------------
+# ML lung mask (if weights provided)
+# ---------------------------------------------------------
+def ml_lung_mask(pil_img):
+    # expects grayscale 1xHxW scaled to 256 side for speed, then resize back
+    img = np.array(pil_img.convert("L"))
+    h, w = img.shape
+    side = 256
+    imr = cv2.resize(img, (side, side))
+    tens = torch.from_numpy(imr[None, None].astype(np.float32) / 255.0).to(DEVICE)
+    with torch.no_grad():
+        logits = lung_net(tens)
+        prob = torch.sigmoid(logits)[0, 0].cpu().numpy()
+    mask = cv2.resize(prob, (IMG_SIZE, IMG_SIZE))
+    # soft refine
+    mask = cv2.GaussianBlur(mask, (15, 15), 0)
     mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
     return mask
 
+def get_lung_mask(pil_img):
+    if lung_net is not None:
+        return ml_lung_mask(pil_img)
+    return classical_lung_mask(pil_img)
 
 # ---------------------------------------------------------
-# ✅ Soft Grad-CAM
+# Soft Grad-CAM (correct layer)
 # ---------------------------------------------------------
 def get_gradcam(model, img_tensor):
     target_layer = model.blocks[-1]
-
     activ, grads = [], []
-
-    def fwd(m, i, o):
-        activ.append(o)
-
-    def bwd(m, gi, go):
-        grads.append(go[0])
-
+    def fwd(m, i, o): activ.append(o)
+    def bwd(m, gi, go): grads.append(go[0])
     h1 = target_layer.register_forward_hook(fwd)
     h2 = target_layer.register_backward_hook(bwd)
-
     logits = model(img_tensor)
     score = torch.sigmoid(logits)[0]
-    model.zero_grad()
-    score.backward()
-
-    h1.remove()
-    h2.remove()
-
+    model.zero_grad(); score.backward()
+    h1.remove(); h2.remove()
     A = activ[0][0].detach().cpu().numpy()
     G = grads[0].detach().cpu().numpy()
-
     weights = G.mean(axis=(1, 2))
     cam = np.zeros(A.shape[1:], dtype=np.float32)
-    for c, w in enumerate(weights):
-        cam += w * A[c]
-
+    for c, w in enumerate(weights): cam += w * A[c]
     cam = np.maximum(cam, 0)
-    cam -= cam.min()
-    cam /= (cam.max() + 1e-8)
-
+    cam -= cam.min(); cam /= (cam.max() + 1e-8)
     cam = cv2.GaussianBlur(cam, (7, 7), 0)
     return cam
 
-
 # ---------------------------------------------------------
-# ✅ Visualization Modes
+# Mode enhancement
 # ---------------------------------------------------------
 def enhance_cam(cam, mode):
-    if mode == "Weak":
-        return cam ** 1.4
-    elif mode == "Strong":
-        return cam ** 0.6
+    if mode == "Weak":   return cam ** 1.4
+    if mode == "Strong": return cam ** 0.6
     return cam ** 0.9  # Medium
 
-
 # ---------------------------------------------------------
-# ✅ Prediction + Visualization
+# Prediction + Visualization
 # ---------------------------------------------------------
 def predict(img, mode, intensity, threshold):
-
-    # Preprocess
     cropped = autocrop_chest(img)
     resized = cropped.resize((IMG_SIZE, IMG_SIZE))
-
     x = val_tf(resized).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
@@ -154,42 +224,35 @@ def predict(img, mode, intensity, threshold):
     # Grad-CAM
     cam = get_gradcam(model, x)
     cam = cv2.resize(cam, (IMG_SIZE, IMG_SIZE))
-
-    # Soft threshold
     cam = np.where(cam >= threshold, cam, cam * 0.3)
     cam = enhance_cam(cam, mode)
-    cam /= (cam.max() + 1e-8)
+    cam = cam / (cam.max() + 1e-8)
 
-    # Heatmap
+    # Heatmap + base (both float32 [0,1])
     heat = cv2.applyColorMap((cam * 255).astype(np.uint8), cv2.COLORMAP_JET)[..., ::-1] / 255.0
-
-    # Convert both to float32 to avoid OpenCV errors
     base = np.array(resized).astype(np.float32) / 255.0
     base_f = base.astype(np.float32)
     heat_f = heat.astype(np.float32)
 
-    # Blended overlay
+    # Overlay
     overlay = cv2.addWeighted(base_f, 1 - intensity, heat_f, intensity, 0)
     overlay = np.clip(overlay * 255, 0, 255).astype(np.uint8)
 
     # Lung-masked CAM
     lung_mask = get_lung_mask(resized)
-    masked = heat_f * lung_mask[..., None]
+    masked = (heat_f * lung_mask[..., None])
     masked = np.clip(masked * 255, 0, 255).astype(np.uint8)
 
-    # Side-by-side results
     combined = np.hstack([
         (base * 255).astype(np.uint8),
         overlay,
         masked
     ])
-
-    text = f"Prediction: {label}\nP(PNEUMONIA)={prob_p:.3f}\nP(NORMAL)={prob_n:.3f}"
+    text = f"{seg_loaded_msg}\nPrediction: {label}\nP(PNEUMONIA)={prob_p:.3f} | P(NORMAL)={prob_n:.3f}"
     return text, combined
 
-
 # ---------------------------------------------------------
-# ✅ Gradio UI
+# UI
 # ---------------------------------------------------------
 demo = gr.Interface(
     fn=predict,
@@ -197,16 +260,14 @@ demo = gr.Interface(
         gr.Image(type="pil", label="Upload Chest X-Ray"),
         gr.Radio(["Weak", "Medium", "Strong"], value="Medium", label="Heatmap Mode"),
         gr.Slider(0.1, 1.0, value=0.70, step=0.05, label="Heatmap Intensity"),
-        gr.Slider(0.0, 1.0, value=0.40, step=0.05, label="Heatmap Threshold")
+        gr.Slider(0.0, 1.0, value=0.40, step=0.05, label="Heatmap Threshold"),
     ],
     outputs=[
-        gr.Text(label="Prediction"),
-        gr.Image(label="Original | Heatmap | Lung-Masked CAM")
+        gr.Text(label="Status + Prediction"),
+        gr.Image(label="Original | Heatmap | Lung-Masked CAM"),
     ],
-    title="Pneumonia Detection (Soft Grad-CAM + Lung Mask)",
-    description=(
-        "Upload an X-ray to view pneumonia predictions with heatmaps and lung-masked Grad-CAM."
-    )
+    title="Pneumonia Detector (Soft Grad-CAM + Lung Mask, Lite)",
+    description="Upload an X-ray. For lung-only CAM, upload 'lung_unet_lite.pt' or 'lung_unet_lite.ts' to this Space."
 )
 
 demo.launch()