#!/usr/bin/env python3
#
# Copyright 2021 Graviti. Licensed under MIT License.
#
"""The implementation of the TensorBay bounding box."""
import math
import warnings
from typing import Dict, Iterable, Mapping, Optional, Tuple, Type, TypeVar
from tensorbay.geometry.transform import Transform3D
from tensorbay.geometry.vector import Vector2D, Vector3D
from tensorbay.utility import MatrixType, ReprMixin, ReprType, UserSequence, common_loads
with warnings.catch_warnings():
warnings.simplefilter("ignore")
from quaternion import quaternion
_B2 = TypeVar("_B2", bound="Box2D")
_B3 = TypeVar("_B3", bound="Box3D")
[docs]class Box2D(UserSequence[float]):
"""This class defines the concept of Box2D.
:class:`Box2D` contains the information of a 2D bounding box, such as the coordinates,
width and height.
It provides :meth:`Box2D.iou` to calculate the intersection over union of two 2D boxes.
Arguments:
xmin: The x coordinate of the top-left vertex of the 2D box.
ymin: The y coordinate of the top-left vertex of the 2D box.
xmax: The x coordinate of the bottom-right vertex of the 2D box.
ymax: The y coordinate of the bottom-right vertex of the 2D box.
Examples:
>>> Box2D(1, 2, 3, 4)
Box2D(1, 2, 3, 4)
"""
_repr_type = ReprType.INSTANCE
_LENGTH = 4
def __init__(
self,
xmin: float,
ymin: float,
xmax: float,
ymax: float,
) -> None:
if xmin >= xmax or ymin >= ymax:
self._data = (0.0,) * Box2D._LENGTH
else:
self._data = (xmin, ymin, xmax, ymax)
def __len__(self) -> int:
return Box2D._LENGTH
def __and__(self, other: "Box2D") -> "Box2D":
"""Calculate the intersect box of two boxes.
Arguments:
other: The other box.
Returns:
The intersect box of the two boxes.
"""
try:
xmin = max(self._data[0], other._data[0])
ymin = max(self._data[1], other._data[1])
xmax = min(self._data[2], other._data[2])
ymax = min(self._data[3], other._data[3])
return Box2D(xmin, ymin, xmax, ymax)
except (TypeError, IndexError, AttributeError):
return NotImplemented
def _loads(self, contents: Mapping[str, float]) -> None:
self._data = (contents["xmin"], contents["ymin"], contents["xmax"], contents["ymax"])
def _repr_head(self) -> str:
"""Return basic information of the Box2D.
Returns:
Basic information of the Box2D.
"""
return f"{self.__class__.__name__}{self._data}"
[docs] @staticmethod
def iou(box1: "Box2D", box2: "Box2D") -> float:
"""Calculate the intersection over union of two 2D boxes.
Arguments:
box1: A 2D box.
box2: A 2D box.
Returns:
The intersection over union between the two input boxes.
Examples:
>>> box2d_1 = Box2D(1, 2, 3, 4)
>>> box2d_2 = Box2D(2, 2, 3, 4)
>>> Box2D.iou(box2d_1, box2d_2)
0.5
"""
area1 = box1.area()
area2 = box2.area()
intersect_box = box1 & box2
intersect = intersect_box.area()
union = area1 + area2 - intersect
return intersect / union
[docs] @classmethod
def from_xywh(cls: Type[_B2], x: float, y: float, width: float, height: float) -> _B2:
"""Create a :class:`Box2D` instance from the top-left vertex and the width and the height.
Arguments:
x: X coordinate of the top left vertex of the box.
y: Y coordinate of the top left vertex of the box.
width: Length of the box along the x axis.
height: Length of the box along the y axis.
Returns:
The created :class:`Box2D` instance.
Examples:
>>> Box2D.from_xywh(1, 2, 3, 4)
Box2D(1, 2, 4, 6)
"""
return cls(x, y, x + width, y + height)
[docs] @classmethod
def loads(cls: Type[_B2], contents: Mapping[str, float]) -> _B2:
"""Load a :class:`Box2D` from a dict containing coordinates of the 2D box.
Arguments:
contents: A dict containing coordinates of a 2D box.
Returns:
The loaded :class:`Box2D` object.
Examples:
>>> contents = {"xmin": 1.0, "ymin": 2.0, "xmax": 3.0, "ymax": 4.0}
>>> Box2D.loads(contents)
Box2D(1.0, 2.0, 3.0, 4.0)
"""
return common_loads(cls, contents)
@property
def xmin(self) -> float:
"""Return the minimum x coordinate.
Returns:
Minimum x coordinate.
Examples:
>>> box2d = Box2D(1, 2, 3, 4)
>>> box2d.xmin
1
"""
return self._data[0]
@property
def ymin(self) -> float:
"""Return the minimum y coordinate.
Returns:
Minimum y coordinate.
Examples:
>>> box2d = Box2D(1, 2, 3, 4)
>>> box2d.ymin
2
"""
return self._data[1]
@property
def xmax(self) -> float:
"""Return the maximum x coordinate.
Returns:
Maximum x coordinate.
Examples:
>>> box2d = Box2D(1, 2, 3, 4)
>>> box2d.xmax
3
"""
return self._data[2]
@property
def ymax(self) -> float:
"""Return the maximum y coordinate.
Returns:
Maximum y coordinate.
Examples:
>>> box2d = Box2D(1, 2, 3, 4)
>>> box2d.ymax
4
"""
return self._data[3]
@property
def tl(self) -> Vector2D: # pylint: disable=invalid-name
"""Return the top left point.
Returns:
The top left point.
Examples:
>>> box2d = Box2D(1, 2, 3, 4)
>>> box2d.tl
Vector2D(1, 2)
"""
return Vector2D(self._data[0], self._data[1])
@property
def br(self) -> Vector2D: # pylint: disable=invalid-name
"""Return the bottom right point.
Returns:
The bottom right point.
Examples:
>>> box2d = Box2D(1, 2, 3, 4)
>>> box2d.br
Vector2D(3, 4)
"""
return Vector2D(self._data[2], self._data[3])
@property
def width(self) -> float:
"""Return the width of the 2D box.
Returns:
The width of the 2D box.
Examples:
>>> box2d = Box2D(1, 2, 3, 6)
>>> box2d.width
2
"""
return self._data[2] - self._data[0]
@property
def height(self) -> float:
"""Return the height of the 2D box.
Returns:
The height of the 2D box.
Examples:
>>> box2d = Box2D(1, 2, 3, 6)
>>> box2d.height
4
"""
return self._data[3] - self._data[1]
[docs] def dumps(self) -> Dict[str, float]:
"""Dumps a 2D box into a dict.
Returns:
A dict containing vertex coordinates of the box.
Examples:
>>> box2d = Box2D(1, 2, 3, 4)
>>> box2d.dumps()
{'xmin': 1, 'ymin': 2, 'xmax': 3, 'ymax': 4}
"""
return {
"xmin": self._data[0],
"ymin": self._data[1],
"xmax": self._data[2],
"ymax": self._data[3],
}
[docs] def area(self) -> float:
"""Return the area of the 2D box.
Returns:
The area of the 2D box.
Examples:
>>> box2d = Box2D(1, 2, 3, 4)
>>> box2d.area()
4
"""
return self.width * self.height
[docs]class Box3D(ReprMixin):
"""This class defines the concept of Box3D.
:class:`Box3D` contains the information of a 3D bounding box such as the transform,
translation, rotation and size.
It provides :meth:`Box3D.iou` to calculate the intersection over union of two 3D boxes.
Arguments:
translation: Translation in a sequence of [x, y, z].
rotation: Rotation in a sequence of [w, x, y, z] or numpy quaternion.
size: Size in a sequence of [x, y, z].
transform_matrix: A 4x4 or 3x4 transform matrix.
Examples:
*Initialization Method 1:* Init from size, translation and rotation.
>>> Box3D([1, 2, 3], [0, 1, 0, 0], [1, 2, 3])
Box3D(
(size): Vector3D(1, 2, 3)
(translation): Vector3D(1, 2, 3),
(rotation): quaternion(0, 1, 0, 0),
)
*Initialization Method 2:* Init from size and transform matrix.
>>> from tensorbay.geometry import Transform3D
>>> matrix = [[1, 0, 0, 1], [0, 1, 0, 2], [0, 0, 1, 3]]
>>> Box3D(size=[1, 2, 3], transform_matrix=matrix)
Box3D(
(size): Vector3D(1, 2, 3)
(translation): Vector3D(1, 2, 3),
(rotation): quaternion(1, -0, -0, -0),
)
"""
_repr_type = ReprType.INSTANCE
_repr_attrs: Tuple[str, ...] = ("size", "translation", "rotation")
def __init__(
self,
size: Iterable[float],
translation: Iterable[float] = (0, 0, 0),
rotation: Transform3D.RotationType = (1, 0, 0, 0),
*,
transform_matrix: Optional[MatrixType] = None,
) -> None:
self._transform = Transform3D(translation, rotation, matrix=transform_matrix)
self._size = Vector3D(*size)
def __eq__(self, other: object) -> bool:
if not isinstance(other, self.__class__):
return False
return self._size.__eq__(other._size) and self._transform.__eq__(other._transform)
def __rmul__(self: _B3, other: Transform3D) -> _B3:
if isinstance(other, (Transform3D, quaternion)):
box: _B3 = object.__new__(self.__class__)
box._transform = other * self._transform
box._size = self._size
return box
return NotImplemented # type: ignore[unreachable]
@staticmethod
def _line_intersect(length1: float, length2: float, midpoint_distance: float) -> float:
"""Calculate the intersect length between two parallel lines.
Arguments:
length1: The length of line1.
length2: the length of line2.
midpoint_distance: The distance between midpoints of the two lines.
Returns:
The intersect length between line1 and line2.
Examples:
>>> Box3D._line_intersect(4, 4, 1)
3.0
"""
line1_min = -length1 / 2
line1_max = length1 / 2
line2_min = -length2 / 2 + midpoint_distance
line2_max = length2 / 2 + midpoint_distance
intersect_length = min(line1_max, line2_max) - max(line1_min, line2_min)
return intersect_length if intersect_length > 0 else 0
def _loads(self, contents: Mapping[str, Mapping[str, float]]) -> None:
self._size = Vector3D.loads(contents["size"])
self._transform = Transform3D.loads(contents)
[docs] @classmethod
def loads(cls: Type[_B3], contents: Mapping[str, Mapping[str, float]]) -> _B3:
"""Load a :class:`Box3D` from a dict containing the coordinates of the 3D box.
Arguments:
contents: A dict containing the coordinates of a 3D box.
Returns:
The loaded :class:`Box3D` object.
Examples:
>>> contents = {
... "size": {"x": 1.0, "y": 2.0, "z": 3.0},
... "translation": {"x": 1.0, "y": 2.0, "z": 3.0},
... "rotation": {"w": 0.0, "x": 1.0, "y": 0.0, "z": 0.0},
... }
>>> Box3D.loads(contents)
Box3D(
(size): Vector3D(1.0, 2.0, 3.0)
(translation): Vector3D(1.0, 2.0, 3.0),
(rotation): quaternion(0, 1, 0, 0),
)
"""
return common_loads(cls, contents)
[docs] @classmethod
def iou(cls, box1: "Box3D", box2: "Box3D", angle_threshold: float = 5) -> float:
"""Calculate the intersection over union between two 3D boxes.
Arguments:
box1: A 3D box.
box2: A 3D box.
angle_threshold: The threshold of the relative angles
between two input 3d boxes in degree.
Returns:
The intersection over union of the two 3D boxes.
Examples:
>>> box3d_1 = Box3D(size=[1, 1, 1])
>>> box3d_2 = Box3D(size=[2, 2, 2])
>>> Box3D.iou(box3d_1, box3d_2)
0.125
"""
box2 = box1.transform.inverse() * box2
if abs(math.degrees(box2.rotation.angle())) > angle_threshold:
return 0
intersect_size = [
cls._line_intersect(*args) for args in zip(box1.size, box2.size, box2.translation)
]
intersect = intersect_size[0] * intersect_size[1] * intersect_size[2]
union = box1.volume() + box2.volume() - intersect
return intersect / union
@property
def translation(self) -> Vector3D:
"""Return the translation of the 3D box.
Returns:
The translation of the 3D box.
Examples:
>>> box3d = Box3D(size=(1, 1, 1), translation=(1, 2, 3))
>>> box3d.translation
Vector3D(1, 2, 3)
"""
return self._transform.translation
@property
def rotation(self) -> quaternion:
"""Return the rotation of the 3D box.
Returns:
The rotation of the 3D box.
Examples:
>>> box3d = Box3D(size=(1, 1, 1), rotation=(0, 1, 0, 0))
>>> box3d.rotation
quaternion(0, 1, 0, 0)
"""
return self._transform.rotation
@property
def transform(self) -> Transform3D:
"""Return the transform of the 3D box.
Returns:
The transform of the 3D box.
Examples:
>>> box3d = Box3D(size=(1, 1, 1), translation=(1, 2, 3), rotation=(1, 0, 0, 0))
>>> box3d.transform
Transform3D(
(translation): Vector3D(1, 2, 3),
(rotation): quaternion(1, 0, 0, 0)
)
"""
return self._transform
@property
def size(self) -> Vector3D:
"""Return the size of the 3D box.
Returns:
The size of the 3D box.
Examples:
>>> box3d = Box3D(size=(1, 1, 1))
>>> box3d.size
Vector3D(1, 1, 1)
"""
return self._size
[docs] def volume(self) -> float:
"""Return the volume of the 3D box.
Returns:
The volume of the 3D box.
Examples:
>>> box3d = Box3D(size=(1, 2, 3))
>>> box3d.volume()
6
"""
return self.size.x * self.size.y * self.size.z
[docs] def dumps(self) -> Dict[str, Dict[str, float]]:
"""Dumps the 3D box into a dict.
Returns:
A dict containing translation, rotation and size information.
Examples:
>>> box3d = Box3D(size=(1, 2, 3), translation=(1, 2, 3), rotation=(0, 1, 0, 0))
>>> box3d.dumps()
{
"translation": {"x": 1, "y": 2, "z": 3},
"rotation": {"w": 0.0, "x": 1.0, "y": 0.0, "z": 0.0},
"size": {"x": 1, "y": 2, "z": 3},
}
"""
contents = self._transform.dumps()
contents["size"] = self.size.dumps()
return contents