LAP/torchlap/utils/bb_generator.py

from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F


def _generate_receptive_area(size: np.ndarray,
                             start_ratio: np.ndarray,
                             end_ratio: np.ndarray,
                             receptive_radius: Optional[int] = 0) -> torch.Tensor:
    """
    Generate the map of bounding box's receptive area

    Args:
        size (np.ndarray): The size of the map
        start_ratio (np.ndarray): ratio of the start point
        end_ratio (np.ndarray): ratio of the end point
        receptive_radius (Optional[int]): radius of the receptive field. Default is 0.
    
    Returns:
        torch.Tensor: map of the receptive area
    """
    x_s, y_s = np.floor(size * start_ratio).astype(int)
    x_e, y_e = np.floor(size * end_ratio).astype(int)

    point = torch.FloatTensor([[[
        [0, 0, 0],
        [0, 1, 0],
        [0, 0, 0],
    ]]])
    point = F.interpolate(point, receptive_radius * 2, mode='bilinear', align_corners=True).squeeze()

    map = torch.zeros(*size)
    map[max(x_s-receptive_radius, 0):x_s, max(y_s-receptive_radius, 0):y_s] = point[max(receptive_radius-x_s, 0):receptive_radius, max(receptive_radius-y_s, 0):receptive_radius]
    map[x_e:min(x_e+receptive_radius, size[0]), max(y_s-receptive_radius, 0):y_s] = point[receptive_radius:min(receptive_radius+size[0]-x_e, 2*receptive_radius), max(receptive_radius-y_s, 0):receptive_radius]
    map[max(x_s-receptive_radius, 0):x_s, y_e:min(y_e+receptive_radius, size[1])] = point[max(receptive_radius-x_s, 0):receptive_radius, receptive_radius:min(receptive_radius+size[1]-y_e, 2*receptive_radius)]
    map[x_e:min(x_e+receptive_radius, size[0]), y_e:min(y_e+receptive_radius, size[1])] = point[receptive_radius:min(receptive_radius+size[0]-x_e, 2*receptive_radius), receptive_radius:min(receptive_radius+size[1]-y_e, 2*receptive_radius)]

    horizontal = torch.FloatTensor([[[[0, 1, 0]]]])
    horizontal = F.interpolate(horizontal, (x_e - x_s, receptive_radius * 2), mode='bilinear', align_corners=True).squeeze()
    map[x_s:x_e, max(y_s-receptive_radius, 0):y_s] = horizontal[:, max(receptive_radius-y_s, 0):receptive_radius]
    map[x_s:x_e, y_e:min(y_e+receptive_radius, size[1])] = horizontal[:, receptive_radius:min(receptive_radius+size[1]-y_e, 2*receptive_radius)]

    vertical = torch.FloatTensor([[[
        [0],
        [1],
        [0],
    ]]])
    vertical = F.interpolate(vertical, (receptive_radius * 2, y_e - y_s), mode='bilinear', align_corners=True).squeeze()
    map[max(x_s-receptive_radius, 0):x_s, y_s:y_e] = vertical[max(receptive_radius-x_s, 0):receptive_radius, :]
    map[x_e:min(x_e+receptive_radius, size[0]), y_s:y_e] = vertical[receptive_radius:min(receptive_radius+size[0]-x_e, 2*receptive_radius), :]

    return map


def _generate_bb(size: np.ndarray,
                 start_ratio: np.ndarray,
                 end_ratio: np.ndarray) -> torch.Tensor:
    """
    Generate the map of the bounding box

    Args:
        size (np.ndarray): The size of the map
        start_ratio (np.ndarray): ratio of the start point
        end_ratio (np.ndarray): ratio of the end point

    Returns:
        torch.Tensor: map of the bounding box
    """
    x_s, y_s = np.floor(size * start_ratio).astype(int)
    x_e, y_e = np.floor(size * end_ratio).astype(int)

    map = torch.zeros(*size)
    map[x_s:x_e, y_s:y_e] = 1

    return map


def generate_bb_map(start_ratios: List[Tuple[float, float]],
                    end_ratios: List[Tuple[float, float]],
                    size: Tuple[int, int],
                    receptive_radius: Optional[int] = 0) -> np.ndarray:
    """
    Generate map of the given bounding box

    Args:
        start_ratios (List[Tuple[float, float]]): list of ratios of each bounding box's start point
        end_ratios (List[Tuple[float, float]]): list of ratios of each bounding box's end point
        size (Tuple[int, int]): Size of the map
        receptive_radius (Optional[int]): radius of the receptive field. Default is 0.

    Returns:
        np.ndarray: map of the bounding box
    """

    map = torch.zeros(size)
    size = np.array(size)

    for start_ratio, end_ratio in zip(start_ratios, end_ratios):
        start_ratio = np.array(start_ratio)
        end_ratio = np.array(end_ratio)
        map += _generate_bb(size, start_ratio, end_ratio)
        if receptive_radius > 0:
            map += _generate_receptive_area(size, start_ratio, end_ratio, receptive_radius)

    return map.clamp(0, 1).numpy()