Accusys
39a2cbc65b
- get_face_groups_handler: COALESCE(tp.name, tn.label) for name consistency - sync_file_status: compare JSON vs pre_chunks (not chunk table) - face consistency: compare frames.len() not total_faces - cleanup 2 ghost records with NULL file_name/file_path - replace identity_agent with face_dedup in pipeline stages - remove identity_agent_api.rs and all references - update required_processors to match actual processors - update AGENTS.md with team responsibilities - add Studio pipeline changes documentation
235 lines
7.1 KiB
Python
235 lines
7.1 KiB
Python
#!/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()) |