| import os |
| import argparse |
| import pandas as pd |
| import cv2 |
| import numpy as np |
|
|
| |
| SKELETON = [ |
| ("left_ankle", "left_knee"), ("left_knee", "left_hip"), |
| ("right_ankle", "right_knee"), ("right_knee", "right_hip"), |
| ("left_hip", "right_hip"), ("left_shoulder", "right_shoulder"), |
| ("left_shoulder", "left_hip"), ("right_shoulder", "right_hip"), |
| ("left_shoulder", "left_elbow"), ("left_elbow", "left_wrist"), |
| ("right_shoulder", "right_elbow"), ("right_elbow", "right_wrist"), |
| ("left_ankle", "left_heel"), ("left_ankle", "left_big_toe"), |
| ("right_ankle", "right_heel"), ("right_ankle", "right_big_toe") |
| ] |
|
|
| def annotate_frames(frames_dir, output_dir, csv_path): |
| """Draw keypoints and skeleton on frames.""" |
| df = pd.read_csv(csv_path) |
| os.makedirs(output_dir, exist_ok=True) |
| |
| frame_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.jpg')]) |
| |
| for _, row in df.iterrows(): |
| f_idx = int(row['frame']) |
| if f_idx >= len(frame_files): continue |
| |
| frame_name = frame_files[f_idx] |
| img = cv2.imread(os.path.join(frames_dir, frame_name)) |
| |
| |
| for start_kp, end_kp in SKELETON: |
| s_x, s_y = row.get(f"{start_kp}_x"), row.get(f"{start_kp}_y") |
| e_x, e_y = row.get(f"{end_kp}_x"), row.get(f"{end_kp}_y") |
| if pd.notna(s_x) and pd.notna(e_x): |
| cv2.line(img, (int(s_x), int(s_y)), (int(e_x), int(e_y)), (0, 255, 0), 2) |
| |
| |
| for col in df.columns: |
| if col.endswith('_x'): |
| kp_name = col[:-2] |
| x, y = row[col], row[f"{kp_name}_y"] |
| if pd.notna(x): |
| cv2.circle(img, (int(x), int(y)), 4, (0, 0, 255), -1) |
| |
| cv2.imwrite(os.path.join(output_dir, frame_name), img) |
| if f_idx % 100 == 0: |
| print(f"Annotated {f_idx} frames") |
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--frames", required=True, help="Input frames directory") |
| parser.add_argument("--csv", required=True, help="Input keypoints CSV") |
| parser.add_argument("--output", required=True, help="Output frames directory") |
| args = parser.parse_args() |
| |
| annotate_frames(args.frames, args.output, args.csv) |
|
|
