diff --git a/main.py b/main.py index 26b1811..55a775d 100644 --- a/main.py +++ b/main.py @@ -2,7 +2,9 @@ import matplotlib.pyplot as plt import pandas as pd +from matplotlib.patches import FancyArrowPatch from matplotlib.widgets import Slider +from mpl_toolkits.mplot3d import proj3d connection = [ (0, 1), @@ -25,6 +27,25 @@ POINT_NAMES: List[str] = [] +# 3D arrow helper to draw arrows between 3D points +class Arrow3D(FancyArrowPatch): + def __init__(self, xs, ys, zs, *args, **kwargs): + super().__init__((0, 0), (0, 0), *args, **kwargs) + self._verts3d = xs, ys, zs + + def draw(self, renderer): + xs3d, ys3d, zs3d = self._verts3d + xs, ys, _ = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.get_proj()) + self.set_positions((xs[0], ys[0]), (xs[1], ys[1])) + super().draw(renderer) + + def do_3d_projection(self, renderer=None): + xs3d, ys3d, zs3d = self._verts3d + xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.get_proj()) + self.set_positions((xs[0], ys[0]), (xs[1], ys[1])) + return max(zs) + + # Data class to hold thumb frame data class ThumbFrameData: def __init__( @@ -113,6 +134,18 @@ text = ax.text(p[0], p[1], p[2], name, fontsize=8) texts.append(text) + # Draw vector arrow from point 4 to 5 (thick red arrow) + arrow = Arrow3D( + [points[4][0], points[5][0]], + [points[4][1], points[5][1]], + [points[4][2], points[5][2]], + mutation_scale=20, + lw=4, + arrowstyle="-|>", + color="r", + ) + ax.add_artist(arrow) + ax.set_xlabel("X") ax.set_ylabel("Y") ax.set_zlabel("Z") @@ -146,6 +179,13 @@ for i, p in enumerate(points): texts[i].set_position((p[0], p[1], p[2])) + # Update arrow from point 4 to 5 + arrow._verts3d = ( + [points[4][0], points[5][0]], + [points[4][1], points[5][1]], + [points[4][2], points[5][2]], + ) + ax.set_title(f"Frame {data[frame_idx].frame_number}") fig.canvas.draw_idle()