tracker/utils/image_saver.py

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)