Spaces:
Sleeping
Sleeping
File size: 3,534 Bytes
0f049b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | 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
|