--- a +++ b/app/backend/Utils.py @@ -0,0 +1,73 @@ +#############################GRAD-Cam############################# + +#@title Grad-CAM code with full-credit to Jimin Tan (https://github.com/tanjimin/grad-cam-pytorch-light) +import torch ,cv2 +from PIL import Image +import numpy as np +import os +class InfoHolder(): + + def __init__(self, heatmap_layer): + self.gradient = None + self.activation = None + self.heatmap_layer = heatmap_layer + + def get_gradient(self, grad): + self.gradient = grad + + def hook(self, model, input, output): + if output.requires_grad: + output.register_hook(self.get_gradient) + self.activation = output.detach() + +def generate_heatmap(weighted_activation): + raw_heatmap = torch.mean(weighted_activation, 0) + heatmap = np.maximum(raw_heatmap.detach().cpu(), 0) + heatmap /= torch.max(heatmap) + 1e-10 + return heatmap.numpy() + +def superimpose(input_img, heatmap): + img = to_RGB(input_img) + heatmap = cv2.resize(heatmap, (img.shape[0], img.shape[1])) + heatmap = np.uint8(255 * heatmap) + heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) + superimposed_img = np.uint8(heatmap * 0.6 + img * 0.4) + pil_img = cv2.cvtColor(superimposed_img,cv2.COLOR_BGR2RGB) + return pil_img + +def to_RGB(tensor): + tensor = (tensor - tensor.min()) + tensor = tensor/(tensor.max() + 1e-10) + image_binary = np.transpose(tensor.numpy(), (1, 2, 0)) + image = np.uint8(255 * image_binary) + return image + +def grad_cam(model, input_tensor, heatmap_layer, truelabel=None): + + info = InfoHolder(heatmap_layer) + heatmap_layer.register_forward_hook(info.hook) + + output = model(input_tensor.unsqueeze(0))[0] + truelabel = truelabel if truelabel else torch.argmax(output) + + output[truelabel].backward() + + weights = torch.mean(info.gradient, [0, 2, 3]) + activation = info.activation.squeeze(0) + + weighted_activation = torch.zeros(activation.shape) + for idx, (weight, activation) in enumerate(zip(weights, activation)): + weighted_activation[idx] = weight * activation + heatmap = generate_heatmap(weighted_activation) + return superimpose(input_tensor, heatmap) + +def use_gradcam(img_path,dest_path,model,transforms): + image=Image.open(img_path).convert('RGB') + layer4=model[0][-1] + heatmap_layer=layer4[2].conv2 + input_tensor=transforms(image) + + #get filename without extension + filename=os.path.splitext(os.path.basename(img_path))[0] + grad_cam_image= grad_cam(model, input_tensor, heatmap_layer) + return cv2.imwrite(os.path.join(dest_path,'(gradcam)'+filename+'.png'),grad_cam_image) \ No newline at end of file