Spaces:
Sleeping
Sleeping
| import torch | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| from data import build_transforms | |
| from model import build_model | |
| import json | |
| from pathlib import Path | |
| def generate_gradcam(img, model, device, class_names, threshold=0.5): | |
| """ | |
| Generate Grad-CAM for a given image and model. | |
| Args: | |
| img (PIL.Image): Input image. | |
| model (torch.nn.Module): Pretrained model. | |
| device (torch.device): Torch device (cpu/cuda). | |
| class_names (list): List of class labels. | |
| threshold (float): Threshold for masking heatmap. | |
| Returns: | |
| overlay_rgb (PIL.Image): Overlay of heatmap on original image. | |
| heatmap (np.ndarray): Heatmap array. | |
| pred_label (str): Predicted class label. | |
| confidence (float): Probability of predicted class. | |
| """ | |
| # Load transforms | |
| _, eval_transforms = build_transforms(224) | |
| # Handle input type | |
| if isinstance(img, str): | |
| img = Image.open(img).convert("RGB") | |
| elif not isinstance(img, Image.Image): | |
| raise ValueError("Input must be a file path or PIL.Image") | |
| img_tensor = eval_transforms(img).unsqueeze(0).to(device) | |
| # Hook storage | |
| feats, grads = [], [] | |
| def fwd_hook(m, i, o): feats.append(o) | |
| def bwd_hook(m, gi, go): grads.append(go[0]) | |
| # Register hooks | |
| layer = model.layer4[1].conv2 | |
| layer.register_forward_hook(fwd_hook) | |
| layer.register_full_backward_hook(bwd_hook) | |
| # Forward pass | |
| out = model(img_tensor) | |
| pred_idx = out.argmax(1).item() | |
| score = out[0, pred_idx] | |
| # Confidence (softmax) | |
| probs = torch.softmax(out, dim=1) | |
| confidence = probs[0, pred_idx].item() | |
| # Backward pass | |
| model.zero_grad() | |
| score.backward() | |
| # Grad-CAM calculation | |
| grad = grads[0][0].detach().cpu().numpy() | |
| feat = feats[0][0].detach().cpu().numpy() | |
| weights = grad.mean(axis=(1, 2)) | |
| cam = np.maximum(np.sum(weights[:, None, None] * feat, axis=0), 0) | |
| cam = cv2.resize(cam, img.size) | |
| cam = cam / cam.max() | |
| # Mask | |
| mask = (cam > threshold).astype(np.uint8) * 255 | |
| # Heatmap | |
| heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_TURBO) | |
| heatmap_masked = cv2.bitwise_and(heatmap, heatmap, mask=mask) | |
| heatmap_rgb = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) | |
| # Overlay | |
| orig_bgr = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) | |
| overlay = cv2.addWeighted(orig_bgr, 0.5, heatmap_masked, 0.5, 0) | |
| overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB) | |
| return Image.fromarray(overlay_rgb), heatmap_rgb, class_names[pred_idx], confidence | |
| if __name__ == "__main__": | |
| # Load labels | |
| with open("models/labels.json", "r") as f: | |
| idx2name = {int(k): v for k, v in json.load(f).items()} | |
| class_names = [idx2name[i] for i in sorted(idx2name.keys())] | |
| # Load model | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = build_model(num_classes=len(class_names), freeze_backbone=False, device=device) | |
| model.load_state_dict(torch.load("models/resnet18_best.pt", map_location=device)) | |
| model.eval() | |
| print("Model loaded.") | |
| # Test with one image | |
| img = Image.open("data/own_images/test/paper/paper_normal_dark_23sec_013.jpg").convert("RGB") | |
| overlay, heatmap, pred_label, conf = generate_gradcam(img, model, device, class_names) | |
| print(f"Final Prediction: {pred_label} ({conf:.2%})") | |
| overlay.show() # Display overlay | |
| heatmap_img = Image.fromarray(heatmap) | |
| heatmap_img.show() # Display heatmap | |