Spaces:
Sleeping
Sleeping
File size: 949 Bytes
46fdf2a |
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 |
from PIL import Image
import json
import torch
from torchvision import transforms
import cv2
import numpy as np
import os
import torch.nn as nn
def show_cam_on_img(img, mask, img_path_save):
heat_map = cv2.applyColorMap(np.uint8(255*mask), cv2.COLORMAP_JET)
heat_map = np.float32(heat_map) / 255
cam = heat_map + np.float32(img)
cam = cam / np.max(cam)
cv2.imwrite(img_path_save, np.uint8(255 * cam))
img_path_read = ""
img_path_save = ""
def main():
img = cv2.imread(img_path_read, flags=1)
img = np.float32(cv2.resize(img, (224, 224))) / 255
# cam_all is the score tensor of shape (B, C, H, W), similar to y_raw in out Figure 1
# cls_idx specifying the i-th class out of C class
# visualize the 0's class heatmap
cls_idx = 0
cam = cam_all[cls_idx]
# cam = nn.ReLU()(cam)
cam = cam / torch.max(cam)
cam = cv2.resize(np.array(cam), (224, 224))
show_cam_on_img(img, cam, img_path_save)
|