Rangement

salvador/grad_cam.py

@@ -1,102 +0,0 @@
-import torch
-import numpy as np
-import torchvision
-from PIL import Image
-from torch import topk
-from torch import nn
-import torch.nn.functional as F
-from torch import topk
-import cv2
-from torchvision import transforms
-import os
-
-class Lambda(nn.Module):
-    "Create a layer that simply calls `func` with `x`"
-    def __init__(self, func): 
-        super().__init__()
-        self.func=func
-    def forward(self, x): return self.func(x)
-
-class SaveFeatures():
-    activations, gradients = None, None
-    def __init__(self, m): 
-        self.forward = m.register_forward_hook(self.forward_hook_fn)
-        self.backward = m.register_backward_hook(self.backward_hook_fn)
-
-    def forward_hook_fn(self, module, input, output): 
-        self.activations = output.cpu().detach()
-
-    def backward_hook_fn(self, module, grad_input, grad_output): 
-        self.gradients = grad_output[0].cpu().detach()
-
-    def remove(self): 
-        self.forward.remove()
-        self.backward.remove()
-
-def main(cam):
-    device = 'cuda:0'
-    model_name = 'resnet50'
-    root = '/mnt/md0/data/cifar10/tmp/cifar/train'
-    _root = 'cifar'
-    
-    os.makedirs(os.path.join(_root + '_CAM'), exist_ok=True)
-    os.makedirs(os.path.join(_root + '_CAM'), exist_ok=True)
-
-    train_transform = transforms.Compose([
-        transforms.ToTensor(),
-    ])
-
-    dataset = torchvision.datasets.ImageFolder(
-        root=root, transform=train_transform,
-    )
-
-    loader = torch.utils.data.DataLoader(dataset, batch_size=1)
-    model = torchvision.models.__dict__[model_name](pretrained=True)
-    flat = list(model.children())
-    body, head = nn.Sequential(*flat[:-2]), nn.Sequential(flat[-2], Lambda(func=lambda x: torch.flatten(x, 1)), nn.Linear(flat[-1].in_features, len(loader.dataset.classes))) 
-    model = nn.Sequential(body, head)
-
-    model.load_state_dict(torch.load('checkpoint.pt', map_location=lambda storage, loc: storage))
-    model = model.to(device)
-    model.eval()
-
-    activated_features = SaveFeatures(model[0])
-
-    for i, (img, target ) in enumerate(loader):
-        img = img.to(device)
-        pred = model(img)
-        import ipdb; ipdb.set_trace()
-        # get the gradient of the output with respect to the parameters of the model
-        pred[:, target.item()].backward()
-
-        # import ipdb; ipdb.set_trace()
-        # pull the gradients out of the model
-        gradients = activated_features.gradients[0]
-
-        pooled_gradients = gradients.mean(1).mean(1)
-
-        # get the activations of the last convolutional layer
-        activations = activated_features.activations[0]
-
-        heatmap = F.relu(((activations*pooled_gradients[...,None,None])).sum(0))
-        heatmap /= torch.max(heatmap)
-
-        heatmap = heatmap.numpy()
-
-        
-        image = cv2.imread(dataset.imgs[i][0])
-        heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))
-        heatmap = np.uint8(255 * heatmap)
-        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-        # superimposed_img = heatmap * 0.3 + image * 0.5
-        superimposed_img = heatmap 
-
-        clss = dataset.imgs[i][0].split(os.sep)[1]
-        name = dataset.imgs[i][0].split(os.sep)[2].split('.')[0]
-        cv2.imwrite(os.path.join(_root+"_CAM", name + '.jpg'), superimposed_img)
-        print(f'{os.path.join(_root+"_CAM", name + ".jpg")} saved')
-    
-    activated_features.remove()
-
-if __name__ == "__main__":
-    main(cam=True)