import json import cv2 import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torchvision import models from pathlib import Path # ─── Architecture (must match training notebook exactly) ─────────────────────── class ChannelAttention(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc = nn.Sequential( nn.Conv2d(channels, channels // reduction, 1, bias=False), nn.ReLU(inplace=True), nn.Conv2d(channels // reduction, channels, 1, bias=False), ) self.sigmoid = nn.Sigmoid() def forward(self, x): return self.sigmoid(self.fc(self.avg_pool(x)) + self.fc(self.max_pool(x))) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super().__init__() self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) return self.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1))) class CBAM(nn.Module): def __init__(self, channels, reduction=16, kernel_size=7): super().__init__() self.channel_attention = ChannelAttention(channels, reduction) self.spatial_attention = SpatialAttention(kernel_size) def forward(self, x): x = x * self.channel_attention(x) x = x * self.spatial_attention(x) return x class DenseNet121ForBinaryClassification(nn.Module): def __init__(self, num_classes=1, pretrained=False, dropout=0.5): super().__init__() self.densenet = models.densenet121(pretrained=pretrained) orig_conv = self.densenet.features.conv0 self.densenet.features.conv0 = nn.Conv2d( 1, orig_conv.out_channels, kernel_size=orig_conv.kernel_size, stride=orig_conv.stride, padding=orig_conv.padding, bias=False ) self.cbam = CBAM(channels=1024, reduction=16, kernel_size=7) num_features = self.densenet.classifier.in_features self.densenet.classifier = nn.Sequential( nn.Dropout(dropout), nn.Linear(num_features, num_classes) ) def forward(self, x): features = self.densenet.features(x) features = self.cbam(features) out = F.adaptive_avg_pool2d(features, (1, 1)) out = out.view(out.size(0), -1) return self.densenet.classifier(out) # ─── GradCAM++ ───────────────────────────────────────────────────────────────── class GradCAMPlusPlus: def __init__(self, model, target_layer_name="cbam"): self.model = model self.gradients = None self.activations = None self.hooks = [] self._register_hooks(target_layer_name) def _register_hooks(self, target_layer_name): def forward_hook(module, input, output): self.activations = output.detach() def backward_hook(module, grad_input, grad_output): self.gradients = grad_output[0].detach() for name, module in self.model.named_modules(): if name == target_layer_name: self.hooks.append(module.register_forward_hook(forward_hook)) self.hooks.append(module.register_backward_hook(backward_hook)) return raise ValueError(f"Layer '{target_layer_name}' not found in model") def generate_cam(self, input_tensor, class_idx=None): self.model.eval() for param in self.model.parameters(): param.requires_grad = True input_tensor = input_tensor.clone().detach().requires_grad_(True) output = self.model(input_tensor) if class_idx is None: class_idx = int(torch.sigmoid(output).round().item()) target_score = output[0, 0] if class_idx == 1 else -output[0, 0] self.model.zero_grad() target_score.backward(retain_graph=True) grads = self.gradients # [1, C, H, W] acts = self.activations # [1, C, H, W] B, C, H, W = grads.shape # ── GradCAM++ alpha computation (all ops stay in [B, C, H, W]) ── grads_sq = grads.pow(2) # [1, C, H, W] grads_cub = grads.pow(3) # [1, C, H, W] # sum over spatial dims H,W → [1, C, 1, 1] then broadcast back spatial_sum = (acts * grads_cub).sum(dim=[2, 3], keepdim=True) # [1, C, 1, 1] alpha_denom = 2.0 * grads_sq + spatial_sum # [1, C, H, W] alpha_denom = torch.where( alpha_denom != 0.0, alpha_denom, torch.ones_like(alpha_denom) ) alpha = grads_sq / (alpha_denom + 1e-7) # [1, C, H, W] # weights: alpha * ReLU(grads), summed over H,W → [1, C, 1, 1] weights = (alpha * torch.relu(grads)).sum(dim=[2, 3], keepdim=True) # [1, C, 1, 1] # CAM: weighted sum of activations → [1, 1, H, W] cam = (weights * acts).sum(dim=1, keepdim=True) # [1, 1, H, W] cam = torch.relu(cam) cam = cam.squeeze().cpu().detach().numpy() # [H, W] cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) return cam def cleanup(self): for hook in self.hooks: hook.remove() # ─── Preprocessing ───────────────────────────────────────────────────────────── def apply_jet_colormap(gray_img): """Apply jet colormap manually without matplotlib (0=blue, 1=red).""" gray_img = np.clip(gray_img, 0, 1) r = np.clip(1.5 - np.abs(gray_img * 2 - 3), 0, 1) g = np.clip(1.5 - np.abs(gray_img * 2 - 2), 0, 1) b = np.clip(1.5 - np.abs(gray_img * 2 - 1), 0, 1) return (np.stack([r, g, b], axis=-1) * 255).astype(np.uint8) def apply_clahe_cv2(image, clip_limit=2.0, tile_grid_size=(8, 8)): """Apply CLAHE using OpenCV (replaces albumentations).""" clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size) return clahe.apply(image) def preprocess_with_cv2(image, image_size=512): """Preprocessing pipeline using pure OpenCV (replaces albumentations).""" # CLAHE image = apply_clahe_cv2(image, clip_limit=2.0, tile_grid_size=(8, 8)) # Center crop 350x350 h, w = image.shape[:2] start_y = (h - 350) // 2 start_x = (w - 350) // 2 image = image[start_y:start_y+350, start_x:start_x+350] # Resize to target size image = cv2.resize(image, (image_size, image_size), interpolation=cv2.INTER_LINEAR) return image def preprocess_image(image_path: str, meta: dict) -> torch.Tensor: """Returns a [1, 1, H, W] tensor ready for inference.""" image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) if image is None: raise ValueError(f"Could not read image: {image_path}") # Use OpenCV preprocessing instead of albumentations image = preprocess_with_cv2(image, image_size=meta["image_size"]) image = image.astype(np.float32) / 255.0 image = (image - meta["global_mean"]) / meta["global_std"] tensor = torch.from_numpy(image).unsqueeze(0).unsqueeze(0) # [1, 1, H, W] return tensor # ─── Model Loader (singleton) ────────────────────────────────────────────────── WEIGHTS_PATH = "trainedmodels/Model.pth" META_PATH = "trainedmodels/Model.json" DEVICE = "cpu" _model_cache = {} def load_model(weights_path: str, device: str, meta_path: str = None) -> tuple: if weights_path in _model_cache: return _model_cache[weights_path] # ── Always resolve paths using Path() to normalize slashes ── weights_path = Path(weights_path).as_posix() if meta_path is None: meta_path = Path(weights_path).with_suffix(".json").as_posix() else: meta_path = Path(meta_path).as_posix() # normalize whatever is passed in if not Path(meta_path).exists(): raise FileNotFoundError(f"Metadata file not found: {meta_path}") with open(meta_path) as f: meta = json.load(f) model = DenseNet121ForBinaryClassification( num_classes=1, pretrained=False, dropout=meta["dropout"] ) state = torch.load(weights_path, map_location=device, weights_only=True) model.load_state_dict(state) model.to(device) model.eval() _model_cache[weights_path] = (model, meta) return model, meta # ─── Single Inference ────────────────────────────────────────────────────────── def run_inference(image_path: str, weights_path: str, device: str = "cpu", generate_gradcam: bool = False) -> dict: model, meta = load_model(weights_path, device) tensor = preprocess_image(image_path, meta).to(device) with torch.no_grad(): logit = model(tensor) prob = torch.sigmoid(logit).item() threshold = meta["best_threshold"] pred_class = int(prob >= threshold) if pred_class == 1: label = "Cancer" display_prob = prob else: label = "Normal" display_prob = 1 - prob result = { "probability": round(display_prob, 4), "predicted_class": pred_class, "label": label, "threshold_used": threshold, "gradcam_overlay": None } if generate_gradcam: gradcam = GradCAMPlusPlus(model, target_layer_name="cbam") try: # GradCAM needs gradients — don't use no_grad here tensor_gc = preprocess_image(image_path, meta).to(device) cam_map = gradcam.generate_cam(tensor_gc, class_idx=pred_class) # Build overlay orig = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) orig = cv2.resize(orig, (meta["image_size"], meta["image_size"])) orig_rgb = cv2.cvtColor(orig, cv2.COLOR_GRAY2RGB) cam_resized = cv2.resize(cam_map, (orig_rgb.shape[1], orig_rgb.shape[0])) # Use matplotlib jet colormap for better quality import matplotlib.cm as mpl_cm heatmap = (mpl_cm.jet(cam_resized)[:, :, :3] * 255).astype(np.uint8) overlay = cv2.addWeighted(orig_rgb, 0.6, heatmap, 0.4, 0) result["gradcam_overlay"] = overlay # numpy array, encode downstream finally: gradcam.cleanup() return result