|
|
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) |
|
|
|
