683 lines
24 KiB
Python
683 lines
24 KiB
Python
"""End-to-end optimized RTMW-L prediction for CMS videos.
|
|
|
|
Pipeline:
|
|
YOLO11s person -> fixed-camera false-positive filters -> lightweight track IDs
|
|
-> RTMW-L raw/smaller-person-isolated views -> temporal fusion -> causal
|
|
occlusion-aware V4 global hand assignment. Reachable regions stay in JSON and are hidden in the main
|
|
comparison video.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import time
|
|
from collections import Counter, defaultdict, deque
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from rtmlib import RTMPose
|
|
from ultralytics import YOLO
|
|
|
|
from compare_rtmw_ml import HAND_SLICES, draw_pose, summarize
|
|
from fuse_rtmw_l_temporal import (
|
|
BODY_SIDE_INDICES,
|
|
PersonTrack,
|
|
adaptive_smooth,
|
|
build_global_hands,
|
|
draw_fused,
|
|
fuse_person,
|
|
)
|
|
from occlusion_aware_pose_v2 import (
|
|
EndpointTrack,
|
|
Measurement,
|
|
calibrated_bone_lengths,
|
|
draw_v2,
|
|
reachable_region,
|
|
resolve_global_hand_assignments,
|
|
solve_two_bone_elbow,
|
|
temporal_recovery_measurement,
|
|
update_endpoint_track,
|
|
)
|
|
from rtmw_l_person_mask_ablation import mask_other_people
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--videos", type=Path, nargs="+", required=True)
|
|
parser.add_argument("--det-model", type=Path, required=True)
|
|
parser.add_argument("--pose-model", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--det-device",
|
|
default="0",
|
|
help="Ultralytics YOLO device, for example 0, cuda:0, or cpu.",
|
|
)
|
|
parser.add_argument(
|
|
"--pose-device",
|
|
default="cpu",
|
|
help="RTMW ONNX Runtime device: cpu or cuda.",
|
|
)
|
|
parser.add_argument("--det-imgsz", type=int, default=960)
|
|
parser.add_argument("--det-conf", type=float, default=0.30)
|
|
parser.add_argument("--det-roi-y-max-ratio", type=float, default=1.0)
|
|
parser.add_argument("--short-occlusion-frames", type=int, default=12)
|
|
parser.add_argument("--medium-occlusion-frames", type=int, default=30)
|
|
parser.add_argument("--max-people", type=int, default=3)
|
|
return parser.parse_args()
|
|
|
|
|
|
def bbox_iou(a: np.ndarray, b: np.ndarray) -> float:
|
|
x1 = max(float(a[0]), float(b[0]))
|
|
y1 = max(float(a[1]), float(b[1]))
|
|
x2 = min(float(a[2]), float(b[2]))
|
|
y2 = min(float(a[3]), float(b[3]))
|
|
intersection = max(x2 - x1, 0.0) * max(y2 - y1, 0.0)
|
|
union = (
|
|
max(float(a[2] - a[0]), 0.0) * max(float(a[3] - a[1]), 0.0)
|
|
+ max(float(b[2] - b[0]), 0.0) * max(float(b[3] - b[1]), 0.0)
|
|
- intersection
|
|
)
|
|
return intersection / max(union, 1.0)
|
|
|
|
|
|
class BoxTracker:
|
|
def __init__(self, max_age: int = 30):
|
|
self.max_age = max_age
|
|
self.next_id = 1
|
|
self.tracks: dict[int, dict] = {}
|
|
|
|
def update(
|
|
self,
|
|
boxes: list[list[float]],
|
|
frame_idx: int,
|
|
image_shape: tuple[int, int],
|
|
) -> list[int]:
|
|
height, width = image_shape
|
|
image_diagonal = math.hypot(width, height)
|
|
active_ids = [
|
|
track_id
|
|
for track_id, track in self.tracks.items()
|
|
if frame_idx - track["last_frame"] <= self.max_age
|
|
]
|
|
pairs = []
|
|
for detection_idx, box in enumerate(boxes):
|
|
box_array = np.asarray(box, dtype=np.float64)
|
|
center = (box_array[:2] + box_array[2:]) * 0.5
|
|
for track_id in active_ids:
|
|
prior = np.asarray(self.tracks[track_id]["box"], dtype=np.float64)
|
|
prior_center = (prior[:2] + prior[2:]) * 0.5
|
|
iou = bbox_iou(box_array, prior)
|
|
center_distance = float(np.linalg.norm(center - prior_center)) / image_diagonal
|
|
if iou >= 0.05 or center_distance <= 0.18:
|
|
score = 2.0 * iou + max(0.0, 1.0 - center_distance / 0.18)
|
|
pairs.append((score, detection_idx, track_id))
|
|
|
|
assigned_detections = set()
|
|
assigned_tracks = set()
|
|
result: list[int | None] = [None] * len(boxes)
|
|
for _, detection_idx, track_id in sorted(pairs, reverse=True):
|
|
if detection_idx in assigned_detections or track_id in assigned_tracks:
|
|
continue
|
|
assigned_detections.add(detection_idx)
|
|
assigned_tracks.add(track_id)
|
|
result[detection_idx] = track_id
|
|
|
|
for detection_idx, box in enumerate(boxes):
|
|
track_id = result[detection_idx]
|
|
if track_id is None:
|
|
track_id = self.next_id
|
|
self.next_id += 1
|
|
result[detection_idx] = track_id
|
|
self.tracks[track_id] = {"box": box, "last_frame": frame_idx}
|
|
return [int(track_id) for track_id in result]
|
|
|
|
|
|
@dataclass
|
|
class LivePersonState:
|
|
fusion: PersonTrack = field(default_factory=PersonTrack)
|
|
endpoints: dict[str, EndpointTrack] = field(
|
|
default_factory=lambda: {"left": EndpointTrack(), "right": EndpointTrack()}
|
|
)
|
|
bone_samples: dict[str, deque] = field(
|
|
default_factory=lambda: {
|
|
"left": deque(maxlen=120),
|
|
"right": deque(maxlen=120),
|
|
}
|
|
)
|
|
last_frame: int = -1
|
|
|
|
|
|
def bbox_diagonal(bbox: list[float]) -> float:
|
|
return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0)
|
|
|
|
|
|
def update_bone_samples(
|
|
state: LivePersonState,
|
|
keypoints: np.ndarray,
|
|
metadata: dict,
|
|
bbox: list[float],
|
|
) -> None:
|
|
diagonal = bbox_diagonal(bbox)
|
|
for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items():
|
|
meta = metadata[side]
|
|
if meta["limb_state"] != "consensus" or not meta["distance_eligible"]:
|
|
continue
|
|
upper = float(np.linalg.norm(keypoints[shoulder] - keypoints[elbow])) / diagonal
|
|
forearm = float(np.linalg.norm(keypoints[elbow] - keypoints[wrist])) / diagonal
|
|
if 0.04 <= upper <= 0.65 and 0.04 <= forearm <= 0.65:
|
|
state.bone_samples[side].append((upper, forearm))
|
|
|
|
|
|
def current_bone_ratios(state: LivePersonState, side: str) -> tuple[float, float]:
|
|
samples = state.bone_samples[side]
|
|
if len(samples) < 10:
|
|
opposite = "right" if side == "left" else "left"
|
|
if len(state.bone_samples[opposite]) >= 10:
|
|
samples = state.bone_samples[opposite]
|
|
if not samples:
|
|
return 0.24, 0.22
|
|
values = np.asarray(samples, dtype=np.float64)
|
|
return float(np.median(values[:, 0])), float(np.median(values[:, 1]))
|
|
|
|
|
|
def update_live_v2(
|
|
frame_idx: int,
|
|
bbox: list[float],
|
|
selected_keypoints: np.ndarray,
|
|
smoothed_keypoints: np.ndarray,
|
|
fusion_metadata: dict,
|
|
state: LivePersonState,
|
|
assigned_measurements: dict[str, Measurement | None] | None,
|
|
identity_assignments: dict[str, dict] | None,
|
|
short_occlusion_frames: int,
|
|
medium_occlusion_frames: int,
|
|
) -> tuple[np.ndarray, dict]:
|
|
frame_gap = 0 if state.last_frame < 0 else max(frame_idx - state.last_frame - 1, 0)
|
|
for endpoint in state.endpoints.values():
|
|
endpoint.missed_frames += frame_gap
|
|
state.last_frame = frame_idx
|
|
update_bone_samples(state, selected_keypoints, fusion_metadata, bbox)
|
|
|
|
output = smoothed_keypoints.copy()
|
|
output_metadata = {}
|
|
diagonal = bbox_diagonal(bbox)
|
|
for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items():
|
|
endpoint = state.endpoints[side]
|
|
meta = fusion_metadata[side]
|
|
start, end = HAND_SLICES[side]
|
|
measurement = (
|
|
assigned_measurements.get(side)
|
|
if assigned_measurements is not None
|
|
else None
|
|
)
|
|
if assigned_measurements is None and meta["alarm_eligible"]:
|
|
measurement = Measurement(
|
|
keypoints=selected_keypoints[start:end],
|
|
precise=bool(meta["distance_eligible"]),
|
|
provisional=not bool(meta["distance_eligible"]),
|
|
quality=float(meta.get("selected_hand_score", 0.0)),
|
|
)
|
|
if measurement is None:
|
|
measurement = temporal_recovery_measurement(
|
|
selected_keypoints,
|
|
meta,
|
|
side,
|
|
bbox,
|
|
endpoint,
|
|
medium_occlusion_frames,
|
|
)
|
|
returning_after_long = (
|
|
measurement is not None
|
|
and endpoint.missed_frames > short_occlusion_frames
|
|
)
|
|
wrist, hand, endpoint_state = update_endpoint_track(
|
|
endpoint,
|
|
measurement,
|
|
frame_idx,
|
|
short_occlusion_frames,
|
|
medium_occlusion_frames,
|
|
)
|
|
shoulder = smoothed_keypoints[shoulder_idx]
|
|
upper_ratio, forearm_ratio = current_bone_ratios(state, side)
|
|
upper_length, forearm_length = calibrated_bone_lengths(
|
|
upper_ratio, forearm_ratio, bbox
|
|
)
|
|
region = None
|
|
if wrist is not None and hand is not None:
|
|
reference_elbow = (
|
|
selected_keypoints[elbow_idx]
|
|
if returning_after_long or endpoint.last_elbow is None
|
|
else endpoint.last_elbow
|
|
)
|
|
use_ik = (
|
|
meta["limb_state"] != "consensus"
|
|
or endpoint_state
|
|
in ("provisional", "recovered", "reassigned", "predicted_short")
|
|
)
|
|
elbow = (
|
|
solve_two_bone_elbow(
|
|
shoulder,
|
|
wrist,
|
|
upper_length,
|
|
forearm_length,
|
|
reference_elbow,
|
|
)
|
|
if use_ik
|
|
else smoothed_keypoints[elbow_idx]
|
|
)
|
|
endpoint.last_elbow = elbow.copy()
|
|
direction = wrist - shoulder
|
|
if float(np.linalg.norm(direction)) > 1e-6:
|
|
endpoint.last_direction = direction / np.linalg.norm(direction)
|
|
output[shoulder_idx] = shoulder
|
|
output[elbow_idx] = elbow
|
|
output[wrist_idx] = wrist
|
|
output[start:end] = hand
|
|
uncertainty = (
|
|
0.025 * diagonal
|
|
if endpoint_state == "observed"
|
|
else (0.04 + 0.015 * endpoint.missed_frames) * diagonal
|
|
)
|
|
precise = endpoint_state == "observed"
|
|
else:
|
|
region = reachable_region(
|
|
shoulder,
|
|
upper_length,
|
|
forearm_length,
|
|
bbox,
|
|
endpoint,
|
|
endpoint_state,
|
|
short_occlusion_frames,
|
|
medium_occlusion_frames,
|
|
)
|
|
uncertainty = region["outer_radius_px"]
|
|
precise = False
|
|
output_metadata[side] = {
|
|
"state": endpoint_state,
|
|
"precise_distance": precise,
|
|
"fast_alarm_candidate": True,
|
|
"risk_mode": (
|
|
"precise_hand"
|
|
if endpoint_state == "observed"
|
|
else (
|
|
"motion_corridor_intersection"
|
|
if endpoint_state == "reachable_medium"
|
|
else (
|
|
"door_intersection_only"
|
|
if endpoint_state == "reachable_long"
|
|
else (
|
|
"reassigned_hand"
|
|
if endpoint_state == "reassigned"
|
|
else "temporal_hand"
|
|
)
|
|
)
|
|
)
|
|
),
|
|
"requires_door_intersection": endpoint_state
|
|
in ("reachable_medium", "reachable_long"),
|
|
"identity_assignment": (
|
|
identity_assignments.get(side)
|
|
if identity_assignments is not None
|
|
else None
|
|
),
|
|
"missed_frames": endpoint.missed_frames,
|
|
"uncertainty_radius_px": round(float(uncertainty), 3),
|
|
"bone_lengths_px": {
|
|
"upper": round(upper_length, 3),
|
|
"forearm": round(forearm_length, 3),
|
|
},
|
|
"reachable_region": region,
|
|
}
|
|
return output, output_metadata
|
|
|
|
|
|
def filter_person_boxes(
|
|
result,
|
|
width: int,
|
|
height: int,
|
|
max_people: int,
|
|
) -> tuple[list[list[float]], list[float], Counter]:
|
|
boxes = result.boxes.xyxy.cpu().numpy()
|
|
scores = result.boxes.conf.cpu().numpy()
|
|
rejected = Counter()
|
|
kept = []
|
|
for box, score in zip(boxes, scores):
|
|
center_x = float(box[0] + box[2]) * 0.5
|
|
box_height = float(box[3] - box[1])
|
|
if center_x > 0.80 * width:
|
|
rejected["right_reflection_roi"] += 1
|
|
continue
|
|
if box_height < 0.20 * height:
|
|
rejected["too_short"] += 1
|
|
continue
|
|
kept.append((box.tolist(), float(score)))
|
|
kept.sort(key=lambda item: item[1], reverse=True)
|
|
if len(kept) > max_people:
|
|
rejected["over_max_people"] += len(kept) - max_people
|
|
kept = kept[:max_people]
|
|
kept.sort(key=lambda item: item[0][0])
|
|
return [item[0] for item in kept], [item[1] for item in kept], rejected
|
|
|
|
|
|
def process_video(
|
|
video_path: Path,
|
|
output_dir: Path,
|
|
detector: YOLO,
|
|
pose: RTMPose,
|
|
det_device: str,
|
|
det_imgsz: int,
|
|
det_conf: float,
|
|
det_roi_y_max_ratio: float,
|
|
short_occlusion_frames: int,
|
|
medium_occlusion_frames: int,
|
|
max_people: int,
|
|
video_label: str,
|
|
) -> dict:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
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(output_dir / "raw_temporal_occlusion_v4.mp4"),
|
|
cv2.VideoWriter_fourcc(*"mp4v"),
|
|
fps,
|
|
(1920, 360),
|
|
)
|
|
|
|
box_tracker = BoxTracker(max_age=30)
|
|
people: dict[int, LivePersonState] = {}
|
|
records = []
|
|
detection_counts = Counter()
|
|
rejection_counts = Counter()
|
|
state_counts = Counter()
|
|
timings = defaultdict(list)
|
|
isolated_inferences = 0
|
|
frame_idx = 0
|
|
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
started = time.perf_counter()
|
|
detection_height = max(1, min(height, round(height * det_roi_y_max_ratio)))
|
|
detection = detector.predict(
|
|
frame[:detection_height],
|
|
imgsz=det_imgsz,
|
|
conf=det_conf,
|
|
iou=0.30,
|
|
classes=[0],
|
|
device=det_device,
|
|
verbose=False,
|
|
)[0]
|
|
timings["detector_ms"].append((time.perf_counter() - started) * 1000.0)
|
|
boxes, confidences, rejected = filter_person_boxes(
|
|
detection, width, height, max_people
|
|
)
|
|
rejection_counts.update(rejected)
|
|
detection_counts[len(boxes)] += 1
|
|
track_ids = box_tracker.update(boxes, frame_idx, (height, width))
|
|
|
|
if boxes:
|
|
started = time.perf_counter()
|
|
raw_keypoints, raw_scores = pose(frame, bboxes=boxes)
|
|
timings["raw_pose_ms"].append((time.perf_counter() - started) * 1000.0)
|
|
isolated_keypoints = raw_keypoints.copy()
|
|
isolated_scores = raw_scores.copy()
|
|
started = time.perf_counter()
|
|
for person_idx, bbox in enumerate(boxes):
|
|
target_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
|
|
needs_isolation = any(
|
|
idx != person_idx
|
|
and max((other[2] - other[0]) * (other[3] - other[1]), 1.0)
|
|
< 0.5 * target_area
|
|
for idx, other in enumerate(boxes)
|
|
)
|
|
if not needs_isolation:
|
|
continue
|
|
masked = mask_other_people(frame, boxes, person_idx)
|
|
kpts, scores = pose(masked, bboxes=[bbox])
|
|
isolated_keypoints[person_idx] = kpts[0]
|
|
isolated_scores[person_idx] = scores[0]
|
|
isolated_inferences += 1
|
|
timings["isolated_extra_ms"].append((time.perf_counter() - started) * 1000.0)
|
|
else:
|
|
raw_keypoints = np.empty((0, 133, 2), dtype=np.float64)
|
|
raw_scores = np.empty((0, 133), dtype=np.float64)
|
|
isolated_keypoints = raw_keypoints.copy()
|
|
isolated_scores = raw_scores.copy()
|
|
|
|
selected_list = []
|
|
smoothed_list = []
|
|
fusion_meta_list = []
|
|
person_states = []
|
|
for person_idx, (bbox, track_id) in enumerate(zip(boxes, track_ids)):
|
|
state = people.setdefault(track_id, LivePersonState())
|
|
selected, fusion_meta = fuse_person(
|
|
raw_keypoints[person_idx],
|
|
raw_scores[person_idx],
|
|
isolated_keypoints[person_idx],
|
|
isolated_scores[person_idx],
|
|
bbox,
|
|
state.fusion,
|
|
)
|
|
smoothed = adaptive_smooth(selected, bbox, state.fusion)
|
|
selected_list.append(selected)
|
|
smoothed_list.append(smoothed)
|
|
fusion_meta_list.append(fusion_meta)
|
|
person_states.append(state)
|
|
|
|
frame_measurements = {}
|
|
endpoint_tracks = {}
|
|
for person_idx, (selected, fusion_meta, state) in enumerate(
|
|
zip(selected_list, fusion_meta_list, person_states)
|
|
):
|
|
for side, (start, end) in HAND_SLICES.items():
|
|
meta = fusion_meta[side]
|
|
frame_measurements[(person_idx, side)] = (
|
|
Measurement(
|
|
keypoints=selected[start:end],
|
|
precise=bool(meta["distance_eligible"]),
|
|
provisional=not bool(meta["distance_eligible"]),
|
|
quality=float(meta.get("selected_hand_score", 0.0)),
|
|
)
|
|
if meta["alarm_eligible"]
|
|
else None
|
|
)
|
|
endpoint_tracks[(person_idx, side)] = state.endpoints[side]
|
|
frame_measurements, assignment_events = resolve_global_hand_assignments(
|
|
frame_measurements,
|
|
endpoint_tracks,
|
|
boxes,
|
|
medium_occlusion_frames,
|
|
)
|
|
|
|
v2_list = []
|
|
v2_meta_list = []
|
|
for person_idx, (bbox, selected, smoothed, fusion_meta, state) in enumerate(
|
|
zip(
|
|
boxes,
|
|
selected_list,
|
|
smoothed_list,
|
|
fusion_meta_list,
|
|
person_states,
|
|
)
|
|
):
|
|
v2, v2_meta = update_live_v2(
|
|
frame_idx,
|
|
bbox,
|
|
selected,
|
|
smoothed,
|
|
fusion_meta,
|
|
state,
|
|
{
|
|
side: frame_measurements[(person_idx, side)]
|
|
for side in ("left", "right")
|
|
},
|
|
{
|
|
side: assignment_events.get((person_idx, side))
|
|
for side in ("left", "right")
|
|
},
|
|
short_occlusion_frames,
|
|
medium_occlusion_frames,
|
|
)
|
|
v2_list.append(v2)
|
|
v2_meta_list.append(v2_meta)
|
|
for side in ("left", "right"):
|
|
state_counts[f"{v2_meta[side]['state']}"] += 1
|
|
|
|
global_hands = build_global_hands(selected_list, fusion_meta_list)
|
|
records.append(
|
|
{
|
|
"frame_index": frame_idx,
|
|
"timestamp_ms": round(frame_idx / fps * 1000.0, 3),
|
|
"track_ids": track_ids,
|
|
"boxes": np.round(np.asarray(boxes), 3).tolist() if boxes else [],
|
|
"det_confidences": [round(value, 4) for value in confidences],
|
|
"raw_keypoints": np.round(raw_keypoints, 3).tolist(),
|
|
"temporal_selected_keypoints": np.round(np.asarray(selected_list), 3).tolist(),
|
|
"temporal_keypoints": np.round(np.asarray(smoothed_list), 3).tolist(),
|
|
"temporal_metadata": fusion_meta_list,
|
|
"v2_keypoints": np.round(np.asarray(v2_list), 3).tolist(),
|
|
"v2_metadata": v2_meta_list,
|
|
"global_hands": global_hands,
|
|
}
|
|
)
|
|
|
|
raw_panel = draw_pose(
|
|
frame,
|
|
boxes,
|
|
raw_keypoints,
|
|
f"{video_label} RTMW-L raw | frame {frame_idx}",
|
|
)
|
|
temporal_panel = draw_fused(
|
|
frame,
|
|
boxes,
|
|
smoothed_list,
|
|
fusion_meta_list,
|
|
f"{video_label} temporal fusion | frame {frame_idx}",
|
|
)
|
|
v2_panel = draw_v2(
|
|
frame,
|
|
boxes,
|
|
np.asarray(v2_list),
|
|
v2_meta_list,
|
|
f"{video_label} occlusion V4 | frame {frame_idx}",
|
|
show_reachable=False,
|
|
)
|
|
panels = [
|
|
cv2.resize(panel, (640, 360), interpolation=cv2.INTER_AREA)
|
|
for panel in (raw_panel, temporal_panel, v2_panel)
|
|
]
|
|
writer.write(np.hstack(panels))
|
|
if frame_idx == 0 and video_path.suffix.lower() in {
|
|
".jpg",
|
|
".jpeg",
|
|
".png",
|
|
".bmp",
|
|
".webp",
|
|
}:
|
|
cv2.imwrite(str(output_dir / "raw_pose.jpg"), raw_panel)
|
|
cv2.imwrite(str(output_dir / "spatial_fusion.jpg"), temporal_panel)
|
|
cv2.imwrite(str(output_dir / "occlusion_v2.jpg"), v2_panel)
|
|
panel_width = 960
|
|
panel_height = round(height * panel_width / width)
|
|
image_panels = [
|
|
cv2.resize(
|
|
panel,
|
|
(panel_width, panel_height),
|
|
interpolation=cv2.INTER_AREA,
|
|
)
|
|
for panel in (raw_panel, temporal_panel, v2_panel)
|
|
]
|
|
cv2.imwrite(
|
|
str(output_dir / "raw_spatial_occlusion_v2.jpg"),
|
|
np.hstack(image_panels),
|
|
)
|
|
if frame_idx % 100 == 0:
|
|
print(f"{video_label}: processed frame {frame_idx}", flush=True)
|
|
frame_idx += 1
|
|
|
|
cap.release()
|
|
writer.release()
|
|
with (output_dir / "predictions.jsonl").open("w", encoding="utf-8") as handle:
|
|
for record in records:
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
summary = {
|
|
"occlusion_algorithm": "v4_global_assignment",
|
|
"video": str(video_path.resolve()),
|
|
"frames": len(records),
|
|
"fps": round(fps, 3),
|
|
"resolution": [width, height],
|
|
"person_filter": {
|
|
"center_x_max_ratio": 0.80,
|
|
"minimum_box_height_ratio": 0.20,
|
|
"max_people": max_people,
|
|
},
|
|
"detector_settings": {
|
|
"imgsz": det_imgsz,
|
|
"confidence": det_conf,
|
|
"roi_y_max_ratio": det_roi_y_max_ratio,
|
|
},
|
|
"filtered_person_count_frames": dict(sorted(detection_counts.items())),
|
|
"rejected_detection_counts": dict(sorted(rejection_counts.items())),
|
|
"track_ids_created": box_tracker.next_id - 1,
|
|
"isolated_pose_inferences": isolated_inferences,
|
|
"short_occlusion_frames": short_occlusion_frames,
|
|
"medium_occlusion_frames": medium_occlusion_frames,
|
|
"v2_state_counts": dict(sorted(state_counts.items())),
|
|
"timing_ms_cpu": {
|
|
name: summarize(values[5:] if len(values) > 5 else values)
|
|
for name, values in timings.items()
|
|
},
|
|
"output_note": (
|
|
"Main video hides reachable-region arcs. Reachable-region geometry remains in "
|
|
"predictions.jsonl for later door-distance lower-bound calculation."
|
|
),
|
|
}
|
|
with (output_dir / "summary.json").open("w", encoding="utf-8") as handle:
|
|
json.dump(summary, handle, ensure_ascii=False, indent=2)
|
|
return summary
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
detector = YOLO(str(args.det_model))
|
|
pose = RTMPose(
|
|
str(args.pose_model),
|
|
model_input_size=(288, 384),
|
|
backend="onnxruntime",
|
|
device=args.pose_device,
|
|
)
|
|
all_summaries = {}
|
|
for video_idx, video_path in enumerate(args.videos, start=1):
|
|
label = f"video{video_idx}"
|
|
output_dir = args.output / label
|
|
all_summaries[label] = process_video(
|
|
video_path,
|
|
output_dir,
|
|
detector,
|
|
pose,
|
|
args.det_device,
|
|
args.det_imgsz,
|
|
args.det_conf,
|
|
args.det_roi_y_max_ratio,
|
|
args.short_occlusion_frames,
|
|
args.medium_occlusion_frames,
|
|
args.max_people,
|
|
label,
|
|
)
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
with (args.output / "summary.json").open("w", encoding="utf-8") as handle:
|
|
json.dump(all_summaries, handle, ensure_ascii=False, indent=2)
|
|
print(json.dumps(all_summaries, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|