337 lines
11 KiB
Python
337 lines
11 KiB
Python
"""General 3D Bounding Box class.
|
|
|
|
Origin: https://github.com/OasisYang/Wild6D/blob/main/lib/box.py
|
|
"""
|
|
|
|
import numpy as np
|
|
from numpy.linalg import lstsq as optimizer
|
|
from scipy.spatial.transform import Rotation as rotation_util
|
|
|
|
EDGES = (
|
|
[1, 5],
|
|
[2, 6],
|
|
[3, 7],
|
|
[4, 8], # lines along x-axis
|
|
[1, 3],
|
|
[5, 7],
|
|
[2, 4],
|
|
[6, 8], # lines along y-axis
|
|
[1, 2],
|
|
[3, 4],
|
|
[5, 6],
|
|
[7, 8], # lines along z-axis
|
|
)
|
|
|
|
# The vertices are ordered according to the left-hand rule, so the normal
|
|
# vector of each face will point inward the box.
|
|
FACES = np.array(
|
|
[
|
|
[5, 6, 8, 7], # +x on yz plane
|
|
[1, 3, 4, 2], # -x on yz plane
|
|
[3, 7, 8, 4], # +y on xz plane = top
|
|
[1, 2, 6, 5], # -y on xz plane
|
|
[2, 4, 8, 6], # +z on xy plane = front
|
|
[1, 5, 7, 3], # -z on xy plane
|
|
]
|
|
)
|
|
|
|
UNIT_BOX = np.asarray(
|
|
[
|
|
[0.0, 0.0, 0.0],
|
|
[-0.5, -0.5, -0.5],
|
|
[-0.5, -0.5, 0.5],
|
|
[-0.5, 0.5, -0.5],
|
|
[-0.5, 0.5, 0.5],
|
|
[0.5, -0.5, -0.5],
|
|
[0.5, -0.5, 0.5],
|
|
[0.5, 0.5, -0.5],
|
|
[0.5, 0.5, 0.5],
|
|
]
|
|
)
|
|
|
|
NUM_KEYPOINTS = 9
|
|
FRONT_FACE_ID = 4
|
|
TOP_FACE_ID = 2
|
|
|
|
|
|
def get_bbox_center_and_corners(bbox):
|
|
x_min, y_min, z_min, x_max, y_max, z_max = (
|
|
bbox.min[0],
|
|
bbox.min[1],
|
|
bbox.min[2],
|
|
bbox.max[0],
|
|
bbox.max[1],
|
|
bbox.max[2],
|
|
)
|
|
|
|
center = np.array([(x_min + x_max) / 2, (y_min + y_max) / 2, (z_min + z_max) / 2])
|
|
|
|
corners = np.array(
|
|
[
|
|
[x_min, y_min, z_min], # 0: minimum corner
|
|
[x_min, y_min, z_max], # 1: z changes
|
|
[x_min, y_max, z_min], # 2: y changes
|
|
[x_min, y_max, z_max], # 3: y and z change
|
|
[x_max, y_min, z_min], # 4: x changes
|
|
[x_max, y_min, z_max], # 5: x and z change
|
|
[x_max, y_max, z_min], # 6: x and y change
|
|
[x_max, y_max, z_max], # 7: maximum corner
|
|
]
|
|
)
|
|
|
|
result = np.vstack([center, corners])
|
|
return result
|
|
|
|
|
|
class Box(object):
|
|
"""General 3D Oriented Bounding Box."""
|
|
|
|
def __init__(self, vertices=None):
|
|
if vertices is None:
|
|
vertices = self.scaled_axis_aligned_vertices(np.array([1.0, 1.0, 1.0]))
|
|
|
|
self._vertices = vertices
|
|
self._rotation = None
|
|
self._translation = None
|
|
self._scale = None
|
|
self._transformation = None
|
|
self._volume = None
|
|
|
|
@classmethod
|
|
def from_transformation(cls, rotation, translation, scale):
|
|
"""Constructs an oriented bounding box from transformation and scale."""
|
|
if rotation.size not in (3, 9):
|
|
raise ValueError("Unsupported rotation, only 3x1 euler angles or 3x3 rotation matrices are supported.")
|
|
if rotation.size == 3:
|
|
rotation = rotation_util.from_rotvec(rotation.tolist()).as_dcm()
|
|
scaled_identity_box = cls.scaled_axis_aligned_vertices(scale)
|
|
vertices = np.zeros((NUM_KEYPOINTS, 3))
|
|
for i in range(NUM_KEYPOINTS):
|
|
vertices[i, :] = np.matmul(rotation, scaled_identity_box[i, :]) + translation.flatten()
|
|
return cls(vertices=vertices)
|
|
|
|
def __repr__(self):
|
|
representation = "Box: "
|
|
for i in range(NUM_KEYPOINTS):
|
|
representation += f"[{i}: {self.vertices[i, 0]}, {self.vertices[i, 1]}, {self.vertices[i, 2]}]"
|
|
return representation
|
|
|
|
def __len__(self):
|
|
return NUM_KEYPOINTS
|
|
|
|
def __name__(self):
|
|
return "Box"
|
|
|
|
def apply_transformation(self, transformation):
|
|
"""Applies transformation on the box.
|
|
|
|
Group multiplication is the same as rotation concatenation. Therefore return
|
|
new box with SE3(R * R2, T + R * T2); Where R2 and T2 are existing rotation
|
|
and translation. Note we do not change the scale.
|
|
|
|
Args:
|
|
transformation: a 4x4 transformation matrix.
|
|
|
|
Returns:
|
|
transformed box.
|
|
"""
|
|
if transformation.shape != (4, 4):
|
|
raise ValueError("Transformation should be a 4x4 matrix.")
|
|
|
|
new_rotation = np.matmul(transformation[:3, :3], self.rotation)
|
|
new_translation = transformation[:3, 3] + (np.matmul(transformation[:3, :3], self.translation))
|
|
return Box.from_transformation(new_rotation, new_translation, self.scale)
|
|
|
|
@classmethod
|
|
def scaled_axis_aligned_vertices(cls, scale):
|
|
"""Returns an axis-aligned set of verticies for a box of the given scale.
|
|
|
|
Args:
|
|
scale: A 3*1 vector, specifiying the size of the box in x-y-z dimension.
|
|
"""
|
|
w = scale[0] / 2.0
|
|
h = scale[1] / 2.0
|
|
d = scale[2] / 2.0
|
|
|
|
# Define the local coordinate system, w.r.t. the center of the box
|
|
aabb = np.array(
|
|
[
|
|
[0.0, 0.0, 0.0],
|
|
[-w, -h, -d],
|
|
[-w, -h, +d],
|
|
[-w, +h, -d],
|
|
[-w, +h, +d],
|
|
[+w, -h, -d],
|
|
[+w, -h, +d],
|
|
[+w, +h, -d],
|
|
[+w, +h, +d],
|
|
]
|
|
)
|
|
return aabb
|
|
|
|
@classmethod
|
|
def fit(cls, vertices):
|
|
"""Estimates a box 9-dof parameters from the given vertices.
|
|
|
|
Directly computes the scale of the box, then solves for orientation and
|
|
translation.
|
|
|
|
Args:
|
|
vertices: A 9*3 array of points. Points are arranged as 1 + 8 (center
|
|
keypoint + 8 box vertices) matrix.
|
|
|
|
Returns:
|
|
orientation: 3*3 rotation matrix.
|
|
translation: 3*1 translation vector.
|
|
scale: 3*1 scale vector.
|
|
"""
|
|
orientation = np.identity(3)
|
|
translation = np.zeros((3, 1))
|
|
scale = np.zeros(3)
|
|
|
|
# The scale would remain invariant under rotation and translation.
|
|
# We can safely estimate the scale from the oriented box.
|
|
for axis in range(3):
|
|
for edge_id in range(4):
|
|
# The edges are stored in quadruples according to each axis
|
|
begin, end = EDGES[axis * 4 + edge_id]
|
|
scale[axis] += np.linalg.norm(vertices[begin, :] - vertices[end, :])
|
|
scale[axis] /= 4.0
|
|
|
|
x = cls.scaled_axis_aligned_vertices(scale)
|
|
system = np.concatenate((x, np.ones((NUM_KEYPOINTS, 1))), axis=1)
|
|
solution, _, _, _ = optimizer(system, vertices, rcond=None)
|
|
orientation = solution[:3, :3].T
|
|
translation = solution[3, :3]
|
|
return orientation, translation, scale
|
|
|
|
def inside(self, point):
|
|
"""Tests whether a given point is inside the box.
|
|
|
|
Brings the 3D point into the local coordinate of the box. In the local
|
|
coordinate, the looks like an axis-aligned bounding box. Next checks if
|
|
the box contains the point.
|
|
Args:
|
|
point: A 3*1 numpy vector.
|
|
|
|
Returns:
|
|
True if the point is inside the box, False otherwise.
|
|
"""
|
|
inv_trans = np.linalg.inv(self.transformation)
|
|
scale = self.scale
|
|
point_w = np.matmul(inv_trans[:3, :3], point) + inv_trans[:3, 3]
|
|
for i in range(3):
|
|
if abs(point_w[i]) > scale[i] / 2.0:
|
|
return False
|
|
return True
|
|
|
|
def sample(self):
|
|
"""Samples a 3D point uniformly inside this box."""
|
|
point = np.random.uniform(-0.5, 0.5, 3) * self.scale
|
|
point = np.matmul(self.rotation, point) + self.translation
|
|
return point
|
|
|
|
@property
|
|
def vertices(self):
|
|
return self._vertices
|
|
|
|
@property
|
|
def rotation(self):
|
|
if self._rotation is None:
|
|
self._rotation, self._translation, self._scale = self.fit(self._vertices)
|
|
return self._rotation
|
|
|
|
@property
|
|
def translation(self):
|
|
if self._translation is None:
|
|
self._rotation, self._translation, self._scale = self.fit(self._vertices)
|
|
return self._translation
|
|
|
|
@property
|
|
def scale(self):
|
|
if self._scale is None:
|
|
self._rotation, self._translation, self._scale = self.fit(self._vertices)
|
|
return self._scale
|
|
|
|
@property
|
|
def volume(self):
|
|
"""Compute the volume of the parallelpiped or the box.
|
|
|
|
For the boxes, this is equivalent to np.prod(self.scale). However for
|
|
parallelpiped, this is more involved. Viewing the box as a linear function
|
|
we can estimate the volume using a determinant. This is equivalent to
|
|
sp.ConvexHull(self._vertices).volume
|
|
|
|
Returns:
|
|
volume (float)
|
|
"""
|
|
if self._volume is None:
|
|
i = self._vertices[2, :] - self._vertices[1, :]
|
|
j = self._vertices[3, :] - self._vertices[1, :]
|
|
k = self._vertices[5, :] - self._vertices[1, :]
|
|
sys = np.array([i, j, k])
|
|
self._volume = abs(np.linalg.det(sys))
|
|
return self._volume
|
|
|
|
@property
|
|
def transformation(self):
|
|
if self._rotation is None:
|
|
self._rotation, self._translation, self._scale = self.fit(self._vertices)
|
|
if self._transformation is None:
|
|
self._transformation = np.identity(4)
|
|
self._transformation[:3, :3] = self._rotation
|
|
self._transformation[:3, 3] = self._translation
|
|
return self._transformation
|
|
|
|
def get_ground_plane(self, gravity_axis=1):
|
|
"""Get ground plane under the box."""
|
|
|
|
gravity = np.zeros(3)
|
|
gravity[gravity_axis] = 1
|
|
|
|
def get_face_normal(face, center):
|
|
"""Get a normal vector to the given face of the box."""
|
|
v1 = self.vertices[face[0], :] - center
|
|
v2 = self.vertices[face[1], :] - center
|
|
normal = np.cross(v1, v2)
|
|
return normal
|
|
|
|
def get_face_center(face):
|
|
"""Get the center point of the face of the box."""
|
|
center = np.zeros(3)
|
|
for vertex in face:
|
|
center += self.vertices[vertex, :]
|
|
center /= len(face)
|
|
return center
|
|
|
|
ground_plane_id = 0
|
|
ground_plane_error = 10.0
|
|
|
|
# The ground plane is defined as a plane aligned with gravity.
|
|
# gravity is the (0, 1, 0) vector in the world coordinate system.
|
|
for i in [0, 2, 4]:
|
|
face = FACES[i, :]
|
|
center = get_face_center(face)
|
|
normal = get_face_normal(face, center)
|
|
w = np.cross(gravity, normal)
|
|
w_sq_norm = np.linalg.norm(w)
|
|
if w_sq_norm < ground_plane_error:
|
|
ground_plane_error = w_sq_norm
|
|
ground_plane_id = i
|
|
|
|
face = FACES[ground_plane_id, :]
|
|
center = get_face_center(face)
|
|
normal = get_face_normal(face, center)
|
|
|
|
# For each face, we also have a parallel face that it's normal is also
|
|
# aligned with gravity vector. We pick the face with lower height (y-value).
|
|
# The parallel to face 0 is 1, face 2 is 3, and face 4 is 5.
|
|
parallel_face_id = ground_plane_id + 1
|
|
parallel_face = FACES[parallel_face_id]
|
|
parallel_face_center = get_face_center(parallel_face)
|
|
parallel_face_normal = get_face_normal(parallel_face, parallel_face_center)
|
|
if parallel_face_center[gravity_axis] < center[gravity_axis]:
|
|
center = parallel_face_center
|
|
normal = parallel_face_normal
|
|
return center, normal
|