import cv2
import numpy as np
import torch
from collections import defaultdict
def tensor_to_numpy(image):
image_np = (image.numpy() * 255).astype('uint8')
return image_np
def tensor_to_np_float(image):
image_np = image.numpy().astype('float32')
return image_np
def detach_to_cpu(x):
return x.detach().cpu()
def transpose_np(x):
return np.transpose(x, [1, 2, 0])
def tensor_to_gray_im(x):
x = detach_to_cpu(x)
x = tensor_to_numpy(x)
x = transpose_np(x)
return x
def tensor_to_im(x):
x = detach_to_cpu(x).clamp(0, 1)
x = tensor_to_numpy(x)
x = transpose_np(x)
return x
# Predefined key <-> caption dict
key_captions = {
'im': 'Image',
'gt': 'GT',
}
"""
Return an image array with captions
keys in dictionary will be used as caption if not provided
values should contain lists of cv2 images
"""
def get_image_array(images, grid_shape, captions={}):
h, w = grid_shape
cate_counts = len(images)
rows_counts = len(next(iter(images.values())))
font = cv2.FONT_HERSHEY_SIMPLEX
output_image = np.zeros([w * cate_counts, h * (rows_counts + 1), 3], dtype=np.uint8)
col_cnt = 0
for k, v in images.items():
# Default as key value itself
caption = captions.get(k, k)
# Handles new line character
dy = 40
for i, line in enumerate(caption.split('\n')):
cv2.putText(output_image, line, (10, col_cnt * w + 100 + i * dy), font, 0.8,
(255, 255, 255), 2, cv2.LINE_AA)
# Put images
for row_cnt, img in enumerate(v):
im_shape = img.shape
if len(im_shape) == 2:
img = img[..., np.newaxis]
img = (img * 255).astype('uint8')
output_image[(col_cnt + 0) * w:(col_cnt + 1) * w,
(row_cnt + 1) * h:(row_cnt + 2) * h, :] = img
col_cnt += 1
return output_image
def base_transform(im, size):
im = tensor_to_np_float(im)
if len(im.shape) == 3:
im = im.transpose((1, 2, 0))
else:
im = im[:, :, None]
# Resize
if im.shape[1] != size:
im = cv2.resize(im, size, interpolation=cv2.INTER_NEAREST)
return im.clip(0, 1)
def im_transform(im, size):
return base_transform(detach_to_cpu(im), size=size)
def mask_transform(mask, size):
return base_transform(detach_to_cpu(mask), size=size)
def logits_transform(mask, size):
return base_transform(detach_to_cpu(torch.sigmoid(mask)), size=size)
def add_attention(mask, pos):
mask = mask[:, :, None].repeat(3, axis=2)
pos = (pos + 1) / 2
for i in range(pos.shape[0]):
y = int(pos[i][0] * mask.shape[0])
x = int(pos[i][1] * mask.shape[1])
y = max(min(y, mask.shape[0] - 1), 0)
x = max(min(x, mask.shape[1] - 1), 0)
# mask[y, x, :] = (255, 0, 0)
cv2.circle(mask, (x, y), 5, (1, 0, 0), -1)
return mask
def vis(images, size, num_objects):
req_images = defaultdict(list)
b, t = images['rgb'].shape[:2]
# limit the number of images saved
b = min(2, b)
# find max num objects
max_num_objects = max(num_objects[:b])
GT_suffix = ''
for bi in range(b):
GT_suffix += ' \n%s' % images['info']['name'][bi][-25:-4]
for bi in range(b):
for ti in range(t):
req_images['RGB'].append(im_transform(images['rgb'][bi, ti], size))
aux = images[f'aux_{max(ti, 1)}'] # no aux_0, use aux_1 for shape
if 'sensory_logits' in aux:
sensory_aux = aux['sensory_logits'][bi].softmax(dim=0)
# batch_size * num_objects * num_levels * H * W
q_mask_aux = aux['q_logits'][bi].softmax(dim=0)
num_levels = q_mask_aux.shape[1]
for oi in range(max_num_objects):
if ti == 0 or oi >= num_objects[bi]:
req_images[f'Mask_{oi}'].append(
mask_transform(images['first_frame_gt'][bi][0, oi], size))
req_images[f'S-Aux_{oi}'].append(
mask_transform(images['first_frame_gt'][bi][0, oi], size))
for l in range(num_levels):
req_images[f'Q-Aux-L{l}_{oi}'].append(
mask_transform(images['first_frame_gt'][bi][0, oi], size))
else:
mask = mask_transform(images[f'masks_{ti}'][bi][oi], size)
req_images[f'Mask_{oi}'].append(mask)
if 'sensory_logits' in aux:
req_images[f'S-Aux_{oi}'].append(mask_transform(sensory_aux[oi + 1], size))
for l in range(num_levels):
mask = mask_transform(q_mask_aux[oi + 1, l], size)
req_images[f'Q-Aux-L{l}_{oi}'].append(mask)
req_images[f'GT_{oi}_{GT_suffix}'].append(
mask_transform(images['cls_gt'][bi, ti, 0] == (oi + 1), size))
return get_image_array(req_images, size, key_captions)
def vis_debug(images, size, num_objects):
req_images = defaultdict(list)
b, t = images['rgb'].shape[:2]
# limit the number of images saved
b = min(2, b)
# find max num objects
max_num_objects = max(num_objects[:b])
GT_suffix = ''
for bi in range(b):
GT_suffix += ' \n%s' % images['info']['name'][bi][-25:-4]
for bi in range(b):
for ti in range(t):
req_images['RGB'].append(im_transform(images['rgb'][bi, ti], size))
aux = images[f'aux_{max(ti, 1)}'] # no aux_0, use aux_1 for shape
sensory_aux = aux['sensory_logits'][bi].softmax(dim=0)
# batch_size * num_objects * num_levels * H * W
q_mask_aux = aux['q_logits'][bi].softmax(dim=0)
attn_mask = aux['attn_mask'][bi]
num_levels = q_mask_aux.shape[1]
num_queries = attn_mask.shape[1]
for oi in range(max_num_objects):
if ti == 0 or oi >= num_objects[bi]:
req_images[f'Mask_{oi}'].append(
mask_transform(images['first_frame_gt'][bi][0, oi], size))
req_images[f'S-Aux_{oi}'].append(
mask_transform(images['first_frame_gt'][bi][0, oi], size))
for l in range(num_levels):
req_images[f'Q-Aux-L{l}_{oi}'].append(
mask_transform(images['first_frame_gt'][bi][0, oi], size))
for q in range(num_queries):
req_images[f'Attn-Mask-Q{q}_{oi}'].append(
mask_transform(images['first_frame_gt'][bi][0, oi], size))
else:
mask = mask_transform(images[f'masks_{ti}'][bi][oi], size)
req_images[f'Mask_{oi}'].append(mask)
req_images[f'S-Aux_{oi}'].append(mask_transform(sensory_aux[oi + 1], size))
for l in range(num_levels):
mask = mask_transform(q_mask_aux[oi + 1, l], size)
req_images[f'Q-Aux-L{l}_{oi}'].append(mask)
for q in range(num_queries):
mask = mask_transform(1 - attn_mask[oi, q].float(), size)
req_images[f'Attn-Mask-Q{q}_{oi}'].append(mask)
req_images[f'GT_{oi}_{GT_suffix}'].append(
mask_transform(images['cls_gt'][bi, ti, 0] == (oi + 1), size))
return get_image_array(req_images, size, key_captions)