File size: 3,169 Bytes
6cd90b7 |
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 |
# coding=utf-8
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Visualization functions."""
import os
import cv2
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torch
# pylint: disable=g-importing-member
from utils.utils import normalize
_VIS_HEIGHT = 512
_VIS_WIDTH = 512
def show_cam_on_image(img, mask):
if img.shape[1] != mask.shape[1]:
mask = cv2.resize(mask, (img.shape[1], img.shape[0]))
heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
heatmap = np.float32(heatmap) / 255
cam = heatmap + np.float32(img)
cam = cam / np.max(cam)
cam = np.uint8(255 * cam)
return cam
def save_img(array, img_name):
numpy_array = array.astype(np.uint8)
image = Image.fromarray(numpy_array, mode="RGB")
image.save(f"{img_name}.png")
def viz_attn(img, attn_map, prefix="vis_results/clipcam_img", img_name="cam"):
"""Visualize attention map."""
num_masks = 1
if len(attn_map.shape) == 3:
num_masks = attn_map.shape[0]
attn_map = attn_map.float().squeeze(1).detach().cpu().numpy()
attn_map = normalize(attn_map)
img = normalize(img)
if num_masks == 1:
vis = show_cam_on_image(img, attn_map)
if not os.path.exists(prefix):
os.makedirs(prefix)
save_img(vis, os.path.join(prefix, f"{img_name}"))
return vis
for i in range(num_masks):
vis = show_cam_on_image(img, attn_map[i])
if not os.path.exists(prefix):
os.makedirs(prefix)
save_img(vis, os.path.join(prefix, f"{img_name}_{i}"))
def vis_mask(mask, gt_mask, img, output_dir, fname):
"""Visualize mask."""
mask_img = torch.zeros((_VIS_WIDTH, _VIS_HEIGHT))
mask_img[mask[0]] = 1
# print(gt_mask.shape, img.size())
# Assume img and gt_mask are also torch.Tensor with size (512, 512)
img = img[0].permute(1, 2, 0).numpy()
gt_mask_img = torch.zeros((_VIS_WIDTH, _VIS_HEIGHT))
gt_mask_img[gt_mask[0]] = 1
_, axs = plt.subplots(
1, 3, figsize=(15, 5)
) # change the figsize if necessary
axs[0].imshow(img) # if image is grayscale, otherwise remove cmap argument
axs[0].axis("off")
axs[0].set_title("Original Image")
axs[1].imshow(
mask_img.numpy(), cmap="jet", alpha=0.5
) # using alpha for transparency
axs[1].axis("off")
axs[1].set_title("Mask")
axs[2].imshow(
gt_mask_img.numpy(), cmap="jet", alpha=0.5
) # using alpha for transparency
axs[2].axis("off")
axs[2].set_title("Ground Truth Mask")
plt.savefig(
os.path.join(output_dir, f"{fname}.jpg"),
bbox_inches="tight",
dpi=300,
pad_inches=0.0,
)
|