diff --git a/.gitignore b/.gitignore index ccc9fd9..13a8af7 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,6 @@ -*.DS_Store \ No newline at end of file +*.DS_Store +__pycache__/ +*.py[cod] +.ipynb_checkpoints/ +.python-version +.vscode/ \ No newline at end of file diff --git a/basketball/freethrow/README.md b/basketball/freethrow/README.md index af3afbf..633ab33 100644 --- a/basketball/freethrow/README.md +++ b/basketball/freethrow/README.md @@ -125,6 +125,37 @@ We provide basic utilities to animate free throw trials. We use the 3D basketbal ``` pip install mplbasketball ``` -The `animate_trial()` is defined in `utils/animate.py`, and produces a GIF like the one below: +The `animate_trial()` function is defined in `animate.py`, and produces a GIF like the one below. + +Recent updates to `animate_trial()`: +- Animation playback now uses each trial's `sampling_rate` (e.g. 30fps and 60fps are handled automatically). +- Joint names are resolved across common schema variants (e.g. `R_HIP` and `RIGHT_HIP`) for better cross-session compatibility. + +Example: +```python +from animate import animate_trial + +anim = animate_trial("./data/2025-12-18/P0001/BB_FT_P0001_T0001.json") +``` + +## Testing + +If `pytest` is not installed in your environment, install it first: + +```bash +python -m pip install pytest +``` + +From the repository root, run: + +```bash +python -m pytest -q tests/test_animate.py +``` + +To run all tests in the `tests/` directory: + +```bash +python -m pytest -q tests +``` diff --git a/basketball/freethrow/animate.py b/basketball/freethrow/animate.py index 6e51667..08ffd15 100644 --- a/basketball/freethrow/animate.py +++ b/basketball/freethrow/animate.py @@ -44,6 +44,44 @@ ("L_1STFINGER", "L_5THFINGER"), ] +JOINT_ALIASES = { + "R_EYE": ["RIGHT_EYE"], + "L_EYE": ["LEFT_EYE"], + "R_EAR": ["RIGHT_EAR"], + "L_EAR": ["LEFT_EAR"], + "R_SHOULDER": ["RIGHT_SHOULDER"], + "L_SHOULDER": ["LEFT_SHOULDER"], + "R_ELBOW": ["RIGHT_ELBOW"], + "L_ELBOW": ["LEFT_ELBOW"], + "R_WRIST": ["RIGHT_WRIST"], + "L_WRIST": ["LEFT_WRIST"], + "R_HIP": ["RIGHT_HIP"], + "L_HIP": ["LEFT_HIP"], + "R_KNEE": ["RIGHT_KNEE"], + "L_KNEE": ["LEFT_KNEE"], + "R_ANKLE": ["RIGHT_ANKLE"], + "L_ANKLE": ["LEFT_ANKLE"], + "R_1STTOE": ["RIGHT_BIG_TOE"], + "L_1STTOE": ["LEFT_BIG_TOE"], + "R_5THTOE": ["RIGHT_SMALL_TOE"], + "L_5THTOE": ["LEFT_SMALL_TOE"], + "R_CALC": ["RIGHT_HEEL"], + "L_CALC": ["LEFT_HEEL"], + "R_1STFINGER": ["RIGHT_FIRST_FINGER_MCP"], + "L_1STFINGER": ["LEFT_FIRST_FINGER_MCP"], + "R_5THFINGER": ["RIGHT_FIFTH_FINGER_MCP"], + "L_5THFINGER": ["LEFT_FIFTH_FINGER_MCP"], +} + + +def _resolve_joint_name(joint_name, available_joints): + if joint_name in available_joints: + return joint_name + for alias in JOINT_ALIASES.get(joint_name, []): + if alias in available_joints: + return alias + return None + def animate_trial( path_to_json, @@ -128,8 +166,28 @@ def animate_trial( # For convenience, we will cast everything to numpy arrays here, but you can keep them as lists if you prefer. for joint in player_joint_dict: - player_joint_dict[joint] = np.array(player_joint_dict[joint]) - ball_data_array = np.array(ball_data_array) + player_joint_dict[joint] = np.array(player_joint_dict[joint], dtype=float) + ball_data_array = np.array(ball_data_array, dtype=float) + available_joints = set(player_joint_dict.keys()) + + resolved_connections = [] + seen_connections = set() + for part1, part2 in connections: + resolved_part1 = _resolve_joint_name(part1, available_joints) + resolved_part2 = _resolve_joint_name(part2, available_joints) + if resolved_part1 is None or resolved_part2 is None: + continue + connection = (resolved_part1, resolved_part2) + if connection in seen_connections: + continue + seen_connections.add(connection) + resolved_connections.append(connection) + + if not resolved_connections: + raise ValueError("No valid connections available for the current trial.") + + right_hip_joint = _resolve_joint_name("R_HIP", available_joints) + left_hip_joint = _resolve_joint_name("L_HIP", available_joints) # Animate the data fig = plt.figure(figsize=(8, 8)) @@ -145,25 +203,30 @@ def animate_trial( ax.view_init(elev=elev, azim=azim) # Prepare the lines to be updated - lines = { - connection: ax.plot([], [], [], c=player_color, lw=player_lw)[0] - for connection in connections - } + lines = [ + ax.plot([], [], [], c=player_color, lw=player_lw)[0] + for _ in resolved_connections + ] (ball,) = ax.plot([], [], [], "o", markersize=ball_size, c=ball_color) def update(frame): # Use the average of the right and left hip to center the view. - rh_xy = player_joint_dict["R_HIP"][frame][:2] - lh_xy = player_joint_dict["L_HIP"][frame][:2] - mh_xy = (rh_xy + lh_xy) / 2 + if right_hip_joint is not None and left_hip_joint is not None: + rh_xy = player_joint_dict[right_hip_joint][frame][:2] + lh_xy = player_joint_dict[left_hip_joint][frame][:2] + mh_xy = (rh_xy + lh_xy) / 2 + else: + # Fallback to any available joint if hip markers are unavailable. + first_joint = next(iter(player_joint_dict)) + mh_xy = player_joint_dict[first_joint][frame][:2] ax.set_xlim([mh_xy[0] - xbuffer, mh_xy[0] + xbuffer]) ax.set_ylim([mh_xy[1] - ybuffer, mh_xy[1] + ybuffer]) # Update the line data for each connection - for connection in connections: + for line, connection in zip(lines, resolved_connections): part1, part2 = connection x = [ player_joint_dict[part1][frame, 0], @@ -177,7 +240,7 @@ def update(frame): player_joint_dict[part1][frame, 2], player_joint_dict[part2][frame, 2], ] - lines[connection].set_data_3d(x, y, z) + line.set_data_3d(x, y, z) # Update ball data x = ball_data_array[frame, 0] @@ -202,5 +265,9 @@ def update(frame): plt.subplots(layout="constrained") plt.close() - anim = FuncAnimation(fig, update, frames=N_frames, interval=1000 / 30) + sampling_rate = data.get("sampling_rate", 30) + if not isinstance(sampling_rate, (float, int)) or sampling_rate <= 0: + sampling_rate = 30 + + anim = FuncAnimation(fig, update, frames=N_frames, interval=1000 / sampling_rate) return anim diff --git a/tests/test_animate.py b/tests/test_animate.py new file mode 100644 index 0000000..7f82cdc --- /dev/null +++ b/tests/test_animate.py @@ -0,0 +1,106 @@ +import importlib.util +import json +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +ANIMATE_PATH = REPO_ROOT / "basketball" / "freethrow" / "animate.py" + + +def _load_animate_module(): + spec = importlib.util.spec_from_file_location("animate_module", ANIMATE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _write_trial(path: Path, sampling_rate): + trial = { + "sampling_rate": sampling_rate, + "tracking": [ + { + "frame": 0, + "time": 0, + "data": { + "ball": [1.0, 2.0, 3.0], + "player": { + "RIGHT_HIP": [0.0, 0.0, 1.0], + "LEFT_HIP": [1.0, 0.0, 1.0], + "RIGHT_KNEE": [0.0, 0.0, 0.5], + "LEFT_KNEE": [1.0, 0.0, 0.5], + }, + }, + }, + { + "frame": 1, + "time": 16, + "data": { + "ball": [1.1, 2.1, 3.1], + "player": { + "RIGHT_HIP": [0.1, 0.0, 1.0], + "LEFT_HIP": [1.1, 0.0, 1.0], + "RIGHT_KNEE": [0.1, 0.0, 0.5], + "LEFT_KNEE": [1.1, 0.0, 0.5], + }, + }, + }, + ], + } + path.write_text(json.dumps(trial)) + + +def test_resolve_joint_name_uses_alias_when_needed(): + animate = _load_animate_module() + available_joints = {"RIGHT_HIP", "LEFT_HIP"} + assert animate._resolve_joint_name("R_HIP", available_joints) == "RIGHT_HIP" + assert animate._resolve_joint_name("L_HIP", available_joints) == "LEFT_HIP" + assert animate._resolve_joint_name("R_WRIST", available_joints) is None + + +def test_animate_trial_uses_trial_sampling_rate_for_interval(tmp_path, monkeypatch): + animate = _load_animate_module() + trial_path = tmp_path / "trial_60fps.json" + _write_trial(trial_path, sampling_rate=60) + + captured = {} + + def fake_func_animation(fig, update, frames, interval): + captured["frames"] = frames + captured["interval"] = interval + return {"ok": True} + + monkeypatch.setattr(animate, "FuncAnimation", fake_func_animation) + + anim = animate.animate_trial( + str(trial_path), + show_court=False, + notebook_mode=False, + connections=[("R_HIP", "R_KNEE"), ("L_HIP", "L_KNEE")], + ) + + assert anim == {"ok": True} + assert captured["frames"] == 2 + assert captured["interval"] == 1000 / 60 + + +def test_animate_trial_falls_back_to_30fps_for_invalid_sampling_rate(tmp_path, monkeypatch): + animate = _load_animate_module() + trial_path = tmp_path / "trial_invalid_fps.json" + _write_trial(trial_path, sampling_rate="bad_value") + + captured = {} + + def fake_func_animation(fig, update, frames, interval): + captured["interval"] = interval + return {"ok": True} + + monkeypatch.setattr(animate, "FuncAnimation", fake_func_animation) + + animate.animate_trial( + str(trial_path), + show_court=False, + notebook_mode=False, + connections=[("R_HIP", "R_KNEE"), ("L_HIP", "L_KNEE")], + ) + + assert captured["interval"] == 1000 / 30