385 lines
14 KiB
Python
385 lines
14 KiB
Python
"""Compare RTMW-M 256x192 and RTMW-L 384x288 on cached YOLO boxes.
|
|
|
|
This is an experiment script. It intentionally keeps the candidate weights
|
|
outside the formal models/ directory and does not change deployment code.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import time
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from rtmlib import RTMPose
|
|
|
|
|
|
BODY_EDGES = ((5, 7), (7, 9), (6, 8), (8, 10))
|
|
HAND_EDGES_LOCAL = (
|
|
(0, 1), (1, 2), (2, 3), (3, 4),
|
|
(0, 5), (5, 6), (6, 7), (7, 8),
|
|
(0, 9), (9, 10), (10, 11), (11, 12),
|
|
(0, 13), (13, 14), (14, 15), (15, 16),
|
|
(0, 17), (17, 18), (18, 19), (19, 20),
|
|
)
|
|
HAND_SLICES = {"left": (91, 112), "right": (112, 133)}
|
|
BODY_WRISTS = {"left": 9, "right": 10}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--video", type=Path, required=True)
|
|
parser.add_argument("--boxes", type=Path, required=True)
|
|
parser.add_argument("--model-m", type=Path, required=True)
|
|
parser.add_argument("--model-l", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--max-people", type=int, default=2)
|
|
return parser.parse_args()
|
|
|
|
|
|
def load_boxes(path: Path, max_people: int) -> dict[int, list[list[float]]]:
|
|
frames: dict[int, list[list[float]]] = {}
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
item = json.loads(line)
|
|
boxes = [entry["xyxy"] for entry in item["boxes"]]
|
|
frames[int(item["frame_index"])] = sorted(boxes, key=lambda b: b[0])[:max_people]
|
|
return frames
|
|
|
|
|
|
def point(pt: np.ndarray) -> tuple[int, int]:
|
|
return int(round(float(pt[0]))), int(round(float(pt[1])))
|
|
|
|
|
|
def draw_pose(
|
|
image: np.ndarray,
|
|
boxes: list[list[float]],
|
|
keypoints: np.ndarray,
|
|
title: str,
|
|
) -> np.ndarray:
|
|
canvas = image.copy()
|
|
for person_idx, (bbox, kpts) in enumerate(zip(boxes, keypoints), start=1):
|
|
x1, y1, x2, y2 = (int(round(value)) for value in bbox)
|
|
cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2)
|
|
cv2.putText(
|
|
canvas,
|
|
f"P{person_idx}",
|
|
(x1 + 4, max(22, y1 + 24)),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.75,
|
|
(0, 220, 255),
|
|
2,
|
|
cv2.LINE_AA,
|
|
)
|
|
|
|
for edge_idx, (a, b) in enumerate(BODY_EDGES):
|
|
color = (40, 40, 255) if edge_idx < 2 else (30, 230, 30)
|
|
cv2.line(canvas, point(kpts[a]), point(kpts[b]), color, 3, cv2.LINE_AA)
|
|
for idx in (5, 6, 7, 8, 9, 10):
|
|
color = (40, 40, 255) if idx in (5, 7, 9) else (30, 230, 30)
|
|
cv2.circle(canvas, point(kpts[idx]), 5, color, -1, cv2.LINE_AA)
|
|
|
|
for side, (start, end) in HAND_SLICES.items():
|
|
hand = kpts[start:end]
|
|
color = (255, 80, 220) if side == "left" else (255, 190, 40)
|
|
for a, b in HAND_EDGES_LOCAL:
|
|
cv2.line(canvas, point(hand[a]), point(hand[b]), color, 2, cv2.LINE_AA)
|
|
for hand_point in hand:
|
|
cv2.circle(canvas, point(hand_point), 2, color, -1, cv2.LINE_AA)
|
|
|
|
cv2.rectangle(canvas, (0, 0), (610, 45), (0, 0, 0), -1)
|
|
cv2.putText(
|
|
canvas,
|
|
title,
|
|
(10, 31),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.8,
|
|
(255, 255, 255),
|
|
2,
|
|
cv2.LINE_AA,
|
|
)
|
|
return canvas
|
|
|
|
|
|
def normalized_keypoints(kpts: np.ndarray, bbox: list[float]) -> np.ndarray:
|
|
x1, y1, x2, y2 = bbox
|
|
width = max(x2 - x1, 1.0)
|
|
height = max(y2 - y1, 1.0)
|
|
result = kpts.copy().astype(np.float64)
|
|
result[:, 0] = (result[:, 0] - x1) / width
|
|
result[:, 1] = (result[:, 1] - y1) / height
|
|
return result
|
|
|
|
|
|
def hand_span(kpts: np.ndarray, bbox: list[float], side: str) -> float:
|
|
start, end = HAND_SLICES[side]
|
|
hand = kpts[start:end]
|
|
span = math.hypot(
|
|
float(np.ptp(hand[:, 0])),
|
|
float(np.ptp(hand[:, 1])),
|
|
)
|
|
x1, y1, x2, y2 = bbox
|
|
return span / max(math.hypot(x2 - x1, y2 - y1), 1.0)
|
|
|
|
|
|
def summarize(values: list[float]) -> dict[str, float | int | None]:
|
|
if not values:
|
|
return {"count": 0, "median": None, "p95": None, "max": None}
|
|
array = np.asarray(values, dtype=np.float64)
|
|
return {
|
|
"count": int(array.size),
|
|
"median": round(float(np.median(array)), 5),
|
|
"p95": round(float(np.percentile(array, 95)), 5),
|
|
"max": round(float(np.max(array)), 5),
|
|
}
|
|
|
|
|
|
def temporal_metrics(records: list[dict]) -> dict:
|
|
by_model_person: dict[tuple[str, int], list[dict]] = defaultdict(list)
|
|
for record in records:
|
|
for model_name in ("M", "L"):
|
|
for person_idx in range(len(record[model_name]["keypoints"])):
|
|
by_model_person[(model_name, person_idx)].append(
|
|
{
|
|
"frame": record["frame_index"],
|
|
"bbox": record["boxes"][person_idx],
|
|
"kpts": np.asarray(
|
|
record[model_name]["keypoints"][person_idx], dtype=np.float64
|
|
),
|
|
}
|
|
)
|
|
|
|
result: dict[str, dict] = {"M": {}, "L": {}}
|
|
for (model_name, person_idx), sequence in by_model_person.items():
|
|
upper_steps: list[float] = []
|
|
wrist_steps = {"left": [], "right": []}
|
|
hand_steps = {"left": [], "right": []}
|
|
hand_spans = {"left": [], "right": []}
|
|
bone_lengths = {"left_upper": [], "left_forearm": [],
|
|
"right_upper": [], "right_forearm": []}
|
|
|
|
previous: np.ndarray | None = None
|
|
for entry in sequence:
|
|
normalized = normalized_keypoints(entry["kpts"], entry["bbox"])
|
|
if previous is not None:
|
|
steps = np.linalg.norm(normalized - previous, axis=1)
|
|
upper_steps.extend(float(steps[idx]) for idx in (5, 6, 7, 8, 9, 10))
|
|
for side, idx in BODY_WRISTS.items():
|
|
wrist_steps[side].append(float(steps[idx]))
|
|
for side, (start, end) in HAND_SLICES.items():
|
|
hand_steps[side].extend(float(value) for value in steps[start:end])
|
|
previous = normalized
|
|
|
|
hand_spans["left"].append(hand_span(entry["kpts"], entry["bbox"], "left"))
|
|
hand_spans["right"].append(hand_span(entry["kpts"], entry["bbox"], "right"))
|
|
bone_lengths["left_upper"].append(float(np.linalg.norm(normalized[5] - normalized[7])))
|
|
bone_lengths["left_forearm"].append(float(np.linalg.norm(normalized[7] - normalized[9])))
|
|
bone_lengths["right_upper"].append(float(np.linalg.norm(normalized[6] - normalized[8])))
|
|
bone_lengths["right_forearm"].append(float(np.linalg.norm(normalized[8] - normalized[10])))
|
|
|
|
span_summary = {}
|
|
for side, values in hand_spans.items():
|
|
median = float(np.median(values))
|
|
gross_limit = max(0.2, median * 2.5)
|
|
span_summary[side] = {
|
|
**summarize(values),
|
|
"gross_limit": round(gross_limit, 5),
|
|
"gross_outlier_frames": int(sum(value > gross_limit for value in values)),
|
|
}
|
|
|
|
bone_summary = {}
|
|
for name, values in bone_lengths.items():
|
|
array = np.asarray(values, dtype=np.float64)
|
|
mean = float(np.mean(array))
|
|
bone_summary[name] = {
|
|
**summarize(values),
|
|
"coefficient_of_variation": round(float(np.std(array) / max(mean, 1e-9)), 5),
|
|
}
|
|
|
|
result[model_name][f"P{person_idx + 1}"] = {
|
|
"upper_body_frame_step": summarize(upper_steps),
|
|
"wrist_frame_step": {side: summarize(values) for side, values in wrist_steps.items()},
|
|
"hand_landmark_frame_step": {
|
|
side: summarize(values) for side, values in hand_steps.items()
|
|
},
|
|
"hand_span_relative_to_person": span_summary,
|
|
"bone_length_relative_to_box": bone_summary,
|
|
}
|
|
return result
|
|
|
|
|
|
def model_disagreement(records: list[dict]) -> tuple[dict, list[int]]:
|
|
values: dict[str, list[float]] = defaultdict(list)
|
|
per_frame: list[tuple[float, int]] = []
|
|
for record in records:
|
|
frame_max = 0.0
|
|
for person_idx, bbox in enumerate(record["boxes"]):
|
|
m = normalized_keypoints(
|
|
np.asarray(record["M"]["keypoints"][person_idx]), bbox
|
|
)
|
|
l = normalized_keypoints(
|
|
np.asarray(record["L"]["keypoints"][person_idx]), bbox
|
|
)
|
|
distance = np.linalg.norm(m - l, axis=1)
|
|
values[f"P{person_idx + 1}_upper"].extend(
|
|
float(distance[idx]) for idx in (5, 6, 7, 8, 9, 10)
|
|
)
|
|
for side, idx in BODY_WRISTS.items():
|
|
values[f"P{person_idx + 1}_{side}_wrist"].append(float(distance[idx]))
|
|
frame_max = max(frame_max, float(np.max(distance[[5, 6, 7, 8, 9, 10]])))
|
|
per_frame.append((frame_max, record["frame_index"]))
|
|
worst_frames = [frame for _, frame in sorted(per_frame, reverse=True)[:12]]
|
|
return {name: summarize(items) for name, items in values.items()}, worst_frames
|
|
|
|
|
|
def write_montage(
|
|
video_path: Path,
|
|
records: list[dict],
|
|
selected_frames: list[int],
|
|
output_path: Path,
|
|
titles: tuple[str, str] = ("RTMW-M 256x192", "RTMW-L 384x288"),
|
|
) -> None:
|
|
selected = set(selected_frames)
|
|
record_map = {record["frame_index"]: record for record in records}
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
rows: list[np.ndarray] = []
|
|
frame_idx = 0
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
if frame_idx in selected:
|
|
record = record_map[frame_idx]
|
|
m = draw_pose(
|
|
frame,
|
|
record["boxes"],
|
|
np.asarray(record["M"]["keypoints"]),
|
|
f"{titles[0]} | frame {frame_idx}",
|
|
)
|
|
l = draw_pose(
|
|
frame,
|
|
record["boxes"],
|
|
np.asarray(record["L"]["keypoints"]),
|
|
f"{titles[1]} | frame {frame_idx}",
|
|
)
|
|
m = cv2.resize(m, (640, 360), interpolation=cv2.INTER_AREA)
|
|
l = cv2.resize(l, (640, 360), interpolation=cv2.INTER_AREA)
|
|
rows.append(np.hstack([m, l]))
|
|
frame_idx += 1
|
|
cap.release()
|
|
if rows:
|
|
cv2.imwrite(str(output_path), np.vstack(rows))
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
cached_boxes = load_boxes(args.boxes, args.max_people)
|
|
|
|
pose_m = RTMPose(
|
|
str(args.model_m), model_input_size=(192, 256), backend="onnxruntime", device="cpu"
|
|
)
|
|
pose_l = RTMPose(
|
|
str(args.model_l), model_input_size=(288, 384), backend="onnxruntime", device="cpu"
|
|
)
|
|
|
|
cap = cv2.VideoCapture(str(args.video))
|
|
fps = float(cap.get(cv2.CAP_PROP_FPS))
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
video_output = args.output / "reference_video_rtmw_m_vs_l.mp4"
|
|
writer = cv2.VideoWriter(
|
|
str(video_output),
|
|
cv2.VideoWriter_fourcc(*"mp4v"),
|
|
fps,
|
|
(width, height // 2),
|
|
)
|
|
|
|
records: list[dict] = []
|
|
timings = {"M": [], "L": []}
|
|
frame_idx = 0
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
boxes = cached_boxes[frame_idx]
|
|
|
|
started = time.perf_counter()
|
|
kpts_m, scores_m = pose_m(frame, bboxes=boxes)
|
|
timings["M"].append((time.perf_counter() - started) * 1000.0)
|
|
|
|
started = time.perf_counter()
|
|
kpts_l, scores_l = pose_l(frame, bboxes=boxes)
|
|
timings["L"].append((time.perf_counter() - started) * 1000.0)
|
|
|
|
records.append(
|
|
{
|
|
"frame_index": frame_idx,
|
|
"boxes": boxes,
|
|
"M": {
|
|
"keypoints": np.round(kpts_m, 3).tolist(),
|
|
"raw_scores": np.round(scores_m, 4).tolist(),
|
|
},
|
|
"L": {
|
|
"keypoints": np.round(kpts_l, 3).tolist(),
|
|
"raw_scores": np.round(scores_l, 4).tolist(),
|
|
},
|
|
}
|
|
)
|
|
|
|
left = draw_pose(frame, boxes, kpts_m, f"RTMW-M 256x192 | frame {frame_idx}")
|
|
right = draw_pose(frame, boxes, kpts_l, f"RTMW-L 384x288 | frame {frame_idx}")
|
|
left = cv2.resize(left, (width // 2, height // 2), interpolation=cv2.INTER_AREA)
|
|
right = cv2.resize(right, (width // 2, height // 2), interpolation=cv2.INTER_AREA)
|
|
writer.write(np.hstack([left, right]))
|
|
|
|
if frame_idx % 50 == 0:
|
|
print(f"processed frame {frame_idx}", flush=True)
|
|
frame_idx += 1
|
|
|
|
cap.release()
|
|
writer.release()
|
|
|
|
jsonl_path = args.output / "reference_video_rtmw_m_vs_l.jsonl"
|
|
with jsonl_path.open("w", encoding="utf-8") as handle:
|
|
for record in records:
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
disagreement, worst_frames = model_disagreement(records)
|
|
fixed_frames = [0, 60, 120, 180, 240, 252, 300, 360, 411]
|
|
montage_frames = sorted(set(fixed_frames + worst_frames))
|
|
write_montage(
|
|
args.video,
|
|
records,
|
|
montage_frames,
|
|
args.output / "rtmw_m_vs_l_full_montage.jpg",
|
|
)
|
|
|
|
summary = {
|
|
"frames": len(records),
|
|
"input_sizes_wh": {"M": [192, 256], "L": [288, 384]},
|
|
"raw_score_note": (
|
|
"M and L raw SimCC peak scores use different scales and must not be compared "
|
|
"or thresholded with the same numeric value."
|
|
),
|
|
"timing_ms_two_people_cpu": {
|
|
name: summarize(values[5:] if len(values) > 5 else values)
|
|
for name, values in timings.items()
|
|
},
|
|
"temporal_metrics_normalized_to_yolo_box": temporal_metrics(records),
|
|
"m_l_position_disagreement_normalized_to_yolo_box": disagreement,
|
|
"largest_upper_body_disagreement_frames": worst_frames,
|
|
"montage_frames": montage_frames,
|
|
}
|
|
with (args.output / "summary.json").open("w", encoding="utf-8") as handle:
|
|
json.dump(summary, handle, ensure_ascii=False, indent=2)
|
|
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|