runtime/rbs_perception/scripts/cuboid.py

# Copyright (c) 2018 NVIDIA Corporation. All rights reserved.
# This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.
# https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode

from enum import IntEnum, unique
import numpy as np
import cv2

# Related to the object's local coordinate system
# @unique
class CuboidVertexType(IntEnum):
    FrontTopRight = 0
    FrontTopLeft = 1
    FrontBottomLeft = 2
    FrontBottomRight = 3
    RearTopRight = 4
    RearTopLeft = 5
    RearBottomLeft = 6
    RearBottomRight = 7
    Center = 8
    TotalCornerVertexCount = 8  # Corner vertexes doesn't include the center point
    TotalVertexCount = 9


# List of the vertex indexes in each line edges of the cuboid
CuboidLineIndexes = [
    # Front face
    [CuboidVertexType.FrontTopLeft, CuboidVertexType.FrontTopRight],
    [CuboidVertexType.FrontTopRight, CuboidVertexType.FrontBottomRight],
    [CuboidVertexType.FrontBottomRight, CuboidVertexType.FrontBottomLeft],
    [CuboidVertexType.FrontBottomLeft, CuboidVertexType.FrontTopLeft],
    # Back face
    [CuboidVertexType.RearTopLeft, CuboidVertexType.RearTopRight],
    [CuboidVertexType.RearTopRight, CuboidVertexType.RearBottomRight],
    [CuboidVertexType.RearBottomRight, CuboidVertexType.RearBottomLeft],
    [CuboidVertexType.RearBottomLeft, CuboidVertexType.RearTopLeft],
    # Left face
    [CuboidVertexType.FrontBottomLeft, CuboidVertexType.RearBottomLeft],
    [CuboidVertexType.FrontTopLeft, CuboidVertexType.RearTopLeft],
    # Right face
    [CuboidVertexType.FrontBottomRight, CuboidVertexType.RearBottomRight],
    [CuboidVertexType.FrontTopRight, CuboidVertexType.RearTopRight],
]


# ========================= Cuboid3d =========================
class Cuboid3d:
    """This class contains a 3D cuboid."""

    # Create a box with a certain size
    def __init__(
        self,
        size3d=[1.0, 1.0, 1.0],
        center_location=[0, 0, 0],
        coord_system=None,
        parent_object=None,
    ):

        # NOTE: This local coordinate system is similar
        # to the intrinsic transform matrix of a 3d object
        self.center_location = center_location
        self.coord_system = coord_system
        self.size3d = size3d
        self._vertices = [0, 0, 0] * CuboidVertexType.TotalVertexCount

        self.generate_vertexes()

    def get_vertex(self, vertex_type):
        """Returns the location of a vertex.

        Args:
            vertex_type: enum of type CuboidVertexType

        Returns:
            Numpy array(3) - Location of the vertex type in the cuboid
        """
        return self._vertices[vertex_type]

    def get_vertices(self):
        return self._vertices

    def generate_vertexes(self):
        width, height, depth = self.size3d

        # By default just use the normal OpenCV coordinate system
        if self.coord_system is None:
            cx, cy, cz = self.center_location
            # X axis point to the right
            right = cx + width / 2.0
            left = cx - width / 2.0
            # Y axis point downward
            top = cy - height / 2.0
            bottom = cy + height / 2.0
            # Z axis point forward
            front = cz + depth / 2.0
            rear = cz - depth / 2.0

            # List of 8 vertices of the box
            self._vertices = [
                [right, top, front],  # Front Top Right
                [left, top, front],  # Front Top Left
                [left, bottom, front],  # Front Bottom Left
                [right, bottom, front],  # Front Bottom Right
                [right, top, rear],  # Rear Top Right
                [left, top, rear],  # Rear Top Left
                [left, bottom, rear],  # Rear Bottom Left
                [right, bottom, rear],  # Rear Bottom Right
                self.center_location,  # Center
            ]
        else:
            sx, sy, sz = self.size3d
            forward = np.array(self.coord_system.forward, dtype=float) * sy * 0.5
            up = np.array(self.coord_system.up, dtype=float) * sz * 0.5
            right = np.array(self.coord_system.right, dtype=float) * sx * 0.5
            center = np.array(self.center_location, dtype=float)
            self._vertices = [
                center + forward + up + right,  # Front Top Right
                center + forward + up - right,  # Front Top Left
                center + forward - up - right,  # Front Bottom Left
                center + forward - up + right,  # Front Bottom Right
                center - forward + up + right,  # Rear Top Right
                center - forward + up - right,  # Rear Top Left
                center - forward - up - right,  # Rear Bottom Left
                center - forward - up + right,  # Rear Bottom Right
                self.center_location,  # Center
            ]