forked from rendezvue/PyDope
-
Notifications
You must be signed in to change notification settings - Fork 0
/
cuboid.py
147 lines (125 loc) · 5.82 KB
/
cuboid.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
# 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
]
def get_projected_cuboid2d(self, cuboid_transform, camera_intrinsic_matrix):
"""
Projects the cuboid into the image plane using camera intrinsics.
Args:
cuboid_transform: the world transform of the cuboid
camera_intrinsic_matrix: camera intrinsic matrix
Returns:
Cuboid2d - the projected cuboid points
"""
world_transform_matrix = cuboid_transform
rvec = [0, 0, 0]
tvec = [0, 0, 0]
dist_coeffs = np.zeros((4, 1))
transformed_vertices = [0, 0, 0] * CuboidVertexType.TotalVertexCount
for vertex_index in range(CuboidVertexType.TotalVertexCount):
vertex3d = self._vertices[vertex_index]
transformed_vertices[vertex_index] = world_transform_matrix * vertex3d
projected_vertices = cv2.projectPoints(transformed_vertices, rvec, tvec,
camera_intrinsic_matrix, dist_coeffs)
return Cuboid2d(projected_vertices)