app.py

import os
import json
import math
import tempfile
from pathlib import Path
from typing import Dict, Tuple

import cv2
import numpy as np
import mediapipe as mp
import gradio as gr

# ------------------------------
# Configuration
# ------------------------------
REFERENCE_POSES_FILE = "reference_poses.json"
DEFAULT_TOLERANCE = 15.0

mp_pose = mp.solutions.pose
LANDMARK = mp_pose.PoseLandmark
JOINT_TRIPLETS = {
    "left_elbow": (LANDMARK.LEFT_SHOULDER, LANDMARK.LEFT_ELBOW, LANDMARK.LEFT_WRIST),
    "right_elbow": (LANDMARK.RIGHT_SHOULDER, LANDMARK.RIGHT_ELBOW, LANDMARK.RIGHT_WRIST),
    "left_shoulder": (LANDMARK.LEFT_HIP, LANDMARK.LEFT_SHOULDER, LANDMARK.LEFT_ELBOW),
    "right_shoulder": (LANDMARK.RIGHT_HIP, LANDMARK.RIGHT_SHOULDER, LANDMARK.RIGHT_ELBOW),
    "left_knee": (LANDMARK.LEFT_HIP, LANDMARK.LEFT_KNEE, LANDMARK.LEFT_ANKLE),
    "right_knee": (LANDMARK.RIGHT_HIP, LANDMARK.RIGHT_KNEE, LANDMARK.RIGHT_ANKLE),
    "left_hip": (LANDMARK.LEFT_SHOULDER, LANDMARK.LEFT_HIP, LANDMARK.LEFT_KNEE),
    "right_hip": (LANDMARK.RIGHT_SHOULDER, LANDMARK.RIGHT_HIP, LANDMARK.RIGHT_KNEE),
}

# ------------------------------
# Utility functions
# ------------------------------
def load_reference_poses(path: str = REFERENCE_POSES_FILE) -> Dict:
    if not os.path.exists(path):
        default = {
            "Warrior II": {
                "left_elbow": 170, "right_elbow": 170,
                "left_shoulder": 90, "right_shoulder": 90,
                "left_knee": 90, "right_knee": 170,
                "left_hip": 170, "right_hip": 170
            },
            "Tree": {
                "left_elbow": 170, "right_elbow": 170,
                "left_shoulder": 120, "right_shoulder": 120,
                "left_knee": 170, "right_knee": 40,
                "left_hip": 170, "right_hip": 40
            },
            "Downward Dog": {
                "left_elbow": 170, "right_elbow": 170,
                "left_shoulder": 70, "right_shoulder": 70,
                "left_knee": 170, "right_knee": 170,
                "left_hip": 160, "right_hip": 160
            }
        }
        with open(path, "w") as f:
            json.dump(default, f, indent=2)
        return default
    with open(path, "r") as f:
        return json.load(f)


def vector(a, b):
    return np.array([b[0] - a[0], b[1] - a[1]])


def angle_between_points(a, b, c):
    v1 = vector(b, a)
    v2 = vector(b, c)
    dot = v1.dot(v2)
    norm = (np.linalg.norm(v1) * np.linalg.norm(v2)) + 1e-8
    cosang = np.clip(dot / norm, -1.0, 1.0)
    return math.degrees(math.acos(cosang))


def landmarks_to_xy(landmark_list, width, height):
    coords = {}
    for idx, lm in enumerate(landmark_list.landmark):
        coords[idx] = (lm.x * width, lm.y * height, lm.visibility)
    return coords


def compute_joint_angles(landmarks_xy: Dict[int, Tuple[float, float, float]]) -> Dict[str, float]:
    angles = {}
    for name, (p_idx, j_idx, c_idx) in JOINT_TRIPLETS.items():
        try:
            pa, jb, ca = landmarks_xy[p_idx], landmarks_xy[j_idx], landmarks_xy[c_idx]
            if pa[2] < 0.3 or jb[2] < 0.3 or ca[2] < 0.3:
                angles[name] = None
            else:
                angles[name] = angle_between_points((pa[0], pa[1]), (jb[0], jb[1]), (ca[0], ca[1]))
        except KeyError:
            angles[name] = None
    return angles


def compare_angles(detected, reference, tolerance=DEFAULT_TOLERANCE):
    per_joint_score, per_joint_diff = {}, {}
    for joint, ref_ang in reference.items():
        det_ang = detected.get(joint)
        if det_ang is None:
            per_joint_score[joint] = None
            per_joint_diff[joint] = None
        else:
            diff = det_ang - ref_ang
            per_joint_diff[joint] = diff
            score = max(0.0, 100.0 * (1 - (abs(diff) / (2 * tolerance))))
            per_joint_score[joint] = float(np.clip(score, 0.0, 100.0))
    valid = [v for v in per_joint_score.values() if v is not None]
    overall = float(np.mean(valid)) if valid else 0.0
    return overall, per_joint_score, per_joint_diff


# ------------------------------
# Video processing
# ------------------------------
def process_video(input_path: str, pose_name: str, tolerance: float = DEFAULT_TOLERANCE):
    ref_poses = load_reference_poses()
    if pose_name not in ref_poses:
        return None, f"Pose '{pose_name}' not found."

    reference = ref_poses[pose_name]
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        return None, "Failed to open uploaded video."

    fps = cap.get(cv2.CAP_PROP_FPS) or 20.0
    width, height = int(cap.get(3)), int(cap.get(4))
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out_path = os.path.join(tempfile.gettempdir(), f"annotated_{Path(input_path).stem}.mp4")
    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))

    pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5, min_tracking_confidence=0.5)
    aggregate_scores = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = pose.process(image_rgb)
        annotated = frame.copy()

        if results.pose_landmarks:
            lm_xy = landmarks_to_xy(results.pose_landmarks, width, height)
            detected_angles = compute_joint_angles(lm_xy)
            percent, per_joint_score, per_joint_diff = compare_angles(detected_angles, reference, tolerance)
            aggregate_scores.append(percent)

            for joint, (p, j, c) in JOINT_TRIPLETS.items():
                if j in lm_xy and p in lm_xy:
                    color = (0, 255, 0)
                    if per_joint_score[joint] is not None:
                        if per_joint_score[joint] < 33:
                            color = (0, 0, 255)
                        elif per_joint_score[joint] < 66:
                            color = (0, 165, 255)
                    cv2.line(annotated, (int(lm_xy[p][0]), int(lm_xy[p][1])),
                             (int(lm_xy[j][0]), int(lm_xy[j][1])), color, 3)
                if j in lm_xy and c in lm_xy:
                    color = (0, 255, 0)
                    if per_joint_score[joint] is not None:
                        if per_joint_score[joint] < 33:
                            color = (0, 0, 255)
                        elif per_joint_score[joint] < 66:
                            color = (0, 165, 255)
                    cv2.line(annotated, (int(lm_xy[j][0]), int(lm_xy[j][1])),
                             (int(lm_xy[c][0]), int(lm_xy[c][1])), color, 3)

            cv2.putText(annotated, f"{pose_name}: {percent:.0f}%", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
        else:
            cv2.putText(annotated, "No pose detected", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

        out.write(annotated)

    cap.release()
    out.release()
    pose.close()

    avg_score = float(np.mean(aggregate_scores)) if aggregate_scores else 0.0
    return out_path, {
        "pose": pose_name,
        "score_percent": avg_score,
        "suggestions": [
            f"Try maintaining stability. Overall correctness: {avg_score:.1f}%."
        ]
    }


# ------------------------------
# Gradio Interface
# ------------------------------
ref_poses = load_reference_poses()
pose_list = list(ref_poses.keys())

with gr.Blocks(title="Yoga Pose Correctness Checker") as demo:
    gr.Markdown("""
    # 🧘 Yoga Pose Correctness Checker
    Upload a short video of your yoga pose.  
    The app will analyze:
    - ✅ Pose correctness percentage  
    - 📊 Joint-by-joint feedback  
    - 💡 Suggestions for improvement
    """)

    video_in = gr.Video(label="Upload a video (MP4/MOV)")
    pose_dropdown = gr.Dropdown(choices=pose_list, value=pose_list[0], label="Select Pose")
    tol_slider = gr.Slider(5, 40, value=DEFAULT_TOLERANCE, step=1, label="Tolerance (degrees)")
    run_btn = gr.Button("Analyze Pose")
    output_video = gr.Video(label="Annotated Video Output")
    output_json = gr.JSON(label="Results and Suggestions")

    def analyze(video_path, pose_name, tolerance):
        if not video_path:
            return None, {"error": "Please upload a video first."}
        annotated_path, result = process_video(video_path, pose_name, tolerance)
        if annotated_path is None:
            return None, {"error": result}
        return annotated_path, result

    run_btn.click(analyze, inputs=[video_in, pose_dropdown, tol_slider], outputs=[output_video, output_json])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))