mmdetection/mmdet/structures/bbox/horizontal_boxes.py

# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Tuple, TypeVar, Union

import cv2
import numpy as np
import torch
from torch import BoolTensor, Tensor

from mmdet.structures.mask.structures import BitmapMasks, PolygonMasks
from .base_boxes import BaseBoxes
from .bbox_overlaps import bbox_overlaps
from .box_type import register_box

T = TypeVar('T')
DeviceType = Union[str, torch.device]
MaskType = Union[BitmapMasks, PolygonMasks]


@register_box(name='hbox')
class HorizontalBoxes(BaseBoxes):
    """The horizontal box class used in MMDetection by default.

    The ``box_dim`` of ``HorizontalBoxes`` is 4, which means the length of
    the last dimension of the data should be 4. Two modes of box data are
    supported in ``HorizontalBoxes``:

    - 'xyxy': Each row of data indicates (x1, y1, x2, y2), which are the
      coordinates of the left-top and right-bottom points.
    - 'cxcywh': Each row of data indicates (x, y, w, h), where (x, y) are the
      coordinates of the box centers and (w, h) are the width and height.

    ``HorizontalBoxes`` only restores 'xyxy' mode of data. If the the data is
    in 'cxcywh' mode, users need to input ``in_mode='cxcywh'`` and The code
    will convert the 'cxcywh' data to 'xyxy' automatically.

    Args:
        data (Tensor or np.ndarray or Sequence): The box data with shape of
            (..., 4).
        dtype (torch.dtype, Optional): data type of boxes. Defaults to None.
        device (str or torch.device, Optional): device of boxes.
            Default to None.
        clone (bool): Whether clone ``boxes`` or not. Defaults to True.
        mode (str, Optional): the mode of boxes. If it is 'cxcywh', the
            `data` will be converted to 'xyxy' mode. Defaults to None.
    """

    box_dim: int = 4

    def __init__(self,
                 data: Union[Tensor, np.ndarray],
                 dtype: torch.dtype = None,
                 device: DeviceType = None,
                 clone: bool = True,
                 in_mode: Optional[str] = None) -> None:
        super().__init__(data=data, dtype=dtype, device=device, clone=clone)
        if isinstance(in_mode, str):
            if in_mode not in ('xyxy', 'cxcywh'):
                raise ValueError(f'Get invalid mode {in_mode}.')
            if in_mode == 'cxcywh':
                self.tensor = self.cxcywh_to_xyxy(self.tensor)

    @staticmethod
    def cxcywh_to_xyxy(boxes: Tensor) -> Tensor:
        """Convert box coordinates from (cx, cy, w, h) to (x1, y1, x2, y2).

        Args:
            boxes (Tensor): cxcywh boxes tensor with shape of (..., 4).

        Returns:
            Tensor: xyxy boxes tensor with shape of (..., 4).
        """
        ctr, wh = boxes.split((2, 2), dim=-1)
        return torch.cat([(ctr - wh / 2), (ctr + wh / 2)], dim=-1)

    @staticmethod
    def xyxy_to_cxcywh(boxes: Tensor) -> Tensor:
        """Convert box coordinates from (x1, y1, x2, y2) to (cx, cy, w, h).

        Args:
            boxes (Tensor): xyxy boxes tensor with shape of (..., 4).

        Returns:
            Tensor: cxcywh boxes tensor with shape of (..., 4).
        """
        xy1, xy2 = boxes.split((2, 2), dim=-1)
        return torch.cat([(xy2 + xy1) / 2, (xy2 - xy1)], dim=-1)

    @property
    def cxcywh(self) -> Tensor:
        """Return a tensor representing the cxcywh boxes."""
        return self.xyxy_to_cxcywh(self.tensor)

    @property
    def centers(self) -> Tensor:
        """Return a tensor representing the centers of boxes."""
        boxes = self.tensor
        return (boxes[..., :2] + boxes[..., 2:]) / 2

    @property
    def areas(self) -> Tensor:
        """Return a tensor representing the areas of boxes."""
        boxes = self.tensor
        return (boxes[..., 2] - boxes[..., 0]) * (
            boxes[..., 3] - boxes[..., 1])

    @property
    def widths(self) -> Tensor:
        """Return a tensor representing the widths of boxes."""
        boxes = self.tensor
        return boxes[..., 2] - boxes[..., 0]

    @property
    def heights(self) -> Tensor:
        """Return a tensor representing the heights of boxes."""
        boxes = self.tensor
        return boxes[..., 3] - boxes[..., 1]

    def flip_(self,
              img_shape: Tuple[int, int],
              direction: str = 'horizontal') -> None:
        """Flip boxes horizontally or vertically in-place.

        Args:
            img_shape (Tuple[int, int]): A tuple of image height and width.
            direction (str): Flip direction, options are "horizontal",
                "vertical" and "diagonal". Defaults to "horizontal"
        """
        assert direction in ['horizontal', 'vertical', 'diagonal']
        flipped = self.tensor
        boxes = flipped.clone()
        if direction == 'horizontal':
            flipped[..., 0] = img_shape[1] - boxes[..., 2]
            flipped[..., 2] = img_shape[1] - boxes[..., 0]
        elif direction == 'vertical':
            flipped[..., 1] = img_shape[0] - boxes[..., 3]
            flipped[..., 3] = img_shape[0] - boxes[..., 1]
        else:
            flipped[..., 0] = img_shape[1] - boxes[..., 2]
            flipped[..., 1] = img_shape[0] - boxes[..., 3]
            flipped[..., 2] = img_shape[1] - boxes[..., 0]
            flipped[..., 3] = img_shape[0] - boxes[..., 1]

    def translate_(self, distances: Tuple[float, float]) -> None:
        """Translate boxes in-place.

        Args:
            distances (Tuple[float, float]): translate distances. The first
                is horizontal distance and the second is vertical distance.
        """
        boxes = self.tensor
        assert len(distances) == 2
        self.tensor = boxes + boxes.new_tensor(distances).repeat(2)

    def clip_(self, img_shape: Tuple[int, int]) -> None:
        """Clip boxes according to the image shape in-place.

        Args:
            img_shape (Tuple[int, int]): A tuple of image height and width.
        """
        boxes = self.tensor
        boxes[..., 0::2] = boxes[..., 0::2].clamp(0, img_shape[1])
        boxes[..., 1::2] = boxes[..., 1::2].clamp(0, img_shape[0])

    def rotate_(self, center: Tuple[float, float], angle: float) -> None:
        """Rotate all boxes in-place.

        Args:
            center (Tuple[float, float]): Rotation origin.
            angle (float): Rotation angle represented in degrees. Positive
                values mean clockwise rotation.
        """
        boxes = self.tensor
        rotation_matrix = boxes.new_tensor(
            cv2.getRotationMatrix2D(center, -angle, 1))

        corners = self.hbox2corner(boxes)
        corners = torch.cat(
            [corners, corners.new_ones(*corners.shape[:-1], 1)], dim=-1)
        corners_T = torch.transpose(corners, -1, -2)
        corners_T = torch.matmul(rotation_matrix, corners_T)
        corners = torch.transpose(corners_T, -1, -2)
        self.tensor = self.corner2hbox(corners)

    def project_(self, homography_matrix: Union[Tensor, np.ndarray]) -> None:
        """Geometric transformat boxes in-place.

        Args:
            homography_matrix (Tensor or np.ndarray]):
                Shape (3, 3) for geometric transformation.
        """
        boxes = self.tensor
        if isinstance(homography_matrix, np.ndarray):
            homography_matrix = boxes.new_tensor(homography_matrix)
        corners = self.hbox2corner(boxes)
        corners = torch.cat(
            [corners, corners.new_ones(*corners.shape[:-1], 1)], dim=-1)
        corners_T = torch.transpose(corners, -1, -2)
        corners_T = torch.matmul(homography_matrix, corners_T)
        corners = torch.transpose(corners_T, -1, -2)
        # Convert to homogeneous coordinates by normalization
        corners = corners[..., :2] / corners[..., 2:3]
        self.tensor = self.corner2hbox(corners)

    @staticmethod
    def hbox2corner(boxes: Tensor) -> Tensor:
        """Convert box coordinates from (x1, y1, x2, y2) to corners ((x1, y1),
        (x2, y1), (x1, y2), (x2, y2)).

        Args:
            boxes (Tensor): Horizontal box tensor with shape of (..., 4).

        Returns:
            Tensor: Corner tensor with shape of (..., 4, 2).
        """
        x1, y1, x2, y2 = torch.split(boxes, 1, dim=-1)
        corners = torch.cat([x1, y1, x2, y1, x1, y2, x2, y2], dim=-1)
        return corners.reshape(*corners.shape[:-1], 4, 2)

    @staticmethod
    def corner2hbox(corners: Tensor) -> Tensor:
        """Convert box coordinates from corners ((x1, y1), (x2, y1), (x1, y2),
        (x2, y2)) to (x1, y1, x2, y2).

        Args:
            corners (Tensor): Corner tensor with shape of (..., 4, 2).

        Returns:
            Tensor: Horizontal box tensor with shape of (..., 4).
        """
        if corners.numel() == 0:
            return corners.new_zeros((0, 4))
        min_xy = corners.min(dim=-2)[0]
        max_xy = corners.max(dim=-2)[0]
        return torch.cat([min_xy, max_xy], dim=-1)

    def rescale_(self, scale_factor: Tuple[float, float]) -> None:
        """Rescale boxes w.r.t. rescale_factor in-place.

        Note:
            Both ``rescale_`` and ``resize_`` will enlarge or shrink boxes
            w.r.t ``scale_facotr``. The difference is that ``resize_`` only
            changes the width and the height of boxes, but ``rescale_`` also
            rescales the box centers simultaneously.

        Args:
            scale_factor (Tuple[float, float]): factors for scaling boxes.
                The length should be 2.
        """
        boxes = self.tensor
        assert len(scale_factor) == 2
        scale_factor = boxes.new_tensor(scale_factor).repeat(2)
        self.tensor = boxes * scale_factor

    def resize_(self, scale_factor: Tuple[float, float]) -> None:
        """Resize the box width and height w.r.t scale_factor in-place.

        Note:
            Both ``rescale_`` and ``resize_`` will enlarge or shrink boxes
            w.r.t ``scale_facotr``. The difference is that ``resize_`` only
            changes the width and the height of boxes, but ``rescale_`` also
            rescales the box centers simultaneously.

        Args:
            scale_factor (Tuple[float, float]): factors for scaling box
                shapes. The length should be 2.
        """
        boxes = self.tensor
        assert len(scale_factor) == 2
        ctrs = (boxes[..., 2:] + boxes[..., :2]) / 2
        wh = boxes[..., 2:] - boxes[..., :2]
        scale_factor = boxes.new_tensor(scale_factor)
        wh = wh * scale_factor
        xy1 = ctrs - 0.5 * wh
        xy2 = ctrs + 0.5 * wh
        self.tensor = torch.cat([xy1, xy2], dim=-1)

    def is_inside(self,
                  img_shape: Tuple[int, int],
                  all_inside: bool = False,
                  allowed_border: int = 0) -> BoolTensor:
        """Find boxes inside the image.

        Args:
            img_shape (Tuple[int, int]): A tuple of image height and width.
            all_inside (bool): Whether the boxes are all inside the image or
                part inside the image. Defaults to False.
            allowed_border (int): Boxes that extend beyond the image shape
                boundary by more than ``allowed_border`` are considered
                "outside" Defaults to 0.
        Returns:
            BoolTensor: A BoolTensor indicating whether the box is inside
            the image. Assuming the original boxes have shape (m, n, 4),
            the output has shape (m, n).
        """
        img_h, img_w = img_shape
        boxes = self.tensor
        if all_inside:
            return (boxes[:, 0] >= -allowed_border) & \
                (boxes[:, 1] >= -allowed_border) & \
                (boxes[:, 2] < img_w + allowed_border) & \
                (boxes[:, 3] < img_h + allowed_border)
        else:
            return (boxes[..., 0] < img_w + allowed_border) & \
                (boxes[..., 1] < img_h + allowed_border) & \
                (boxes[..., 2] > -allowed_border) & \
                (boxes[..., 3] > -allowed_border)

    def find_inside_points(self,
                           points: Tensor,
                           is_aligned: bool = False) -> BoolTensor:
        """Find inside box points. Boxes dimension must be 2.

        Args:
            points (Tensor): Points coordinates. Has shape of (m, 2).
            is_aligned (bool): Whether ``points`` has been aligned with boxes
                or not. If True, the length of boxes and ``points`` should be
                the same. Defaults to False.

        Returns:
            BoolTensor: A BoolTensor indicating whether a point is inside
            boxes. Assuming the boxes has shape of (n, 4), if ``is_aligned``
            is False. The index has shape of (m, n). If ``is_aligned`` is
            True, m should be equal to n and the index has shape of (m, ).
        """
        boxes = self.tensor
        assert boxes.dim() == 2, 'boxes dimension must be 2.'

        if not is_aligned:
            boxes = boxes[None, :, :]
            points = points[:, None, :]
        else:
            assert boxes.size(0) == points.size(0)

        x_min, y_min, x_max, y_max = boxes.unbind(dim=-1)
        return (points[..., 0] >= x_min) & (points[..., 0] <= x_max) & \
            (points[..., 1] >= y_min) & (points[..., 1] <= y_max)

    def create_masks(self, img_shape: Tuple[int, int]) -> BitmapMasks:
        """
        Args:
            img_shape (Tuple[int, int]): A tuple of image height and width.

        Returns:
            :obj:`BitmapMasks`: Converted masks
        """
        img_h, img_w = img_shape
        boxes = self.tensor

        xmin, ymin = boxes[:, 0:1], boxes[:, 1:2]
        xmax, ymax = boxes[:, 2:3], boxes[:, 3:4]
        gt_masks = np.zeros((len(boxes), img_h, img_w), dtype=np.uint8)
        for i in range(len(boxes)):
            gt_masks[i,
                     int(ymin[i]):int(ymax[i]),
                     int(xmin[i]):int(xmax[i])] = 1
        return BitmapMasks(gt_masks, img_h, img_w)

    @staticmethod
    def overlaps(boxes1: BaseBoxes,
                 boxes2: BaseBoxes,
                 mode: str = 'iou',
                 is_aligned: bool = False,
                 eps: float = 1e-6) -> Tensor:
        """Calculate overlap between two set of boxes with their types
        converted to ``HorizontalBoxes``.

        Args:
            boxes1 (:obj:`BaseBoxes`): BaseBoxes with shape of (m, box_dim)
                or empty.
            boxes2 (:obj:`BaseBoxes`): BaseBoxes with shape of (n, box_dim)
                or empty.
            mode (str): "iou" (intersection over union), "iof" (intersection
                over foreground). Defaults to "iou".
            is_aligned (bool): If True, then m and n must be equal. Defaults
                to False.
            eps (float): A value added to the denominator for numerical
                stability. Defaults to 1e-6.

        Returns:
            Tensor: shape (m, n) if ``is_aligned`` is False else shape (m,)
        """
        boxes1 = boxes1.convert_to('hbox')
        boxes2 = boxes2.convert_to('hbox')
        return bbox_overlaps(
            boxes1.tensor,
            boxes2.tensor,
            mode=mode,
            is_aligned=is_aligned,
            eps=eps)

    @staticmethod
    def from_instance_masks(masks: MaskType) -> 'HorizontalBoxes':
        """Create horizontal boxes from instance masks.

        Args:
            masks (:obj:`BitmapMasks` or :obj:`PolygonMasks`): BitmapMasks or
                PolygonMasks instance with length of n.

        Returns:
            :obj:`HorizontalBoxes`: Converted boxes with shape of (n, 4).
        """
        num_masks = len(masks)
        boxes = np.zeros((num_masks, 4), dtype=np.float32)
        if isinstance(masks, BitmapMasks):
            x_any = masks.masks.any(axis=1)
            y_any = masks.masks.any(axis=2)
            for idx in range(num_masks):
                x = np.where(x_any[idx, :])[0]
                y = np.where(y_any[idx, :])[0]
                if len(x) > 0 and len(y) > 0:
                    # use +1 for x_max and y_max so that the right and bottom
                    # boundary of instance masks are fully included by the box
                    boxes[idx, :] = np.array(
                        [x[0], y[0], x[-1] + 1, y[-1] + 1], dtype=np.float32)
        elif isinstance(masks, PolygonMasks):
            for idx, poly_per_obj in enumerate(masks.masks):
                # simply use a number that is big enough for comparison with
                # coordinates
                xy_min = np.array([masks.width * 2, masks.height * 2],
                                  dtype=np.float32)
                xy_max = np.zeros(2, dtype=np.float32)
                for p in poly_per_obj:
                    xy = np.array(p).reshape(-1, 2).astype(np.float32)
                    xy_min = np.minimum(xy_min, np.min(xy, axis=0))
                    xy_max = np.maximum(xy_max, np.max(xy, axis=0))
                boxes[idx, :2] = xy_min
                boxes[idx, 2:] = xy_max
        else:
            raise TypeError(
                '`masks` must be `BitmapMasks`  or `PolygonMasks`, '
                f'but got {type(masks)}.')
        return HorizontalBoxes(boxes)