Spaces:
Sleeping
Sleeping
File size: 4,574 Bytes
fb91fdc | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """
Main-figure pose extractor — gold 4x1000 subset.
Tests upgrade (a) from the pose investigation: instead of averaging all
skeletons in an image (extract_pose.py), store ONLY the largest detected
figure's torso-normalised skeleton, plus quality indicators so garbage
detections can be filtered in analysis:
filename
n_persons — YOLO detections in the image
main_conf — box confidence of the largest figure
main_kpt_conf — mean keypoint confidence of the largest figure
main_torso_px — torso height in pixels (tiny => unreliable normalisation)
main_area — bbox area fraction of the image
main_skel_0..33 — 17 kpts x (dx,dy), torso-normalised, 0 where invisible
Output: data/features/pose.parquet (checkpointed, resumable)
Usage: python features/extract_pose_main_gold.py [--model yolov8m-pose.pt]
"""
import argparse
import sys
from pathlib import Path
import numpy as np
import pandas as pd
from tqdm import tqdm
ROOT = Path(__file__).resolve().parent.parent
IMAGES = ROOT / "data/images"
SELECTED = ROOT / "data/artwork_metadata.csv"
OUTPUT = ROOT / "data/features/pose.parquet"
KPT_THR = 0.3
L_SHOULDER, R_SHOULDER, L_HIP, R_HIP = 5, 6, 11, 12
def normalise_skeleton(kpts, conf):
"""Torso-normalised skeleton (matches extract_pose.py) + torso height px."""
vis = conf >= KPT_THR
anchors = [L_SHOULDER, R_SHOULDER, L_HIP, R_HIP]
if sum(vis[i] for i in anchors) >= 2:
mid_shoulder = (kpts[L_SHOULDER] + kpts[R_SHOULDER]) / 2.0
mid_hip = (kpts[L_HIP] + kpts[R_HIP]) / 2.0
centre = (mid_shoulder + mid_hip) / 2.0
height = float(np.linalg.norm(mid_shoulder - mid_hip))
else:
visible = kpts[vis]
if len(visible) == 0:
return np.zeros(34, dtype=np.float32), 0.0
centre = visible.mean(axis=0)
height = float(np.linalg.norm(visible.max(axis=0) - visible.min(axis=0)))
norm = (kpts - centre) / max(height, 1.0)
norm[~vis] = 0.0
return norm.ravel().astype(np.float32), height
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="yolov8m-pose.pt")
ap.add_argument("--save-every", type=int, default=200)
args = ap.parse_args()
from ultralytics import YOLO
model = YOLO(args.model)
gold = (pd.read_csv(SELECTED, dtype=str)
.drop_duplicates("filename")[["filename"]])
existing = (pd.read_parquet(OUTPUT) if OUTPUT.exists()
else pd.DataFrame(columns=["filename"]))
have = set(existing["filename"])
todo = gold[~gold["filename"].isin(have)]["filename"].tolist()
print(f"gold={len(gold)} done={len(have)} todo={len(todo)}")
rows, failed = [], 0
for filename in tqdm(todo):
try:
res = model(str(IMAGES / filename), verbose=False)[0]
row = {"filename": filename, "n_persons": 0, "main_conf": 0.0,
"main_kpt_conf": 0.0, "main_torso_px": 0.0, "main_area": 0.0,
**{f"main_skel_{i}": 0.0 for i in range(34)}}
if res.boxes is not None and len(res.boxes) > 0:
areas = ((res.boxes.xyxy[:, 2] - res.boxes.xyxy[:, 0])
* (res.boxes.xyxy[:, 3] - res.boxes.xyxy[:, 1]))
j = int(areas.argmax())
kpts = res.keypoints.xy[j].cpu().numpy()
conf = (res.keypoints.conf[j].cpu().numpy()
if res.keypoints.conf is not None else np.ones(17))
skel, torso = normalise_skeleton(kpts, conf)
ih, iw = res.orig_shape
row.update({"n_persons": len(res.boxes),
"main_conf": float(res.boxes.conf[j]),
"main_kpt_conf": float(conf.mean()),
"main_torso_px": torso,
"main_area": float(areas[j]) / (ih * iw)})
row.update({f"main_skel_{i}": float(v) for i, v in enumerate(skel)})
rows.append(row)
except Exception as e:
failed += 1
sys.stderr.write(f"FAIL {filename}: {e}\n")
if len(rows) >= args.save_every:
existing = pd.concat([existing, pd.DataFrame(rows)], ignore_index=True)
existing.to_parquet(OUTPUT, index=False)
rows = []
if rows:
existing = pd.concat([existing, pd.DataFrame(rows)], ignore_index=True)
existing.to_parquet(OUTPUT, index=False)
print(f"Wrote {OUTPUT}: {len(existing)} rows. Failures: {failed}")
if __name__ == "__main__":
main()
|