app.py

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)}