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