vision-anti-pinch/pipeline/rtmw_l_person_mask_ablation.py

"""Test whether other people inside a top-down crop cause RTMW-L mix-ups."""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import cv2
import numpy as np
from rtmlib import RTMPose

from compare_rtmw_ml import (
    draw_pose,
    model_disagreement,
    summarize,
    temporal_metrics,
    write_montage,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--video", type=Path, required=True)
    parser.add_argument("--baseline-jsonl", type=Path, required=True)
    parser.add_argument("--model-l", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def mask_other_people(
    frame: np.ndarray,
    boxes: list[list[float]],
    target_idx: int,
) -> np.ndarray:
    masked = frame.copy()
    height, width = frame.shape[:2]
    mask = np.zeros((height, width), dtype=bool)
    target = boxes[target_idx]
    target_area = max((target[2] - target[0]) * (target[3] - target[1]), 1.0)
    for idx, bbox in enumerate(boxes):
        if idx == target_idx:
            continue
        other_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
        # The failure in this video is the small background person being
        # absorbed by the much larger foreground person's 3:4 pose crop.
        # Mask only markedly smaller people.  Masking a larger occluder while
        # estimating a small person removes genuine target pixels at overlap.
        if other_area >= 0.5 * target_area:
            continue
        x1, y1, x2, y2 = (int(round(value)) for value in bbox)
        x1, x2 = max(0, x1), min(width - 1, x2)
        y1, y2 = max(0, y1), min(height - 1, y2)
        mask[y1:y2 + 1, x1:x2 + 1] = True

    masked[mask] = (124, 116, 104)
    return masked


def main() -> None:
    args = parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    baseline = [
        json.loads(line)
        for line in args.baseline_jsonl.read_text(encoding="utf-8").splitlines()
    ]
    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))
    writer = cv2.VideoWriter(
        str(args.output / "rtmw_l_raw_vs_person_masked.mp4"),
        cv2.VideoWriter_fourcc(*"mp4v"),
        fps,
        (width, height // 2),
    )

    comparison: list[dict] = []
    timings: list[float] = []
    frame_idx = 0
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        baseline_record = baseline[frame_idx]
        boxes = baseline_record["boxes"][:2]
        masked_keypoints = []
        masked_scores = []
        started = time.perf_counter()
        for person_idx, bbox in enumerate(boxes):
            masked_frame = mask_other_people(frame, boxes, person_idx)
            keypoints, scores = pose_l(masked_frame, bboxes=[bbox])
            masked_keypoints.append(keypoints[0])
            masked_scores.append(scores[0])
        timings.append((time.perf_counter() - started) * 1000.0)
        masked_keypoints_array = np.asarray(masked_keypoints)

        raw_keypoints = np.asarray(baseline_record["L"]["keypoints"][:2])
        comparison.append(
            {
                "frame_index": frame_idx,
                "boxes": boxes,
                "M": {"keypoints": raw_keypoints.tolist()},
                "L": {"keypoints": np.round(masked_keypoints_array, 3).tolist()},
                "masked_raw_scores": np.round(np.asarray(masked_scores), 4).tolist(),
            }
        )

        left = draw_pose(frame, boxes, raw_keypoints, f"RTMW-L raw | frame {frame_idx}")
        right = draw_pose(
            frame,
            boxes,
            masked_keypoints_array,
            f"RTMW-L mask other person | 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()

    with (args.output / "rtmw_l_raw_vs_person_masked.jsonl").open(
        "w", encoding="utf-8"
    ) as handle:
        for record in comparison:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")

    disagreement, worst_frames = model_disagreement(comparison)
    fixed_frames = [0, 60, 120, 132, 136, 180, 240, 252, 300, 360, 411]
    montage_frames = sorted(set(fixed_frames + worst_frames[:8]))
    write_montage(
        args.video,
        comparison,
        montage_frames,
        args.output / "rtmw_l_raw_vs_person_masked_montage.jpg",
        titles=("RTMW-L raw", "RTMW-L smaller-person masked"),
    )

    summary = {
        "frames": len(comparison),
        "labels": {
            "M": "RTMW-L raw crop",
            "L": "RTMW-L with markedly smaller people masked",
        },
        "timing_ms_two_people_cpu": summarize(timings[5:]),
        "temporal_metrics_normalized_to_yolo_box": temporal_metrics(comparison),
        "position_change_after_masking": disagreement,
        "largest_change_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()