File size: 5,590 Bytes
43abac3 | 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | import matplotlib.pyplot as plt
import numpy as np
import cv2
def plot_greenery_overlay(image_np, mask_classifications, output_path=None):
"""
Plots the original image and an overlay where greenery is highlighted.
"""
plt.figure(figsize=(10, 10))
plt.imshow(image_np)
ax = plt.gca()
ax.set_autoscale_on(False)
# Create an aggregate mask for greenery
h, w = image_np.shape[:2]
greenery_overlay = np.zeros((h, w, 4))
green_color = np.array([0, 1, 0, 0.5]) # Semi-transparent green
for item in mask_classifications:
if item['is_green']:
seg = item['segmentation']
greenery_overlay[seg] = green_color
ax.imshow(greenery_overlay)
plt.axis('off')
if output_path:
plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
plt.close()
else:
plt.show()
def create_greenery_overlay(image_np, mask_classifications):
"""
Creates a semi-transparent greenery overlay composite image.
Returns:
tuple: (composite_rgba, green_masks) where composite_rgba is the
overlay image and green_masks is the list of greenery items.
"""
h, w = image_np.shape[:2]
greenery_overlay = np.zeros((h, w, 4))
green_color = np.array([0, 1, 0, 0.5]) # Semi-transparent green
green_masks = []
for item in mask_classifications:
if item['is_green']:
green_masks.append(item)
seg = item['segmentation']
greenery_overlay[seg] = green_color
base_img_rgba = cv2.cvtColor(image_np, cv2.COLOR_RGB2RGBA)
alpha_mask = greenery_overlay[:, :, 3] > 0
composite = base_img_rgba.copy()
composite[alpha_mask] = (
base_img_rgba[alpha_mask] * 0.5 +
greenery_overlay[alpha_mask] * 255 * 0.5
).astype(np.uint8)
return composite, green_masks
def apply_grad_cam(model, input_tensor, target_class=None):
"""
Generates a Grad-CAM heatmap for the given input tensor using the model's
final convolutional layer (layer4 for ResNet-50).
Args:
model: A GreeneryClassifier instance (wraps ResNet-50 via model.model).
input_tensor: A preprocessed image tensor of shape (1, 3, 64, 64).
target_class: The class index to generate the heatmap for.
If None, uses the model's predicted class.
Returns:
heatmap: A numpy array (H, W) with values in [0, 1] representing
the class activation map.
predicted_class: The class index the model predicted.
"""
import torch
import torch.nn.functional as F
model.eval()
device = next(model.parameters()).device
input_tensor = input_tensor.to(device)
# Storage for hooked values
activations = []
gradients = []
# Hook into the last convolutional block of the inner ResNet-50
target_layer = model.model.layer4
def forward_hook(module, input, output):
activations.append(output.detach())
def backward_hook(module, grad_input, grad_output):
gradients.append(grad_output[0].detach())
fwd_handle = target_layer.register_forward_hook(forward_hook)
bwd_handle = target_layer.register_full_backward_hook(backward_hook)
try:
# Forward pass
output = model(input_tensor)
predicted_class = output.argmax(dim=1).item()
if target_class is None:
target_class = predicted_class
# Backward pass for the target class
model.zero_grad()
class_score = output[0, target_class]
class_score.backward()
# Compute Grad-CAM
act = activations[0] # (1, C, H', W')
grad = gradients[0] # (1, C, H', W')
# Global average pool the gradients to get per-channel weights
weights = grad.mean(dim=[2, 3], keepdim=True) # (1, C, 1, 1)
# Weighted combination of activation maps
cam = (weights * act).sum(dim=1, keepdim=True) # (1, 1, H', W')
cam = F.relu(cam) # Only positive contributions
# Upsample to input size
cam = F.interpolate(cam, size=input_tensor.shape[2:], mode='bilinear', align_corners=False)
cam = cam.squeeze().cpu().numpy()
# Normalize to [0, 1]
if cam.max() > 0:
cam = (cam - cam.min()) / (cam.max() - cam.min())
return cam, predicted_class
finally:
# Guarantee hook cleanup even if an exception fires mid-computation
fwd_handle.remove()
bwd_handle.remove()
def save_grad_cam_overlay(image_np, heatmap, output_path, alpha=0.4):
"""
Overlays a Grad-CAM heatmap on the original image and saves to disk.
Args:
image_np: Original RGB image as numpy array (H, W, 3), values 0-255.
heatmap: Grad-CAM heatmap as numpy array (H, W), values in [0, 1].
output_path: File path to save the result.
alpha: Transparency of the heatmap overlay.
"""
# Resize heatmap to match the image dimensions
heatmap_resized = cv2.resize(heatmap, (image_np.shape[1], image_np.shape[0]))
heatmap_colored = cv2.applyColorMap(np.uint8(255 * heatmap_resized), cv2.COLORMAP_JET)
heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
overlay = np.uint8(image_np * (1 - alpha) + heatmap_colored * alpha)
plt.figure(figsize=(10, 10))
plt.imshow(overlay)
plt.axis('off')
plt.title('Grad-CAM Heatmap Overlay')
plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
plt.close()
print(f"Grad-CAM overlay saved to {output_path}")
|