49 lines
2.1 KiB
Python
49 lines
2.1 KiB
Python
import torchvision.utils
|
|
from matplotlib.pyplot import get_cmap
|
|
import torch
|
|
import warnings
|
|
from torch.nn.functional import interpolate
|
|
|
|
|
|
def attention_colored_map(attentions, size=None, cmap_name="jet"):
|
|
assert attentions.dim() == 4 and attentions.size(1) == 1
|
|
|
|
min_attentions = attentions.view(attentions.size(0), -1).min(1, keepdim=True)[0].view(attentions.size(0), 1, 1, 1)
|
|
attentions -= min_attentions
|
|
attentions /= attentions.view(attentions.size(0), -1).max(1, keepdim=True)[0].view(attentions.size(0), 1, 1, 1)
|
|
|
|
if size is not None and attentions.size()[-2:] != size:
|
|
assert len(size) == 2, "for interpolate, size must be (x, y), have two dim"
|
|
attentions = interpolate(attentions, size, mode="bilinear", align_corners=False)
|
|
cmap = get_cmap(cmap_name)
|
|
ca = cmap(attentions.squeeze(1).cpu())[:, :, :, :3]
|
|
return torch.from_numpy(ca).permute(0, 3, 1, 2).contiguous()
|
|
|
|
|
|
def fuse_attention_map(images, attentions, cmap_name="jet", alpha=0.5):
|
|
"""
|
|
|
|
:param images: B x H x W
|
|
:param attentions: B x Ha x Wa
|
|
:param cmap_name:
|
|
:param alpha:
|
|
:return:
|
|
"""
|
|
if attentions.size(0) != images.size(0):
|
|
warnings.warn(f"attentions: {attentions.size()} and images: {images.size} do not have same batch_size")
|
|
return images
|
|
if attentions.size(1) != 1:
|
|
warnings.warn(f"attentions's channels should be 1 but got {attentions.size(1)}")
|
|
return images
|
|
colored_attentions = attention_colored_map(attentions, images.size()[-2:], cmap_name).to(images.device)
|
|
return images * alpha + colored_attentions * (1 - alpha)
|
|
|
|
|
|
def make_2d_grid(tensors, padding=0, normalize=True, range=None, scale_each=False, pad_value=0):
|
|
# merge image in a batch in `y` direction first.
|
|
grids = [torchvision.utils.make_grid(img_batch, padding=padding, nrow=1, normalize=normalize, range=range,
|
|
scale_each=scale_each, pad_value=pad_value)
|
|
for img_batch in tensors]
|
|
# merge images in `x` direction.
|
|
return torchvision.utils.make_grid(grids, padding=0, nrow=len(grids))
|