import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from pathlib import Path
import re
def parse_trc(trc_path):
"""Parse a .trc file and return marker names, frame data, and metadata."""
with open(trc_path, 'r') as f:
lines = f.readlines()
meta_keys = lines[2].strip().split('\t')
meta_vals = lines[3].strip().split('\t')
metadata = dict(zip(meta_keys, meta_vals))
marker_line = lines[3].strip().split('\t')
header_lines = 0
for i, line in enumerate(lines):
if line.strip() and not line.startswith(('PathFileType', 'DataRate',
'Frame', '\t')):
try:
float(line.strip().split('\t')[0])
header_lines = i
break
except ValueError:
continue
raw_markers = lines[3].strip().split('\t')
markers = [m for m in raw_markers if m and m not in ('Frame#', 'Time')]
marker_names = []
for m in markers:
if m and (not marker_names or m != marker_names[-1]):
marker_names.append(m)
data_lines = lines[header_lines:]
data = []
for line in data_lines:
vals = line.strip().split('\t')
if len(vals) > 2:
try:
row = [float(v) if v else np.nan for v in vals]
data.append(row)
except ValueError:
continue
data = np.array(data)
return marker_names, data, metadata
trc_files = sorted((work_dir / "pose-3d").glob("*.trc"))
print(f"š Found {len(trc_files)} TRC file(s):")
for f in trc_files:
print(f" {f.name}")
trc_file = None
for f in trc_files:
if 'filt' in f.name.lower() and 'augm' not in f.name.lower():
trc_file = f
break
if trc_file is None and trc_files:
trc_file = trc_files[0]
if trc_file:
print(f"\nš Visualizing: {trc_file.name}")
marker_names, data, metadata = parse_trc(trc_file)
print(f" Markers: {len(marker_names)}")
print(f" Frames: {data.shape[0]}")
print(f" Marker names: {marker_names[:10]}{'...' if len(marker_names) > 10 else ''}")
frames = data[:, 0].astype(int) if data.shape[1] > 0 else []
times = data[:, 1] if data.shape[1] > 1 else []
coords = data[:, 2:]
n_markers = len(marker_names)
mid_frame = len(data) // 2
fig = plt.figure(figsize=(16, 6))
ax1 = fig.add_subplot(131, projection='3d')
xs = coords[mid_frame, 0::3][:n_markers]
ys = coords[mid_frame, 1::3][:n_markers]
zs = coords[mid_frame, 2::3][:n_markers]
ax1.scatter(xs, ys, zs, c="dodgerblue", s=40, alpha=0.8, edgecolors="navy")
for i, name in enumerate(marker_names[:len(xs)]):
if i % 3 == 0:
ax1.text(xs[i], ys[i], zs[i], f' {name}', fontsize=6, alpha=0.7)
ax1.set_xlabel('X (m)')
ax1.set_ylabel('Y (m)')
ax1.set_zlabel('Z (m)')
ax1.set_title(f'3D Keypoints (Frame {int(frames[mid_frame])})', fontsize=10)
ax2 = fig.add_subplot(132)
key_markers = ['RAnkle', 'LAnkle', 'RWrist', 'LWrist', 'Nose']
colors_map = {'RAnkle': 'red', 'LAnkle': 'blue', 'RWrist': 'orange',
'LWrist': 'green', 'Nose': 'purple'}
for mkr in key_markers:
if mkr in marker_names:
idx = marker_names.index(mkr)
z_col = idx * 3 + 2
if z_col < coords.shape[1]:
ax2.plot(times, coords[:, z_col],
label=mkr, color=colors_map.get(mkr, 'gray'),
linewidth=1.2, alpha=0.8)
ax2.set_xlabel('Time (s)')
ax2.set_ylabel('Z position (m)')
ax2.set_title('Vertical Trajectories', fontsize=10)
ax2.legend(fontsize=8, loc="best")
ax2.grid(True, alpha=0.3)
ax3 = fig.add_subplot(133)
if len(times) > 1:
dt = np.diff(times)
dt[dt == 0] = 1e-6
for mkr in ['RAnkle', 'RWrist']:
if mkr in marker_names:
idx = marker_names.index(mkr)
x_col, y_col, z_col = idx * 3, idx * 3 + 1, idx * 3 + 2
if z_col < coords.shape[1]:
dx = np.diff(coords[:, x_col])
dy = np.diff(coords[:, y_col])
dz = np.diff(coords[:, z_col])
speed = np.sqrt(dx**2 + dy**2 + dz**2) / dt
speed = np.clip(speed, 0, np.nanpercentile(speed, 99))
ax3.plot(times[1:], speed, label=mkr,
color=colors_map.get(mkr, 'gray'),
linewidth=0.8, alpha=0.7)
ax3.set_xlabel('Time (s)')
ax3.set_ylabel('Speed (m/s)')
ax3.set_title('Marker Speeds (Quality Check)', fontsize=10)
ax3.legend(fontsize=8)
ax3.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(work_dir / 'trajectory_analysis.png', dpi=150, bbox_inches="tight")
plt.show()
print("ā
Trajectory plots saved to trajectory_analysis.png")
else:
print("ā No TRC file found to visualize.")