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 |