import argparse import sys from pathlib import Path import cv2 import mediapipe as mp import numpy as np import pandas as pd PROJECT_ROOT = Path(__file__).resolve().parents[2] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from scripts.evaluate.rep_counting_methods import EXERCISE_CONFIGS, FixedThresholdFSMCounter, SmoothingBuffer, extract_primary_angle, normalize_exercise_name def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--manifest-file", required=True) parser.add_argument("--output-dir", default="results/eval_rep_counting") parser.add_argument("--min-visibility", type=float, default=0.5) return parser.parse_args() def load_pose_module(): return mp.solutions.pose def build_landmark_indices(mp_pose): names = [ "LEFT_SHOULDER", "RIGHT_SHOULDER", "LEFT_HIP", "RIGHT_HIP", "LEFT_KNEE", "RIGHT_KNEE", "LEFT_ELBOW", "RIGHT_ELBOW", "LEFT_WRIST", "RIGHT_WRIST", "LEFT_ANKLE", "RIGHT_ANKLE", ] return {name: mp_pose.PoseLandmark[name].value for name in names} def extract_landmark_points(results, landmark_indices, min_visibility): if not results.pose_landmarks: return None points = {} for name, index in landmark_indices.items(): landmark = results.pose_landmarks.landmark[index] if landmark.visibility >= min_visibility: points[name] = np.array([landmark.x, landmark.y, landmark.z], dtype=np.float32) else: points[name] = np.array([0.0, 0.0, 0.0], dtype=np.float32) return points def evaluate_video(video_path, exercise_label, pose_estimator, landmark_indices, min_visibility): config = EXERCISE_CONFIGS[exercise_label] fixed_counter = FixedThresholdFSMCounter(config.fixed_low, config.fixed_high, config.min_state_frames) smoothing = SmoothingBuffer(window_size=config.smoothing_window) capture = cv2.VideoCapture(str(video_path)) processed_frames = 0 valid_angle_frames = 0 while capture.isOpened(): read_ok, frame_bgr = capture.read() if not read_ok: break frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) pose_result = pose_estimator.process(frame_rgb) points = extract_landmark_points(pose_result, landmark_indices, min_visibility) if points is None: processed_frames += 1 continue raw_angle = extract_primary_angle(points, config) smoothed_angle = smoothing.update(raw_angle) if not np.isnan(smoothed_angle): valid_angle_frames += 1 fixed_counter.update(smoothed_angle) processed_frames += 1 capture.release() return { "paper_fsm": fixed_counter.reps, "processed_frames": processed_frames, "valid_angle_frames": valid_angle_frames, } def compute_error_metrics(predicted_reps, true_reps): absolute_error = abs(predicted_reps - true_reps) relative_error_percent = (absolute_error / true_reps) * 100.0 if true_reps > 0 else 0.0 missed_reps = max(0, true_reps - predicted_reps) false_reps = max(0, predicted_reps - true_reps) return absolute_error, relative_error_percent, missed_reps, false_reps def validate_manifest_columns(dataframe): required = {"video_path", "exercise_label"} missing = required - set(dataframe.columns) if missing: missing_text = ", ".join(sorted(missing)) raise ValueError(f"Manifest is missing required columns: {missing_text}") def main(): args = parse_args() output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) manifest = pd.read_csv(args.manifest_file) validate_manifest_columns(manifest) mp_pose = load_pose_module() landmark_indices = build_landmark_indices(mp_pose) per_video_rows = [] has_true_reps = "true_reps" in manifest.columns with mp_pose.Pose( static_image_mode=False, model_complexity=1, enable_segmentation=False, min_detection_confidence=0.5, min_tracking_confidence=0.5, ) as pose_estimator: for row in manifest.itertuples(index=False): video_path = Path(row.video_path) raw_exercise = str(row.exercise_label) exercise_label = normalize_exercise_name(raw_exercise) true_reps = int(row.true_reps) if has_true_reps and not pd.isna(row.true_reps) else None if exercise_label not in EXERCISE_CONFIGS: raise ValueError(f"Unsupported exercise '{raw_exercise}' from manifest row: {video_path}") if not video_path.exists(): raise FileNotFoundError(f"Video file not found: {video_path}") counts = evaluate_video( video_path=video_path, exercise_label=exercise_label, pose_estimator=pose_estimator, landmark_indices=landmark_indices, min_visibility=args.min_visibility, ) pred_reps = counts["paper_fsm"] row_data = { "video_path": str(video_path), "exercise_label": exercise_label, "method": "paper_fsm", "predicted_reps": pred_reps, "processed_frames": counts["processed_frames"], "valid_angle_frames": counts["valid_angle_frames"], } if true_reps is not None: abs_err, rel_err, missed, false = compute_error_metrics(pred_reps, true_reps) row_data["true_reps"] = true_reps row_data["absolute_count_error"] = abs_err row_data["relative_error_percent"] = rel_err row_data["missed_reps"] = missed row_data["false_reps"] = false per_video_rows.append(row_data) print( f"{video_path.name} | {exercise_label} | " f"predicted_reps={counts['paper_fsm']}" ) per_video_df = pd.DataFrame(per_video_rows) if "absolute_count_error" in per_video_df.columns: summary_df = ( per_video_df.groupby("method", as_index=False) .agg( videos=("video_path", "count"), mean_absolute_count_error=("absolute_count_error", "mean"), mean_relative_error_percent=("relative_error_percent", "mean"), total_missed_reps=("missed_reps", "sum"), total_false_reps=("false_reps", "sum"), ) .sort_values(by=["mean_absolute_count_error", "total_false_reps"], ascending=[True, True]) ) else: summary_df = per_video_df.groupby("method", as_index=False).agg(videos=("video_path", "count")) per_video_file = output_dir / "rep_counting_per_video.csv" summary_file = output_dir / "rep_counting_summary.csv" per_video_df.to_csv(per_video_file, index=False) summary_df.to_csv(summary_file, index=False) print("\nRep counting summary") print(summary_df.to_string(index=False)) print(f"\nSaved: {per_video_file}") print(f"Saved: {summary_file}") if __name__ == "__main__": main()