#!/opt/homebrew/bin/python3.11 """ Assign trace_ids to poses by matching with face traces. Uses bbox IoU matching to find corresponding face traces. Input: face_traced.json, pose.json Output: pose_traced.json """ import json import os import argparse from typing import Dict, List, Optional, Any def calculate_iou(bbox1: Dict, bbox2: Dict) -> float: """Calculate Intersection over Union for two bboxes.""" x1 = max(bbox1["x"], bbox2["x"]) y1 = max(bbox1["y"], bbox2["y"]) x2 = min(bbox1["x"] + bbox1["width"], bbox2["x"] + bbox2["width"]) y2 = min(bbox1["y"] + bbox1["height"], bbox2["y"] + bbox2["height"]) if x2 <= x1 or y2 <= y1: return 0.0 intersection = (x2 - x1) * (y2 - y1) area1 = bbox1["width"] * bbox1["height"] area2 = bbox2["width"] * bbox2["height"] union = area1 + area2 - intersection return intersection / union if union > 0 else 0.0 def is_face_center_in_pose(face_bbox: Dict, pose_bbox: Dict) -> bool: """Check if face center is within pose bbox.""" face_cx = face_bbox["x"] + face_bbox["width"] // 2 face_cy = face_bbox["y"] + face_bbox["height"] // 2 return ( pose_bbox["x"] <= face_cx <= pose_bbox["x"] + pose_bbox["width"] and pose_bbox["y"] <= face_cy <= pose_bbox["y"] + pose_bbox["height"] ) def build_face_lookup(face_traced: Dict) -> Dict[int, List[Dict]]: """Build frame -> faces lookup from face_traced.json.""" lookup = {} for trace_id_str, trace in face_traced.get("traces", {}).items(): trace_id = int(trace_id_str) for face in trace.get("path", []): frame = face["frame"] if frame not in lookup: lookup[frame] = [] lookup[frame].append({ "trace_id": trace_id, "bbox": face["bbox"], "confidence": face.get("confidence", 0.5) }) return lookup def find_closest_faces( face_lookup: Dict[int, List[Dict]], target_frame: int, max_distance: int = 10 ) -> List[Dict]: """ Find faces at the closest frame to target_frame. Search within max_distance frames. """ # Check exact frame first if target_frame in face_lookup: return face_lookup[target_frame] # Find closest frame face_frames = sorted(face_lookup.keys()) closest_frame = None closest_distance = max_distance + 1 for frame in face_frames: distance = abs(frame - target_frame) if distance < closest_distance: closest_distance = distance closest_frame = frame if closest_frame is not None and closest_distance <= max_distance: return face_lookup[closest_frame] return [] def match_pose_to_traces( pose_person: Dict, faces_at_frame: List[Dict], frame: int, iou_threshold: float = 0.05 ) -> Dict: """ Match a pose person to face traces. Uses two strategies: 1. IoU matching (lower threshold for body vs face) 2. Face center containment (face center within pose bbox) """ matched_traces = [] for face in faces_at_frame: iou = calculate_iou(pose_person["bbox"], face["bbox"]) # Strategy 1: IoU matching (lower threshold) if iou > iou_threshold: matched_traces.append({ "trace_id": face["trace_id"], "iou": iou, "method": "iou" }) # Strategy 2: Face center in pose bbox elif is_face_center_in_pose(face["bbox"], pose_person["bbox"]): matched_traces.append({ "trace_id": face["trace_id"], "iou": iou, "method": "center_containment" }) # Sort by IoU descending matched_traces.sort(key=lambda x: x["iou"], reverse=True) # Assign trace_ids trace_ids = [t["trace_id"] for t in matched_traces] # Generate pose_id using first trace_id if trace_ids: pose_id = f"pose_{trace_ids[0]}_{frame}" else: pose_id = f"pose_none_{frame}" # Update pose person pose_person["pose_id"] = pose_id pose_person["trace_ids"] = trace_ids return pose_person def assign_pose_traces( face_traced_path: str, pose_path: str, output_path: str, iou_threshold: float = 0.3 ) -> Dict: """ Main function: assign trace_ids to poses. """ # Load face_traced.json print(f"[PoseTrace] Loading face_traced.json: {face_traced_path}") with open(face_traced_path) as f: face_traced = json.load(f) # Load pose.json print(f"[PoseTrace] Loading pose.json: {pose_path}") with open(pose_path) as f: pose_data = json.load(f) # Build face lookup face_lookup = build_face_lookup(face_traced) print(f"[PoseTrace] Built face lookup: {len(face_lookup)} frames with faces") # Process each frame total_poses = 0 matched_poses = 0 for frame_data in pose_data.get("frames", []): frame = frame_data["frame"] # Find closest faces (within 10 frames) faces_at_frame = find_closest_faces(face_lookup, frame, max_distance=10) for person in frame_data.get("persons", []): total_poses += 1 # Match pose to traces matched_person = match_pose_to_traces( person, faces_at_frame, frame, iou_threshold ) if matched_person.get("trace_ids"): matched_poses += 1 print(f"[PoseTrace] Matched {matched_poses}/{total_poses} poses to traces") # Update metadata pose_data["trace_matching"] = { "total_poses": total_poses, "matched_poses": matched_poses, "iou_threshold": iou_threshold } # Save pose_traced.json print(f"[PoseTrace] Saving to: {output_path}") with open(output_path, "w") as f: json.dump(pose_data, f, indent=2, ensure_ascii=False) return pose_data def main(): parser = argparse.ArgumentParser(description="Assign trace_ids to poses") parser.add_argument("--uuid", required=True, help="Video file UUID") parser.add_argument("--iou-threshold", type=float, default=0.3, help="IoU threshold for matching") parser.add_argument("--output-dir", help="Output directory (default: from env)") args = parser.parse_args() output_dir = args.output_dir or os.environ.get("MOMENTRY_OUTPUT_DIR", "/Users/accusys/momentry/output") face_traced_path = os.path.join(output_dir, f"{args.uuid}.face_traced.json") pose_path = os.path.join(output_dir, f"{args.uuid}.pose.json") output_path = os.path.join(output_dir, f"{args.uuid}.pose_traced.json") # Check input files exist if not os.path.exists(face_traced_path): print(f"[PoseTrace] Error: face_traced.json not found: {face_traced_path}") return 1 if not os.path.exists(pose_path): print(f"[PoseTrace] Error: pose.json not found: {pose_path}") return 1 # Run matching assign_pose_traces(face_traced_path, pose_path, output_path, args.iou_threshold) return 0 if __name__ == "__main__": exit(main())