File size: 3,924 Bytes
2d140d5 95cdfc0 2d140d5 | 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 | """Shared model/Grad-CAM code, extracted from
mit-group8-explainable-defect-detection.ipynb (sections 3 and 6) so the
Space app and the training script can both import it without duplication."""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
import matplotlib
CLASS_NAMES = [
"crazing",
"inclusion",
"patches",
"pitted_surface",
"rolled-in_scale",
"scratches",
]
IMAGE_SIZE = 128
MEAN = (0.5, 0.5, 0.5)
STD = (0.25, 0.25, 0.25)
eval_transform = T.Compose([
T.Resize((IMAGE_SIZE, IMAGE_SIZE)),
T.ToTensor(),
T.Normalize(MEAN, STD),
])
_mean_t = torch.tensor(MEAN).view(3, 1, 1)
_std_t = torch.tensor(STD).view(3, 1, 1)
def to_image(tensor):
return (tensor.detach().cpu() * _std_t + _mean_t).clamp(0, 1).permute(1, 2, 0).numpy()
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride):
super().__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride,
padding=1, bias=False)
self.norm = nn.BatchNorm2d(out_channels)
self.activation = nn.ReLU()
def forward(self, x):
return self.activation(self.norm(self.conv(x)))
class DefectCNN(nn.Module):
def __init__(self, number_of_classes=6):
super().__init__()
self.features = nn.Sequential(
ConvBlock(3, 24, stride=2),
ConvBlock(24, 48, stride=2),
ConvBlock(48, 96, stride=2),
ConvBlock(96, 128, stride=1),
)
self.pool = nn.AdaptiveAvgPool2d(1)
self.classifier = nn.Sequential(
nn.Dropout(0.20),
nn.Linear(128, number_of_classes),
)
def forward(self, x):
features = self.features(x)
return self.classifier(self.pool(features).flatten(1))
class GradCAM:
def __init__(self, model, target_layer):
self.model = model
self.target_layer = target_layer
self.activations = None
self.gradients = None
self.forward_handle = target_layer.register_forward_hook(self._capture_activations)
def _capture_activations(self, module, inputs, output):
self.activations = output
if output.requires_grad:
output.register_hook(self._capture_gradients)
def _capture_gradients(self, gradient):
self.gradients = gradient
def __call__(self, images, target_classes):
self.model.eval()
self.model.zero_grad(set_to_none=True)
with torch.enable_grad():
logits = self.model(images)
target_classes = target_classes.to(logits.device, dtype=torch.long)
selected_logits = logits.gather(1, target_classes[:, None]).sum()
selected_logits.backward()
weights = self.gradients.mean(dim=(2, 3), keepdim=True)
cam = torch.relu((weights * self.activations).sum(dim=1, keepdim=True))
cam = F.interpolate(cam, size=images.shape[-2:], mode="bilinear", align_corners=False)
flat = cam.flatten(start_dim=1)
minimum = flat.min(dim=1).values[:, None, None, None]
maximum = flat.max(dim=1).values[:, None, None, None]
cam = (cam - minimum) / (maximum - minimum + 1e-8)
return cam[:, 0].detach(), logits.detach()
def close(self):
self.forward_handle.remove()
def overlay_heatmap(image_tensor, heatmap, heatmap_weight=0.45):
image = to_image(image_tensor)
color = matplotlib.colormaps["inferno"](heatmap.detach().cpu().numpy())[..., :3]
return np.clip((1 - heatmap_weight) * image + heatmap_weight * color, 0, 1)
def load_model(checkpoint_path, device="cpu"):
model = DefectCNN(number_of_classes=len(CLASS_NAMES)).to(device)
state_dict = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(state_dict)
model.eval()
return model
|