Create app.py

app.py

@@ -0,0 +1,317 @@
+# app.py
+import os
+import json
+import math
+import tempfile
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import cv2
+import numpy as np
+import mediapipe as mp
+import gradio as gr
+
+# --- Config / reference poses (angles in degrees) ---
+REFERENCE_POSES_FILE = "reference_poses.json"
+
+# Mediapipe utils
+mp_pose = mp.solutions.pose
+mp_drawing = mp.solutions.drawing_utils
+
+# Useful landmark indices from MediaPipe Pose
+LANDMARK = mp_pose.PoseLandmark
+# Example joints we will compute angles for (triplet: parent, joint, child)
+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),
+}
+
+# thresholds (degrees) for "correct" per joint
+DEFAULT_TOLERANCE = 15.0
+
+# --- Helper functions ---
+def load_reference_poses(path: str = REFERENCE_POSES_FILE) -> Dict:
+    if not os.path.exists(path):
+        # create a default one if missing
+        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: Tuple[float, float], b: Tuple[float, float]) -> np.ndarray:
+    return np.array([b[0]-a[0], b[1]-a[1]])
+
+def angle_between_points(a, b, c) -> float:
+    """
+    Returns the angle ABC (in degrees) formed at point b by points a-b-c.
+    Points are (x, y).
+    """
+    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)
+    ang = math.degrees(math.acos(cosang))
+    return ang
+
+def landmarks_to_xy(landmark_list, image_width, image_height):
+    coords = {}
+    for idx, lm in enumerate(landmark_list.landmark):
+        coords[idx] = (lm.x * image_width, lm.y * image_height, lm.visibility if hasattr(lm, "visibility") else 1.0)
+    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 = landmarks_xy[p_idx]
+            jb = landmarks_xy[j_idx]
+            ca = landmarks_xy[c_idx]
+            # ignore if visibility low (z could be used too)
+            if pa[2] < 0.3 or jb[2] < 0.3 or ca[2] < 0.3:
+                angles[name] = None
+            else:
+                ang = angle_between_points((pa[0], pa[1]), (jb[0], jb[1]), (ca[0], ca[1]))
+                angles[name] = ang
+        except KeyError:
+            angles[name] = None
+    return angles
+
+def compare_angles(detected: Dict[str, float], reference: Dict[str, float], 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 = abs(det_ang - ref_ang)
+            per_joint_diff[joint] = det_ang - ref_ang
+            # score: linear falloff: diff 0 -> 100, diff >= 2*tolerance -> 0
+            score = max(0.0, 100.0 * (1 - (diff / (2 * tolerance))))
+            per_joint_score[joint] = float(np.clip(score, 0.0, 100.0))
+    # final percent: average of available joint scores
+    valid_scores = [v for v in per_joint_score.values() if v is not None]
+    final_percent = float(np.mean(valid_scores)) if valid_scores else 0.0
+    return final_percent, per_joint_score, per_joint_diff
+
+def suggest_corrections(per_joint_diff: Dict[str, float], tol=DEFAULT_TOLERANCE) -> List[str]:
+    suggestions = []
+    for joint, diff in per_joint_diff.items():
+        if diff is None:
+            suggestions.append(f"{joint}: can't detect reliably.")
+            continue
+        if abs(diff) <= tol:
+            suggestions.append(f"{joint}: good (within ±{tol}°).")
+        else:
+            if diff > 0:
+                # detected angle larger than reference -> joint more open than desired
+                suggestions.append(f"{joint}: decrease angle by {abs(diff):.0f}° (e.g. bend more).")
+            else:
+                suggestions.append(f"{joint}: increase angle by {abs(diff):.0f}° (e.g. straighten more).")
+    return suggestions
+
+# --- Video processing ---
+def process_video(input_path: str, pose_name: str, tolerance: float = DEFAULT_TOLERANCE):
+    # load reference poses
+    ref_poses = load_reference_poses()
+    if pose_name not in ref_poses:
+        return None, f"Pose '{pose_name}' not found in reference poses."
+
+    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 = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 640)
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 480)
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    tmp_out = os.path.join(tempfile.gettempdir(), f"annotated_{Path(input_path).stem}.mp4")
+    out = cv2.VideoWriter(tmp_out, fourcc, fps, (width, height))
+
+    pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5, min_tracking_confidence=0.5)
+    frame_idx = 0
+    aggregate_scores = []
+    joint_scores_over_time = []
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame_idx += 1
+        image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        results = pose.process(image_rgb)
+        annotated = frame.copy()
+
+        if results.pose_landmarks:
+            landmark_xy = landmarks_to_xy(results.pose_landmarks, width, height)
+            detected_angles = compute_joint_angles(landmark_xy)
+            final_percent, per_joint_score, per_joint_diff = compare_angles(detected_angles, reference, tolerance)
+            aggregate_scores.append(final_percent)
+            joint_scores_over_time.append(per_joint_score)
+
+            # draw skeleton - color joints green if within tolerance else red
+            for joint, (p_idx, j_idx, c_idx) in JOINT_TRIPLETS.items():
+                # draw lines parent->joint and joint->child
+                if j_idx in landmark_xy and p_idx in landmark_xy:
+                    x1, y1, v1 = landmark_xy[p_idx]
+                    x2, y2, v2 = landmark_xy[j_idx]
+                    score = per_joint_score.get(joint)
+                    if score is None:
+                        color = (0, 255, 255)  # yellow for unknown
+                    else:
+                        color = (0, 255, 0) if score >= 66 else (0, 165, 255) if score >= 33 else (0, 0, 255)
+                    cv2.line(annotated, (int(x1), int(y1)), (int(x2), int(y2)), color, 3)
+
+                if j_idx in landmark_xy and c_idx in landmark_xy:
+                    x2, y2, v2 = landmark_xy[j_idx]
+                    x3, y3, v3 = landmark_xy[c_idx]
+                    score = per_joint_score.get(joint)
+                    if score is None:
+                        color = (0, 255, 255)
+                    else:
+                        color = (0, 255, 0) if score >= 66 else (0, 165, 255) if score >= 33 else (0, 0, 255)
+                    cv2.line(annotated, (int(x2), int(y2)), (int(x3), int(y3)), color, 3)
+
+            # draw circles at joints with ang value and highlight bad ones
+            for joint, (_, j_idx, _) in JOINT_TRIPLETS.items():
+                if j_idx in landmark_xy:
+                    x, y, v = landmark_xy[j_idx]
+                    score = per_joint_score.get(joint)
+                    if score is None:
+                        cv2.circle(annotated, (int(x), int(y)), 6, (0, 255, 255), -1)
+                    else:
+                        color = (0, 255, 0) if score >= 66 else (0, 165, 255) if score >= 33 else (0, 0, 255)
+                        cv2.circle(annotated, (int(x), int(y)), 8, color, -1)
+                        # put text of angle difference small
+                        diff = per_joint_diff.get(joint)
+                        if diff is not None:
+                            txt = f"{diff:+.0f}°"
+                            cv2.putText(annotated, txt, (int(x)+6, int(y)-6), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+            # frame-level overlay of percent and pose name
+            cv2.putText(annotated, f"{pose_name} - {final_percent:.0f}% correct", (10, 30),
+                        cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 2, cv2.LINE_AA)
+        else:
+            # no landmarks, show message
+            cv2.putText(annotated, "No person detected", (10,30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,0,255), 2)
+
+        out.write(annotated)
+
+    cap.release()
+    out.release()
+    pose.close()
+
+    # aggregate results
+    overall_percent = float(np.mean(aggregate_scores)) if aggregate_scores else 0.0
+    # use last frame joint scores to produce suggestions (or averaged)
+    last_joint_scores = joint_scores_over_time[-1] if joint_scores_over_time else {}
+    # compute last diffs using detected angles from last frame - but we saved diffs only inside loop
+    # For simplicity, recompute suggestions by re-reading last frame's per_joint_diff from process? We'll use the last computed per_joint_diff stored implicitly above:
+    # To keep consistent, re-open video and compute final detected angles on last frame:
+    cap2 = cv2.VideoCapture(input_path)
+    last_frame = None
+    while True:
+        ret, f = cap2.read()
+        if not ret:
+            break
+        last_frame = f
+    cap2.release()
+
+    suggestions = ["(no reliable pose detected)"]
+    if last_frame is not None:
+        h, w = last_frame.shape[:2]
+        with mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5) as pose2:
+            res = pose2.process(cv2.cvtColor(last_frame, cv2.COLOR_BGR2RGB))
+            if res.pose_landmarks:
+                landmark_xy = landmarks_to_xy(res.pose_landmarks, w, h)
+                detected_angles = compute_joint_angles(landmark_xy)
+                _, _, per_joint_diff = compare_angles(detected_angles, reference, tolerance)
+                suggestions = suggest_corrections(per_joint_diff, tol=tolerance)
+            else:
+                suggestions = ["No person detected in final frame to produce suggestions."]
+
+    # return annotated video path and a JSON-like result
+    result = {
+        "pose": pose_name,
+        "score_percent": overall_percent,
+        "suggestions": suggestions
+    }
+    return tmp_out, result
+
+# --- Gradio UI ---
+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 or use your webcam. The app will analyze each frame, compute joint angles via MediaPipe,
+        compare them to a reference pose, and return a percentage correctness plus per-joint corrections.
+        """
+    )
+    with gr.Row():
+        video_in = gr.Video(source="webcam", label="Webcam (or upload a video file)", type="filepath")
+        with gr.Column():
+            pose_dropdown = gr.Dropdown(choices=pose_list, value=pose_list[0], label="Reference Pose")
+            tol_slider = gr.Slider(5, 40, value=DEFAULT_TOLERANCE, step=1, label="Tolerance (degrees)")
+            run_btn = gr.Button("Analyze")
+            output_video = gr.Video(label="Annotated video (downloadable)")
+            output_json = gr.JSON(label="Results and suggestions")
+
+    def analyze(video_path, pose_name, tolerance):
+        if not video_path:
+            return None, {"error": "No input video provided"}
+        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)))