mmdet3d_plugin/datasets/utils/data_structure.py

from typing import Type

import numpy as np

import utils.transform as box_transform
import utils.rotation as object_rotation


class Box3D:
    def __init__(
        self,
        x=None,
        y=None,
        z=None,
        h=None,
        w=None,
        l=None,
        pitch=None,
        yaw=None,
        roll=None,
        s=None,
    ):
        self.x = x  # center x
        self.y = y  # center y
        self.z = z  # center z
        self.h = h  # height
        self.w = w  # width
        self.l = l  # length
        self.pitch = pitch  # orientation
        self.yaw = yaw  # orientation
        self.roll = roll  # orientation
        self.s = s  # detection score
        self.corners_3d_cam = None

    def __str__(self):
        return "x: {}, y: {}, z: {}, heading: {}, length: {}, width: {}, height: {}, score: {}".format(
            self.x, self.y, self.z, self.yaw, self.l, self.w, self.h, self.s
        )

    @classmethod
    def bbox2dict(cls, bbox):
        return {
            "center_x": bbox.x,
            "center_y": bbox.y,
            "center_z": bbox.z,
            "height": bbox.h,
            "width": bbox.w,
            "length": bbox.l,
            "yaw": bbox.yaw,
        }

    @classmethod
    def bbox2array(cls, bbox):
        if bbox.s is None:
            return np.array([bbox.x, bbox.y, bbox.z, bbox.yaw, bbox.l, bbox.w, bbox.h])
        else:
            return np.array(
                [bbox.x, bbox.y, bbox.z, bbox.yaw, bbox.l, bbox.w, bbox.h, bbox.s]
            )

    @classmethod
    def bbox2array_raw(cls, bbox):
        if bbox.s is None:
            return np.array([bbox.h, bbox.w, bbox.l, bbox.x, bbox.y, bbox.z, bbox.yaw])
        else:
            return np.array(
                [bbox.h, bbox.w, bbox.l, bbox.x, bbox.y, bbox.z, bbox.yaw, bbox.s]
            )

    @classmethod
    def array2bbox_xyzlwhyaw(cls, data):
        # take the format of data of [x,y,z,l,w,h,yaw]

        bbox = Box3D()
        bbox.x, bbox.y, bbox.z, bbox.l, bbox.w, bbox.h, bbox.yaw = data[:7]
        if len(data) == 8:
            bbox.s = data[-1]
        return bbox

    @classmethod
    def array2bbox_raw(cls, data):
        # take the format of data of [h,w,l,x,y,z,yaw]

        bbox = Box3D()
        bbox.h, bbox.w, bbox.l, bbox.x, bbox.y, bbox.z, bbox.yaw = data[:7]
        if len(data) == 8:
            bbox.s = data[-1]
        return bbox

    @classmethod
    def array2bbox(cls, data):
        # take the format of data of [x,y,z,yaw,l,w,h]

        bbox = Box3D()
        bbox.x, bbox.y, bbox.z, bbox.yaw, bbox.l, bbox.w, bbox.h = data[:7]
        if len(data) == 8:
            bbox.s = data[-1]
        return bbox

    @classmethod
    def array2bbox_9dof(
        cls, location: np.ndarray, rotation: np.ndarray, size: np.ndarray
    ):
        # take the format of data of [h,w,l,x,y,z,yaw]

        bbox = Box3D()
        bbox.x, bbox.y, bbox.z = location
        bbox.pitch, bbox.ry, bbox.roll = rotation
        bbox.l, bbox.w, bbox.h = size

        # TODO: change to world
        bbox.corners_3d_cam: np.ndarray = box_transform.get_box_in_world(
            location, rotation, size
        )  # (3, )

        return bbox

    @classmethod
    def box2corners3d_lidar(cls, bbox):
        """Convert the box to the 8 corners in the lidar coordinate

        Note that carla lidar coordinate is
        +x -> left
        +y -> front
        +z -> up
            4 -------- 5
           /|         /|
          7 -------- 6 .
          | |        | |
          . 0 -------- 1
          |/         |/
          3 -------- 2
        """

        # canonical box
        size = np.array([bbox.l, bbox.w, bbox.h])  # (3, )
        box3d: np.ndarray = box_transform.create_bb_points(size)  # 8 x 4
        # print(box3d)
        # zxc

        # with rotation
        R: np.ndarray = object_rotation.rotz(bbox.yaw)  # 3 x 3
        # print(R)
        box3d = np.dot(R, box3d[:, :3].transpose())  # 3 x 8

        # with translation
        box3d[0, :] = box3d[0, :] + bbox.x
        box3d[1, :] = box3d[1, :] + bbox.y
        box3d[2, :] = box3d[2, :] + bbox.z

        # homogeneous
        box3d = box3d.transpose()  # 8 x 3
        box3d = np.concatenate((box3d, np.ones((8, 1))), axis=1)  # 8 x 4

        return box3d

    # @classmethod
    # def box2corners3d_camcoord(cls, bbox):
    #     """Takes an object's 3D box with the representation of [x,y,z,theta,l,w,h] and
    #     convert it to the 8 corners of the 3D box, the box is in the camera coordinate
    #     with right x, down y, front z

    #     Returns:
    #         corners_3d: (8,3) array in in rect camera coord

    #     box corner order is like follows
    #             1 -------- 0         top is bottom because y direction is negative
    #            /|         /|
    #           2 -------- 3 .
    #           | |        | |
    #           . 5 -------- 4
    #           |/         |/
    #           6 -------- 7

    #     rect/ref camera coord:
    #     right x, down y, front z

    #     x -> w, z -> l, y -> h
    #     """

    #     # if already computed before, then skip it
    #     if bbox.corners_3d_cam is not None:
    #         return bbox.corners_3d_cam

    #     # compute rotational matrix around yaw axis
    #     # -1.57 means straight, so there is a rotation here
    #     R = roty(bbox.ry)

    #     # 3d bounding box dimensions
    #     l, w, h = bbox.l, bbox.w, bbox.h

    #     # 3d bounding box corners
    #     x_corners = [l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2]
    #     y_corners = [0, 0, 0, 0, -h, -h, -h, -h]
    #     z_corners = [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2]

    #     # rotate and translate 3d bounding box
    #     corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
    #     corners_3d[0, :] = corners_3d[0, :] + bbox.x
    #     corners_3d[1, :] = corners_3d[1, :] + bbox.y
    #     corners_3d[2, :] = corners_3d[2, :] + bbox.z
    #     corners_3d = np.transpose(corners_3d)
    #     bbox.corners_3d_cam = corners_3d

    #     return corners_3d