Spaces:
Sleeping
Sleeping
import numpy as np | |
import cv2 | |
import plotly.graph_objects as go | |
from math import cos, sin | |
from scipy.spatial.transform import Rotation as R | |
def plot_3D_rotation(rotation_matrix): | |
fig = go.Figure() | |
# Original axis orientation | |
axes_points = np.array([ | |
[1, 0, 0, 0], | |
[0, 1, 0, 0], | |
[0, 0, 1, 0] | |
], dtype=np.float64) | |
# Plot original axes | |
fig.add_trace(go.Scatter3d( | |
x=[0, axes_points[0, 0]], | |
y=[0, axes_points[1, 0]], | |
z=[0, axes_points[2, 0]], | |
mode='lines+text', | |
line=dict(color='blue', width=6), | |
name='Canonical X-axis', | |
text=['', 'X axis'], | |
textposition='middle center', | |
)) | |
fig.add_trace(go.Scatter3d( | |
x=[0, axes_points[0, 1]], | |
y=[0, axes_points[1, 1]], | |
z=[0, axes_points[2, 1]], | |
mode='lines+text', | |
line=dict(color='blue', width=6), | |
name='Canonical Z-axis', | |
text=['', 'Z axis'], | |
textposition='middle center', | |
)) | |
fig.add_trace(go.Scatter3d( | |
x=[0, axes_points[0, 2]], | |
y=[0, axes_points[1, 2]], | |
z=[0, axes_points[2, 2]], | |
mode='lines+text', | |
line=dict(color='blue', width=6), | |
name='Canonical Y-axis', | |
text=['', 'Y axis'], | |
textposition='middle center', | |
)) | |
# Apply rotation | |
axes_points = rotation_matrix @ axes_points | |
# Plot rotated axes | |
fig.add_trace(go.Scatter3d( | |
x=[0, axes_points[0, 0]], | |
y=[0, axes_points[1, 0]], | |
z=[0, axes_points[2, 0]], | |
mode='lines+text', | |
line=dict(color='red', width=6), | |
name='Rotated X\'-axis', | |
text=['', 'Rotated X axis'], | |
textposition='middle center', | |
)) | |
fig.add_trace(go.Scatter3d( | |
x=[0, axes_points[0, 1]], | |
y=[0, axes_points[1, 1]], | |
z=[0, axes_points[2, 1]], | |
mode='lines+text', | |
line=dict(color='red', width=6), | |
name='Rotated Z\'-axis', | |
text=['', 'Rotated Z axis'], | |
textposition='middle center', | |
)) | |
fig.add_trace(go.Scatter3d( | |
x=[0, axes_points[0, 2]], | |
y=[0, axes_points[1, 2]], | |
z=[0, axes_points[2, 2]], | |
mode='lines+text', | |
line=dict(color='red', width=6), | |
name='Rotated Y\'-axis', | |
text=['', 'Rotated Y axis'], | |
textposition='middle center', | |
)) | |
# Retrieve pitch, yaw, roll from rotation matrix | |
r = R.from_matrix(rotation_matrix) | |
pitch, yaw, roll = r.as_euler('xzy', degrees=True) | |
# Set layout | |
fig.update_layout( | |
scene=dict( | |
xaxis=dict(title='X-axis', range=[-1.2, 1.2]), | |
yaxis=dict(title='Z-axis', range=[-1.2, 1.2]), | |
zaxis=dict(title='Y-axis', range=[-1.2, 1.2]), | |
xaxis_tickvals=np.arange(-1.2, 1.2, 0.6), | |
yaxis_tickvals=np.arange(-1.2, 1.2, 0.5), | |
zaxis_tickvals=np.arange(-1.2, 1.2, 0.5), | |
aspectmode='cube', | |
aspectratio=dict(x=1, y=1, z=1), | |
), | |
margin=dict(l=0, r=0, t=0, b=30), | |
) | |
# add annotation | |
fig.add_annotation(dict(font=dict(color='black',size=15), | |
x=-30, | |
y=50, | |
showarrow=False, | |
text=f"Pitch: {int(pitch)} - Yaw: {int(yaw)} - Roll: {int(roll)}", | |
textangle=0, | |
xanchor='left', | |
xref="paper", | |
yref="paper")) | |
return fig | |
def draw_2D_axes(img, roll, pitch, yaw, tdx=None, tdy=None, size=150.): | |
# Input is a cv2 image | |
# pose_params: (pitch, yaw, roll, tdx, tdy) | |
# Where (tdx, tdy) is the translation of the face. | |
# For pose we have [pitch yaw roll tdx tdy tdz scale_factor] | |
p = pitch * np.pi / 180 | |
y = (yaw * np.pi / 180) | |
r = -roll * np.pi / 180 | |
if tdx != None and tdy != None: | |
face_x = tdx - 0.50 * size | |
face_y = tdy - 0.50 * size | |
else: | |
height, width = img.shape[:2] | |
face_x = width / 2 - 0.5 * size | |
face_y = height / 2 - 0.5 * size | |
x1 = size * (cos(y) * cos(r)) + face_x | |
y1 = size * (cos(p) * sin(r) + cos(r) * sin(p) * sin(y)) + face_y | |
x2 = size * (-cos(y) * sin(r)) + face_x | |
y2 = size * (cos(p) * cos(r) - sin(p) * sin(y) * sin(r)) + face_y | |
x3 = size * (sin(y)) + face_x | |
y3 = size * (-cos(y) * sin(p)) + face_y | |
# Draw base in red | |
cv2.line(img, (int(face_x), int(face_y)), (int(x1),int(y1)),(0,0,255),3) | |
cv2.line(img, (int(face_x), int(face_y)), (int(x2),int(y2)),(0,0,255),3) | |
cv2.line(img, (int(x2), int(y2)), (int(x2+x1-face_x),int(y2+y1-face_y)),(0,0,255),3) | |
cv2.line(img, (int(x1), int(y1)), (int(x1+x2-face_x),int(y1+y2-face_y)),(0,0,255),3) | |
# Draw pillars in blue | |
cv2.line(img, (int(face_x), int(face_y)), (int(x3),int(y3)),(255,0,0),2) | |
cv2.line(img, (int(x1), int(y1)), (int(x1+x3-face_x),int(y1+y3-face_y)),(255,0,0),2) | |
cv2.line(img, (int(x2), int(y2)), (int(x2+x3-face_x),int(y2+y3-face_y)),(255,0,0),2) | |
cv2.line(img, (int(x2+x1-face_x),int(y2+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(255,0,0),2) | |
# Draw top in green | |
cv2.line(img, (int(x3+x1-face_x),int(y3+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2) | |
cv2.line(img, (int(x2+x3-face_x),int(y2+y3-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2) | |
cv2.line(img, (int(x3), int(y3)), (int(x3+x1-face_x),int(y3+y1-face_y)),(0,255,0),2) | |
cv2.line(img, (int(x3), int(y3)), (int(x3+x2-face_x),int(y3+y2-face_y)),(0,255,0),2) | |
return img |