"""Frame-extraction service — HF Space CPU worker. GET /video-meta ?key= Probes video without full download: S3 HeadObject (size) + ffprobe (duration, fps, resolution). Returns { size_bytes, duration_s, native_fps, width, height }. The TS video_planner node calls this to decide target_fps before the GPU job. POST /extract-frames { assessment_id, video_s3_key?, frames?: [int,...], target_fps?: float } Seeks to each requested frame (or all if omitted), uploads raw JPEGs. If target_fps < native_fps, resamples the video first so frame indices match track.json. Fast path for interrupt-1: only the ~15 excluded frames. POST /render { assessment_id, video_s3_key?, track_json_s3_key, bundle_s3_keys, target_fps?: float } Decodes all frames, draws skeleton overlay (green=qualified / red=excluded), stitches debug.mp4, uploads everything. For the results view. GET /health → {"ok": true} Self-contained — zero ergo_agent imports. S3 and cv2 are injected at startup. """ from __future__ import annotations import json import logging import math import os import shutil import subprocess import tempfile import time from concurrent.futures import ThreadPoolExecutor, as_completed from typing import List, Optional import boto3 from fastapi import FastAPI, Query from pydantic import BaseModel logging.basicConfig( level=os.environ.get("LOG_LEVEL", "INFO"), format="%(asctime)s %(levelname)s [%(name)s] %(message)s", ) logger = logging.getLogger("render_worker") # --------------------------------------------------------------------------- # Config (own — no ergo_agent.config) # --------------------------------------------------------------------------- BUCKET: str = os.environ["S3_BUCKET"] _PREFIX: str = os.environ.get("S3_ASSESSMENT_PREFIX", "industrial/assessments").rstrip("/") def _make_s3(): return boto3.client( "s3", endpoint_url=os.environ.get("S3_ENDPOINT_URL") or None, region_name=os.environ.get("S3_REGION") or None, aws_access_key_id=os.environ["S3_ACCESS_KEY_ID"], aws_secret_access_key=os.environ["S3_SECRET_ACCESS_KEY"], ) _s3 = None # initialised lazily on first request def _get_s3(): global _s3 if _s3 is None: _s3 = _make_s3() return _s3 def _assessment_prefix(assessment_id: str) -> str: return f"{_PREFIX}/{assessment_id}" def _video_key(assessment_id: str, explicit: Optional[str]) -> str: return explicit or f"{_assessment_prefix(assessment_id)}/video/input.mp4" # --------------------------------------------------------------------------- # Adaptive FPS (spec: docs/VIDEO_SAMPLING_SPEC.md) # --------------------------------------------------------------------------- def adaptive_fps(duration_s: float, native_fps: float, size_bytes: int) -> float: """Return the target FPS for processing based on video characteristics. Never upsamples. Floor at 5fps (below this REBA temporal reasoning breaks). """ # Duration-based table if duration_s <= 30: target = min(native_fps, 30.0) elif duration_s <= 60: target = 15.0 elif duration_s <= 120: target = 10.0 else: target = 8.0 # Large-file penalty: >500MB AND duration >60s → reduce by 2 more if size_bytes > 500_000_000 and duration_s > 60: target = target - 2.0 # Hard constraints target = max(target, 5.0) # floor target = min(target, native_fps) # never upsample return float(target) # --------------------------------------------------------------------------- # Video resampling (ensures frame indices match track.json) # --------------------------------------------------------------------------- def _resample_video(src: str, target_fps: float, out_dir: str) -> str: """Re-encode video at target_fps. Returns path to resampled file.""" out_path = os.path.join(out_dir, "resampled.mp4") cmd = [ "ffmpeg", "-y", "-loglevel", "error", "-i", src, "-vf", f"fps={target_fps:.6g}", "-an", "-c:v", "libx264", "-preset", "veryfast", "-crf", "20", "-pix_fmt", "yuv420p", out_path, ] proc = subprocess.run(cmd, capture_output=True, text=True) if proc.returncode != 0: raise RuntimeError( f"ffmpeg resample failed (exit {proc.returncode}): " f"{proc.stderr.strip() or ''}" ) return out_path def _transcode_h264(src: str, dst: str) -> None: """Re-encode src (any codec) to H.264 in a browser-compatible MP4.""" cmd = [ "ffmpeg", "-y", "-loglevel", "error", "-i", src, "-c:v", "libx264", "-preset", "veryfast", "-crf", "23", "-pix_fmt", "yuv420p", # required for Safari + broad browser compat "-movflags", "+faststart", # moov atom at front — playback starts immediately "-an", dst, ] proc = subprocess.run(cmd, capture_output=True, text=True) if proc.returncode != 0: raise RuntimeError( f"ffmpeg H.264 transcode failed (exit {proc.returncode}): " f"{proc.stderr.strip() or ''}" ) def _probe_video(local_path: str) -> dict: """Run ffprobe on a local file; returns {duration_s, native_fps, width, height}.""" cmd = [ "ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", "-show_format", local_path, ] proc = subprocess.run(cmd, capture_output=True, text=True) if proc.returncode != 0: raise RuntimeError(f"ffprobe failed: {proc.stderr.strip()}") info = json.loads(proc.stdout) stream = next((s for s in info.get("streams", []) if "width" in s), {}) duration_s = float(info.get("format", {}).get("duration", 0)) r_fps = stream.get("r_frame_rate", "0/1") num, den = (int(x) for x in r_fps.split("/")) native_fps = num / den if den else 0.0 return { "duration_s": duration_s, "native_fps": native_fps, "width": stream.get("width", 0), "height": stream.get("height", 0), } # --------------------------------------------------------------------------- # Skeleton — per-edge severity coloring # --------------------------------------------------------------------------- # Each edge: (joint_a, joint_b, part_name) # part_name maps to rec["parts"][part]["score"] in L1. # Joints: NECK=69, NOSE=0, L_SH=5, R_SH=6, L_EL=7, R_EL=8, # L_HIP=9, R_HIP=10, L_KNEE=11, R_KNEE=12, L_ANK=13, R_ANK=14, # L_WRIST=62, R_WRIST=41 _EDGES = [ (69, 5, "trunk"), # neck → left shoulder (69, 6, "trunk"), # neck → right shoulder ( 5, 7, "upper_arm"), # left shoulder → left elbow ( 7, 62, "lower_arm"), # left elbow → left wrist ( 6, 8, "upper_arm"), # right shoulder → right elbow ( 8, 41, "lower_arm"), # right elbow → right wrist ( 5, 9, "trunk"), # left shoulder → left hip ( 6, 10, "trunk"), # right shoulder → right hip ( 9, 10, "trunk"), # left hip → right hip ( 9, 11, "legs"), # left hip → left knee (11, 13, "legs"), # left knee → left ankle (10, 12, "legs"), # right hip → right knee (12, 14, "legs"), # right knee → right ankle ( 0, 69, "neck"), # nose → neck ] # BGR — score thresholds match REBA risk bands: negligible=1, low=2-3, medium=4-7, high=8-10, very_high=11+ _SCORE_COLOR = [ (1, (0, 200, 0)), # 1 negligible — green (3, (0, 200, 160)), # 2–3 low — teal (7, (0, 180, 255)), # 4–7 medium — amber (10, (0, 80, 255)), # 8–10 high — orange (99, (40, 40, 220)), # 11+ very_high — red ] _EXCLUDED_COLOR = (80, 80, 80) # grey — disqualified frame, don't read score def _part_color(score: int) -> tuple: for threshold, color in _SCORE_COLOR: if score <= threshold: return color return _SCORE_COLOR[-1][1] def overlay_segments( q: list, qv: list, w: int, h: int, parts: dict, vis_floor: float = 0.3 ) -> list: """Pure: normalised 2D joints → colored pixel line segments. Each edge is colored by the REBA score of its body part. Returns [((x1,y1),(x2,y2), color), ...] in pixel coords. """ segs = [] for a, b, part in _EDGES: if a >= len(q) or b >= len(q): continue if qv[a] < vis_floor or qv[b] < vis_floor: continue pa = (int(q[a][0] * w), int(q[a][1] * h)) pb = (int(q[b][0] * w), int(q[b][1] * h)) score = (parts.get(part) or {}).get("score") or 1 segs.append((pa, pb, _part_color(score))) return segs # --------------------------------------------------------------------------- # S3 helpers # --------------------------------------------------------------------------- def _put_frame(rbase: str, filename: str, jpg_bytes: bytes) -> None: _get_s3().put_object( Bucket=BUCKET, Key=f"{rbase}/frames/{filename}", Body=jpg_bytes, ContentType="image/jpeg", ) def _put_video(rbase: str, local_path: str) -> None: if not os.path.exists(local_path): return _get_s3().put_object( Bucket=BUCKET, Key=f"{rbase}/debug.mp4", Body=open(local_path, "rb").read(), ContentType="video/mp4", ) def _open_video(out_dir: str, s3_key: str, target_fps: Optional[float]): """Download video from S3, optionally resample, return (cv2_cap, vid_path).""" import cv2 # lazy vid_path = os.path.join(out_dir, "v.mp4") _get_s3().download_file(BUCKET, s3_key, vid_path) if target_fps is not None: # probe native fps to decide whether resampling is actually needed probe = _probe_video(vid_path) if target_fps < probe["native_fps"] - 0.01: vid_path = _resample_video(vid_path, target_fps, out_dir) return cv2.VideoCapture(vid_path), vid_path # --------------------------------------------------------------------------- # Pydantic models # --------------------------------------------------------------------------- class ExtractReq(BaseModel): assessment_id: str video_s3_key: Optional[str] = None frames: Optional[List[int]] = None # None → all frames target_fps: Optional[float] = None # None → use native fps as-is class RenderReq(BaseModel): assessment_id: str video_s3_key: Optional[str] = None track_json_s3_key: str bundle_s3_keys: dict target_fps: Optional[float] = None # --------------------------------------------------------------------------- # App # --------------------------------------------------------------------------- app = FastAPI() @app.get("/health") def health() -> dict: return {"ok": True} @app.get("/video-meta") def video_meta(key: str = Query(..., description="S3 key of the video")) -> dict: """Probe video metadata without full download. S3 HeadObject → size_bytes. Downloads a small header via ffprobe for duration/fps/resolution. The TS video_planner node calls this once before the GPU job to decide target_fps. """ logger.info("video-meta: probing key=%s", key) t0 = time.monotonic() head = _get_s3().head_object(Bucket=BUCKET, Key=key) size_bytes = head["ContentLength"] # Download just enough for ffprobe: stream from S3 to a temp file # (ffprobe only needs the container header, but boto3 download_file # fetches the whole thing — acceptable for metadata since it's async/cheap) out = tempfile.mkdtemp(prefix="meta_") try: vid_path = os.path.join(out, "v.mp4") _get_s3().download_file(BUCKET, key, vid_path) probe = _probe_video(vid_path) finally: shutil.rmtree(out, ignore_errors=True) result = { "size_bytes": size_bytes, "duration_s": probe["duration_s"], "native_fps": probe["native_fps"], "width": probe["width"], "height": probe["height"], } logger.info( "video-meta: done in %.1fs — %.1fs @ %.2ffps %dx%d (%.1fMB)", time.monotonic() - t0, probe["duration_s"], probe["native_fps"], probe["width"], probe["height"], size_bytes / 1e6, ) return result def _extract_specific(cv2, cap, pool, rbase: str, wanted: list) -> int: """Seek to each requested frame index and upload. Returns count uploaded.""" futures = [] n = 0 for idx in wanted: cap.set(cv2.CAP_PROP_POS_FRAMES, idx) ok, frame = cap.read() if not ok: logger.warning("extract-frames: seek to frame %d failed", idx) continue ok2, jpg = cv2.imencode(".jpg", frame) if ok2: futures.append(pool.submit(_put_frame, rbase, f"f{idx:04d}.jpg", jpg.tobytes())) n += 1 for fut in as_completed(futures): fut.result() return n def _extract_all(cv2, cap, pool, rbase: str) -> int: """Read every frame sequentially and upload. Returns count uploaded.""" futures = [] n = 0 i = 0 while True: ok, frame = cap.read() if not ok: break ok2, jpg = cv2.imencode(".jpg", frame) if ok2: futures.append(pool.submit(_put_frame, rbase, f"f{i:04d}.jpg", jpg.tobytes())) n += 1 i += 1 for fut in as_completed(futures): fut.result() return n @app.post("/extract-frames") def do_extract(req: ExtractReq) -> dict: rbase = f"{_assessment_prefix(req.assessment_id)}/render" n_requested = len(req.frames) if req.frames is not None else None logger.info( "extract-frames: assessment=%s frames=%s target_fps=%s", req.assessment_id, f"{n_requested} specific" if n_requested is not None else "all", req.target_fps, ) t0 = time.monotonic() out = tempfile.mkdtemp(prefix="extract_") try: t_dl = time.monotonic() cap, _ = _open_video(out, _video_key(req.assessment_id, req.video_s3_key), req.target_fps) logger.info("extract-frames: video ready in %.1fs", time.monotonic() - t_dl) import cv2 # lazy — stubbed in tests wanted = sorted(set(req.frames)) if req.frames is not None else None try: with ThreadPoolExecutor(max_workers=8) as pool: n = ( _extract_specific(cv2, cap, pool, rbase, wanted) if wanted is not None else _extract_all(cv2, cap, pool, rbase) ) finally: cap.release() finally: shutil.rmtree(out, ignore_errors=True) logger.info( "extract-frames: done — %d frames uploaded to %s in %.1fs", n, rbase, time.monotonic() - t0, ) return {"render_s3_prefix": rbase, "n_frames": n} def _overlay_frame(cv2, frame: object, i: int, track_frames: list, l1: list, w: int, h: int) -> int: """Draw per-part severity skeleton + frame score label. Returns segments drawn.""" tf = track_frames[i] if i < len(track_frames) else {} rec = l1[i] if i < len(l1) else {} q = tf.get("q") or [] qv = tf.get("qv") or [] if not (q and qv): return 0 qualified = rec.get("qualification", {}).get("qualified", True) parts = rec.get("parts") or {} final = rec.get("final") if not qualified: # Excluded frame: draw everything grey so it's visible but clearly flagged segs = overlay_segments(q, qv, w, h, {}) for pa, pb, _ in segs: cv2.line(frame, pa, pb, _EXCLUDED_COLOR, 2) label = f"f{i:04d} [excl]" cv2.putText(frame, label, (8, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.65, _EXCLUDED_COLOR, 2) return len(segs) segs = overlay_segments(q, qv, w, h, parts) for pa, pb, color in segs: cv2.line(frame, pa, pb, color, 2) # Frame label: index + overall REBA score colored by frame risk frame_score = final or 1 label_color = _part_color(frame_score) risk = rec.get("risk", "") label = f"f{i:04d} REBA {frame_score} {risk}" cv2.putText(frame, label, (8, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.65, label_color, 2) return len(segs) @app.post("/render") def do_render(req: RenderReq) -> dict: rbase = f"{_assessment_prefix(req.assessment_id)}/render" logger.info( "render: assessment=%s track=%s target_fps=%s", req.assessment_id, req.track_json_s3_key, req.target_fps, ) t0 = time.monotonic() out = tempfile.mkdtemp(prefix="render_") try: t_dl = time.monotonic() cap, _ = _open_video(out, _video_key(req.assessment_id, req.video_s3_key), req.target_fps) import cv2 # lazy — stubbed in tests track_raw = _get_s3().get_object(Bucket=BUCKET, Key=req.track_json_s3_key)["Body"].read() track = json.loads(track_raw) track_frames = track.get("frames") or [] l1_raw = _get_s3().get_object(Bucket=BUCKET, Key=req.bundle_s3_keys["l1"])["Body"].read() l1 = [json.loads(line) for line in l1_raw.decode().splitlines() if line.strip()] logger.info( "render: data loaded in %.1fs — %d track frames, %d l1 records", time.monotonic() - t_dl, len(track_frames), len(l1), ) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) or 15.0 logger.info("render: video %dx%d @ %.2ffps", w, h, fps) # Write frames into a raw mp4v intermediate — OpenCV's avc1 is unavailable # on most Linux builds. We transcode to H.264 with ffmpeg afterward so # browsers can play the result in a