File size: 2,356 Bytes
6b5b22f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import os
import argparse
import pandas as pd
import cv2
import numpy as np
# Skeleton definition
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))
# Draw Skeleton
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)
# Draw Keypoints
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)
|