Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,211 @@
+import io
+import math
+import os
+import tempfile
+import uuid
+import zipfile
+
+import cv2
+import mediapipe as mp
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import StreamingResponse, JSONResponse
+from mediapipe.tasks.python import BaseOptions
+from mediapipe.tasks.python.vision import (
+    PoseLandmarker,
+    PoseLandmarkerOptions,
+    RunningMode,
+)
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+MODEL_PATH = os.path.join(SCRIPT_DIR, "pose_landmarker_heavy.task")
+
+# Landmark indices
+LEFT_SHOULDER = 11
+RIGHT_SHOULDER = 12
+LEFT_HIP = 23
+RIGHT_HIP = 24
+
+TARGETS = {
+    "front": 0,
+    "front_45_clockwise": 45,
+    "right_side": 90,
+    "back_45_clockwise": 135,
+    "back": 180,
+    "back_45_anticlockwise": -135,
+    "left_side": -90,
+    "front_45_anticlockwise": -45,
+}
+
+app = FastAPI(title="Pose Frame Extractor API")
+
+
+def estimate_body_angle(world_landmarks):
+    ls = world_landmarks[LEFT_SHOULDER]
+    rs = world_landmarks[RIGHT_SHOULDER]
+    lh = world_landmarks[LEFT_HIP]
+    rh = world_landmarks[RIGHT_HIP]
+
+    s_dx = ls.x - rs.x
+    s_dz = ls.z - rs.z
+    h_dx = lh.x - rh.x
+    h_dz = lh.z - rh.z
+
+    dx = (s_dx + h_dx) / 2
+    dz = (s_dz + h_dz) / 2
+
+    angle_rad = math.atan2(dz, dx)
+    return math.degrees(angle_rad)
+
+
+def angle_distance(a, b):
+    diff = (a - b + 180) % 360 - 180
+    return abs(diff)
+
+
+def process_video(video_path: str):
+    """Process video and return dict of pose_name -> (png_bytes, metadata)."""
+    if not os.path.exists(MODEL_PATH):
+        raise HTTPException(status_code=500, detail="Pose model file not found on server.")
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise HTTPException(status_code=400, detail="Cannot open uploaded video.")
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    if fps == 0:
+        cap.release()
+        raise HTTPException(status_code=400, detail="Invalid video: 0 FPS detected.")
+
+    options = PoseLandmarkerOptions(
+        base_options=BaseOptions(model_asset_path=MODEL_PATH),
+        running_mode=RunningMode.VIDEO,
+        num_poses=1,
+        min_pose_detection_confidence=0.5,
+        min_pose_presence_confidence=0.5,
+        min_tracking_confidence=0.5,
+    )
+    landmarker = PoseLandmarker.create_from_options(options)
+
+    best = {
+        name: {"diff": float("inf"), "frame": None, "angle": None, "frame_idx": -1}
+        for name in TARGETS
+    }
+
+    frame_idx = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
+        timestamp_ms = int(frame_idx * 1000 / fps)
+        results = landmarker.detect_for_video(mp_image, timestamp_ms)
+
+        if results.pose_world_landmarks and len(results.pose_world_landmarks) > 0:
+            world_lms = results.pose_world_landmarks[0]
+            angle = estimate_body_angle(world_lms)
+
+            for name, target in TARGETS.items():
+                diff = angle_distance(angle, target)
+                if diff < best[name]["diff"]:
+                    best[name] = {
+                        "diff": diff,
+                        "frame": frame.copy(),
+                        "angle": angle,
+                        "frame_idx": frame_idx,
+                    }
+
+        frame_idx += 1
+
+    cap.release()
+    landmarker.close()
+
+    # Encode frames as PNGs
+    results_out = {}
+    for name, target_angle in TARGETS.items():
+        info = best[name]
+        if info["frame"] is None:
+            continue
+        suffix = "" if info["diff"] <= 15 else "_approx"
+        filename = f"{name}{suffix}.png"
+        _, buf = cv2.imencode(".png", info["frame"])
+        results_out[filename] = {
+            "png_bytes": buf.tobytes(),
+            "frame_idx": info["frame_idx"],
+            "detected_angle": round(info["angle"], 1),
+            "target_angle": target_angle,
+            "error": round(info["diff"], 1),
+        }
+
+    return results_out, total_frames, fps
+
+
+@app.post("/extract-poses")
+async def extract_poses(video: UploadFile = File(...)):
+    """Upload a video and get back a ZIP of extracted pose frames."""
+    # Save uploaded video to a temp file
+    suffix = os.path.splitext(video.filename or "video.mp4")[1]
+    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
+        tmp.write(await video.read())
+        tmp_path = tmp.name
+
+    try:
+        results, total_frames, fps = process_video(tmp_path)
+    finally:
+        os.unlink(tmp_path)
+
+    if not results:
+        raise HTTPException(status_code=422, detail="No poses detected in video.")
+
+    # Build a ZIP in memory
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
+        for filename, data in results.items():
+            zf.writestr(filename, data["png_bytes"])
+
+    zip_buffer.seek(0)
+    return StreamingResponse(
+        zip_buffer,
+        media_type="application/zip",
+        headers={"Content-Disposition": "attachment; filename=pose_frames.zip"},
+    )
+
+
+@app.post("/extract-poses-json")
+async def extract_poses_json(video: UploadFile = File(...)):
+    """Upload a video and get back JSON metadata (no images, just info)."""
+    suffix = os.path.splitext(video.filename or "video.mp4")[1]
+    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
+        tmp.write(await video.read())
+        tmp_path = tmp.name
+
+    try:
+        results, total_frames, fps = process_video(tmp_path)
+    finally:
+        os.unlink(tmp_path)
+
+    summary = []
+    for filename, data in results.items():
+        summary.append({
+            "filename": filename,
+            "frame_idx": data["frame_idx"],
+            "detected_angle": data["detected_angle"],
+            "target_angle": data["target_angle"],
+            "error_degrees": data["error"],
+        })
+
+    return JSONResponse({
+        "video_info": {
+            "total_frames": total_frames,
+            "fps": round(fps, 1),
+            "duration_seconds": round(total_frames / fps, 2),
+        },
+        "poses": summary,
+    })
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model_loaded": os.path.exists(MODEL_PATH)}