diff --git a/main.py b/main.py index 55a775d..e1fe8cd 100644 --- a/main.py +++ b/main.py @@ -26,6 +26,14 @@ ] POINT_NAMES: List[str] = [] +# Configure vectors to draw (list of (start_index, end_index) using 0-based indices) +# Edit this list to add/remove vectors to display +VECTOR_PAIRS = [(4, 5), (12, 13)] +# Optional colors (will cycle if fewer colors than VECTOR_PAIRS) +VECTOR_COLORS = ["r", "b"] +# Line width for vector arrows +VECTOR_LW = 4 + # 3D arrow helper to draw arrows between 3D points class Arrow3D(FancyArrowPatch): @@ -134,17 +142,21 @@ 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) + # Draw vector arrows as specified in VECTOR_PAIRS + arrows = [] + for i_pair, (idx1, idx2) in enumerate(VECTOR_PAIRS): + color = VECTOR_COLORS[i_pair % len(VECTOR_COLORS)] if VECTOR_COLORS else "r" + a = Arrow3D( + [points[idx1][0], points[idx2][0]], + [points[idx1][1], points[idx2][1]], + [points[idx1][2], points[idx2][2]], + mutation_scale=20, + lw=VECTOR_LW, + arrowstyle="-|>", + color=color, + ) + ax.add_artist(a) + arrows.append(a) ax.set_xlabel("X") ax.set_ylabel("Y") @@ -179,12 +191,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]], - ) + # Update all configured arrows + for i_pair, (idx1, idx2) in enumerate(VECTOR_PAIRS): + arrows[i_pair]._verts3d = ( + [points[idx1][0], points[idx2][0]], + [points[idx1][1], points[idx2][1]], + [points[idx1][2], points[idx2][2]], + ) ax.set_title(f"Frame {data[frame_idx].frame_number}") fig.canvas.draw_idle()