momentry_core/scripts/pose_processor_v2.py

#!/opt/homebrew/bin/python3.11
"""
Pose Expansion Processor V2

Calls swift_pose_expansion which:
1. Reads face_traced.json (from face tracking with trace_id)
2. Expands pose detection from trace frames
3. Stops when 3 consecutive frames have no pose
4. Outputs at 8Hz sampling (floor(fps/8))

Flow:
    face_processor.py → face.json
    store_traced_faces.py → face_traced.json (with trace_id)
    pose_processor.py → pose.json (this script)
    appearance_processor.py → appearance.json
"""

import sys
import os
import json
import argparse
import subprocess
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from redis_publisher import RedisPublisher

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SWIFT_BIN = os.path.join(SCRIPT_DIR, "swift_processors", ".build", "release", "swift_pose_expansion")
SWIFT_BIN_DEBUG = os.path.join(SCRIPT_DIR, "swift_processors", ".build", "debug", "swift_pose_expansion")

OUTPUT_DIR = os.environ.get("MOMENTRY_OUTPUT_DIR", "/Users/accusys/momentry/output")


def process_pose(
    video_path: str,
    output_path: str,
    uuid: str = "",
    publisher: RedisPublisher = None,
) -> dict:
    """Process pose expansion from face traces.
    
    Args:
        video_path: Path to video file
        output_path: Path to output pose.json
        uuid: File UUID for logging
        publisher: Redis publisher for progress updates
    """
    # Check if pose.json already exists
    if os.path.exists(output_path):
        with open(output_path) as f:
            data = json.load(f)
        frame_count = len(data.get("frames", []))
        print(f"[Pose] Output exists: {output_path} ({frame_count} frames)", file=sys.stderr)
        if publisher:
            publisher.progress("pose", 100, 100, f"{frame_count} frames (exists)")
        return data
    
    # Determine file_uuid from output_path
    file_uuid = os.path.basename(output_path).replace(".pose.json", "")
    
    # Find face_traced.json
    face_traced_path = os.path.join(OUTPUT_DIR, f"{file_uuid}.face_traced.json")
    face_json_path = os.path.join(OUTPUT_DIR, f"{file_uuid}.face.json")
    
    # Prefer face_traced.json (has trace_id), fallback to face.json
    input_face_path = None
    if os.path.exists(face_traced_path):
        input_face_path = face_traced_path
        print(f"[Pose] Using face_traced.json: {face_traced_path}", file=sys.stderr)
    elif os.path.exists(face_json_path):
        # Try to run face tracking to generate face_traced.json
        print(f"[Pose] face_traced.json not found, running face tracker...", file=sys.stderr)
        try:
            tracker_script = os.path.join(SCRIPT_DIR, "store_traced_faces.py")
            if os.path.exists(tracker_script):
                result = subprocess.run(
                    ["python3", tracker_script, "--file-uuid", file_uuid],
                    capture_output=True, text=True, timeout=300
                )
                if result.returncode == 0 and os.path.exists(face_traced_path):
                    input_face_path = face_traced_path
                    print(f"[Pose] Face tracing completed: {face_traced_path}", file=sys.stderr)
                else:
                    print(f"[Pose] Face tracing failed, falling back to face.json", file=sys.stderr)
                    input_face_path = face_json_path
            else:
                input_face_path = face_json_path
        except Exception as e:
            print(f"[Pose] Face tracking error: {e}, falling back to face.json", file=sys.stderr)
            input_face_path = face_json_path
        
        if input_face_path == face_json_path:
            print(f"[Pose] WARNING: Using face.json without trace_id", file=sys.stderr)
    else:
        print(f"[Pose] ERROR: No face.json found for {file_uuid}", file=sys.stderr)
        # Return empty result
        empty_result = {"frame_count": 0, "fps": 0.0, "frames": []}
        with open(output_path, "w") as f:
            json.dump(empty_result, f)
        return empty_result
    
    # Build swift_pose_expansion if needed
    swift_bin = SWIFT_BIN if os.path.exists(SWIFT_BIN) else SWIFT_BIN_DEBUG
    if not os.path.exists(swift_bin):
        build_dir = os.path.join(SCRIPT_DIR, "swift_processors")
        print(f"[Pose] Building swift_pose_expansion in {build_dir}...", file=sys.stderr)
        result = subprocess.run(
            ["swift", "build", "-c", "release", "--product", "swift_pose_expansion"],
            cwd=build_dir, capture_output=True, text=True
        )
        if result.returncode != 0:
            print(f"[Pose] Build failed: {result.stderr}", file=sys.stderr)
            raise RuntimeError("Failed to build swift_pose_expansion")
        swift_bin = SWIFT_BIN
    
    # Run swift_pose_expansion
    cmd = [
        swift_bin,
        video_path,
        input_face_path,
        output_path,
    ]
    if uuid:
        cmd.extend(["--uuid", uuid])
    
    print(f"[Pose] Running: {' '.join(cmd)}", file=sys.stderr)
    
    t0 = time.time()
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    
    # Monitor progress
    last_progress = ""
    while proc.poll() is None:
        time.sleep(5)
        # Read stderr for progress
        try:
            # Non-blocking read
            import select
            if select.select([proc.stderr], [], [], 0)[0]:
                line = proc.stderr.readline().strip()
                if line and line != last_progress:
                    last_progress = line
                    print(f"[Pose] {line}", file=sys.stderr)
                    if publisher and "frames" in line:
                        publisher.progress("pose", 50, 100, line)
        except Exception:
            pass
    
    # Read remaining output
    stdout, stderr = proc.communicate()
    if stdout:
        print(stdout, file=sys.stderr)
    if stderr:
        print(stderr, file=sys.stderr)
    
    elapsed = time.time() - t0
    
    if proc.returncode != 0:
        print(f"[Pose] ERROR: swift_pose_expansion exited with code {proc.returncode}", file=sys.stderr)
        if publisher:
            publisher.error("pose", f"Process failed with code {proc.returncode}")
        raise RuntimeError(f"swift_pose_expansion failed: {proc.returncode}")
    
    # Load result
    if not os.path.exists(output_path):
        print(f"[Pose] ERROR: Output file not created: {output_path}", file=sys.stderr)
        raise RuntimeError("Pose output not created")
    
    with open(output_path) as f:
        result = json.load(f)
    
    frame_count = len(result.get("frames", []))
    print(f"[Pose] Done: {frame_count} frames in {elapsed:.1f}s", file=sys.stderr)
    
    if publisher:
        publisher.progress("pose", 100, 100, f"{frame_count} frames")
        publisher.complete("pose", f"{frame_count} frames")
    
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Pose Expansion Processor")
    parser.add_argument("video_path", help="Video file path")
    parser.add_argument("output_path", help="Output pose.json path")
    parser.add_argument("--uuid", "-u", default="", help="File UUID for logging")
    args = parser.parse_args()
    
    publisher = RedisPublisher(args.uuid) if args.uuid else None
    if publisher:
        publisher.info("pose", "POSE_START")
    
    result = process_pose(args.video_path, args.output_path, args.uuid, publisher)
    print(f"Pose: {len(result.get('frames', []))} frames with poses")