File size: 1,881 Bytes
c4b2b37 |
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 |
#
# Copyright (c) 2019 Idiap Research Institute, http://www.idiap.ch/
# Written by Suraj Srinivas <suraj.srinivas@idiap.ch>
#
""" Misc helper functions """
import subprocess
import cv2
import numpy as np
import torch
import torchvision.transforms as transforms
class NormalizeInverse(transforms.Normalize):
# Undo normalization on images
def __init__(self, mean, std):
mean = torch.as_tensor(mean)
std = torch.as_tensor(std)
std_inv = 1 / (std + 1e-7)
mean_inv = -mean * std_inv
super(NormalizeInverse, self).__init__(mean=mean_inv, std=std_inv)
def __call__(self, tensor):
return super(NormalizeInverse, self).__call__(tensor.clone())
def create_folder(folder_name):
try:
subprocess.call(["mkdir", "-p", folder_name])
except OSError:
None
def save_saliency_map(image, saliency_map, filename):
"""
Save saliency map on image.
Args:
image: Tensor of size (3,H,W)
saliency_map: Tensor of size (1,H,W)
filename: string with complete path and file extension
"""
image = image.data.cpu().numpy()
saliency_map = saliency_map.data.cpu().numpy()
saliency_map = saliency_map - saliency_map.min()
saliency_map = saliency_map / saliency_map.max()
saliency_map = saliency_map.clip(0, 1)
saliency_map = np.uint8(saliency_map * 255).transpose(1, 2, 0)
saliency_map = cv2.resize(saliency_map, (224, 224))
image = np.uint8(image * 255).transpose(1, 2, 0)
image = cv2.resize(image, (224, 224))
# Apply JET colormap
color_heatmap = cv2.applyColorMap(saliency_map, cv2.COLORMAP_JET)
# Combine image with heatmap
img_with_heatmap = np.float32(color_heatmap) + np.float32(image)
img_with_heatmap = img_with_heatmap / np.max(img_with_heatmap)
cv2.imwrite(filename, np.uint8(255 * img_with_heatmap))
|