|
|
import os |
|
|
import argparse |
|
|
import pandas as pd |
|
|
import numpy as np |
|
|
from scipy.ndimage import gaussian_filter1d |
|
|
from pathlib import Path |
|
|
|
|
|
|
|
|
COCO_BODY_17 = [ |
|
|
"nose", "left_eye", "right_eye", "left_ear", "right_ear", |
|
|
"left_shoulder", "right_shoulder", "left_elbow", "right_elbow", |
|
|
"left_wrist", "right_wrist", "left_hip", "right_hip", |
|
|
"left_knee", "right_knee", "left_ankle", "right_ankle" |
|
|
] |
|
|
FEET_6_LABELS = [ |
|
|
"left_heel", "left_big_toe", "left_little_toe", |
|
|
"right_heel", "right_big_toe", "right_little_toe" |
|
|
] |
|
|
ALL_KEYPOINTS = COCO_BODY_17 + FEET_6_LABELS |
|
|
|
|
|
def parse_labels_to_df(labels_dir, img_width, img_height): |
|
|
records = [] |
|
|
label_files = sorted([f for f in os.listdir(labels_dir) if f.endswith('.txt')]) |
|
|
for frame_idx, label_file in enumerate(label_files): |
|
|
label_path = os.path.join(labels_dir, label_file) |
|
|
with open(label_path, 'r') as f: |
|
|
lines = f.readlines() |
|
|
for line in lines: |
|
|
parts = line.strip().split() |
|
|
if len(parts) < 5: continue |
|
|
keypoint_data = parts[5:] |
|
|
kp_dict = {} |
|
|
for i in range(0, len(keypoint_data), 3): |
|
|
kp_idx = i // 3 |
|
|
if kp_idx >= len(ALL_KEYPOINTS): break |
|
|
x_norm, y_norm, conf = float(keypoint_data[i]), float(keypoint_data[i+1]), float(keypoint_data[i+2]) |
|
|
if conf > 0: |
|
|
kp_dict[ALL_KEYPOINTS[kp_idx]] = {'x': x_norm * img_width, 'y': y_norm * img_height, 'conf': conf} |
|
|
records.append({'frame': frame_idx, 'keypoints': kp_dict}) |
|
|
return pd.DataFrame(records) |
|
|
|
|
|
def process_keypoints(df, sigma=2.0): |
|
|
"""Expand, interpolate, and smooth keypoints.""" |
|
|
|
|
|
rows = [] |
|
|
for _, row in df.iterrows(): |
|
|
flat = {'frame': row['frame']} |
|
|
for kp, vals in row['keypoints'].items(): |
|
|
flat[f"{kp}_x"] = vals['x'] |
|
|
flat[f"{kp}_y"] = vals['y'] |
|
|
flat[f"{kp}_conf"] = vals['conf'] |
|
|
rows.append(flat) |
|
|
|
|
|
expanded = pd.DataFrame(rows) |
|
|
|
|
|
for kp in ALL_KEYPOINTS: |
|
|
for suffix in ['_x', '_y']: |
|
|
col = f"{kp}{suffix}" |
|
|
if col in expanded.columns: |
|
|
expanded[col] = expanded[col].interpolate(method='linear').ffill().bfill() |
|
|
expanded[col] = gaussian_filter1d(expanded[col], sigma=sigma) |
|
|
|
|
|
return expanded |
|
|
|
|
|
if __name__ == "__main__": |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("--labels", required=True, help="Input labels directory") |
|
|
parser.add_argument("--output", required=True, help="Output CSV path") |
|
|
parser.add_argument("--width", type=int, default=1920) |
|
|
parser.add_argument("--height", type=int, default=1080) |
|
|
parser.add_argument("--sigma", type=float, default=2.0) |
|
|
args = parser.parse_args() |
|
|
|
|
|
print("Post-processing keypoints...") |
|
|
df_raw = parse_labels_to_df(args.labels, args.width, args.height) |
|
|
df_clean = process_keypoints(df_raw, sigma=args.sigma) |
|
|
df_clean.to_csv(args.output, index=False) |
|
|
print(f"✅ Cleaned keypoints saved to {args.output}") |
|
|
|