Create app.py

app.py

@@ -0,0 +1,298 @@
+import os
+import math
+import cv2
+import numpy as np
+import mediapipe as mp
+
+import gradio as gr
+import tempfile
+import shutil
+
+
+# -----------------------
+# Core pipeline function
+# -----------------------
+def analyze_pushup_video(video_path: str, save_annotated: bool = True, annotated_out_path: str | None = None):
+    """
+    Runs MediaPipe Pose on a video, counts pushup reps, and returns:
+    {
+      "ok": bool,
+      "error": str | None,
+      "rep_count": int,
+      "rep_events": list[dict],
+      "annotated_video_path": str | None
+    }
+    """
+    if not os.path.exists(video_path):
+        return {
+            "ok": False,
+            "error": f"Could not find input video: {video_path}",
+            "rep_count": 0,
+            "rep_events": [],
+            "annotated_video_path": None,
+        }
+
+    # ---------- Math helpers ----------
+    def clamp(x, lo=0.0, hi=1.0):
+        return max(lo, min(hi, x))
+
+    def angle_deg(a, b, c):
+        """Angle ABC in degrees using points a,b,c as (x,y)."""
+        a = np.array(a, dtype=np.float32)
+        b = np.array(b, dtype=np.float32)
+        c = np.array(c, dtype=np.float32)
+        ba = a - b
+        bc = c - b
+        denom = (np.linalg.norm(ba) * np.linalg.norm(bc) + 1e-9)
+        cosang = float(np.dot(ba, bc) / denom)
+        cosang = max(-1.0, min(1.0, cosang))
+        return float(np.degrees(np.arccos(cosang)))
+
+    def score_from_range(val, good_lo, good_hi, ok_lo, ok_hi):
+        """
+        Returns 1 if val in [good_lo, good_hi],
+        fades to 0 by the time it reaches ok_lo/ok_hi.
+        """
+        if good_lo <= val <= good_hi:
+            return 1.0
+        if val < good_lo:
+            return clamp((val - ok_lo) / (good_lo - ok_lo))
+        else:
+            return clamp((ok_hi - val) / (ok_hi - good_hi))
+
+    def ema(prev, x, a=0.25):
+        return x if prev is None else (a * x + (1 - a) * prev)
+
+    # ---------- Pose setup ----------
+    mp_pose = mp.solutions.pose
+    pose = mp_pose.Pose(
+        static_image_mode=False,
+        model_complexity=1,
+        smooth_landmarks=True,
+        enable_segmentation=False,
+        min_detection_confidence=0.5,
+        min_tracking_confidence=0.5,
+    )
+
+    # ---------- Video I/O ----------
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        pose.close()
+        return {
+            "ok": False,
+            "error": "OpenCV could not open the video. Try a different mp4 encoding.",
+            "rep_count": 0,
+            "rep_events": [],
+            "annotated_video_path": None,
+        }
+
+    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
+    W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) or 0
+    H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 0
+
+    # Output path handling
+    annotated_path = None
+    writer = None
+    if save_annotated:
+        if annotated_out_path is None:
+            annotated_out_path = os.path.join(tempfile.mkdtemp(), "annotated.mp4")
+        annotated_path = annotated_out_path
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        writer = cv2.VideoWriter(annotated_path, fourcc, fps, (W, H))
+
+    # ---------- Pushup detection logic ----------
+    state = "UNKNOWN"  # "UP" or "DOWN"
+    rep_events = []
+    current_rep = None
+    rep_count = 0
+
+    ema_elbow = None
+    ema_straight = None
+    ema_vis = None
+    alpha = 0.25
+
+    UP_ELBOW_DEG = 155
+    DOWN_ELBOW_DEG = 105
+    MIN_VIS = 0.45
+    MIN_REP_TIME_S = 0.35
+
+    frame_idx = -1
+
+    try:
+        while True:
+            ok, frame = cap.read()
+            if not ok:
+                break
+            frame_idx += 1
+
+            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            res = pose.process(rgb)
+
+            frame_prob = 0.0
+            debug_txt = "No pose"
+
+            if res.pose_landmarks:
+                lms = res.pose_landmarks.landmark
+
+                # Choose side: whichever shoulder has higher visibility
+                Ls = lms[mp_pose.PoseLandmark.LEFT_SHOULDER.value]
+                Rs = lms[mp_pose.PoseLandmark.RIGHT_SHOULDER.value]
+                left_side = (Ls.visibility >= Rs.visibility)
+
+                if left_side:
+                    shoulder = lms[mp_pose.PoseLandmark.LEFT_SHOULDER.value]
+                    elbow = lms[mp_pose.PoseLandmark.LEFT_ELBOW.value]
+                    wrist = lms[mp_pose.PoseLandmark.LEFT_WRIST.value]
+                    hip = lms[mp_pose.PoseLandmark.LEFT_HIP.value]
+                    ankle = lms[mp_pose.PoseLandmark.LEFT_ANKLE.value]
+                else:
+                    shoulder = lms[mp_pose.PoseLandmark.RIGHT_SHOULDER.value]
+                    elbow = lms[mp_pose.PoseLandmark.RIGHT_ELBOW.value]
+                    wrist = lms[mp_pose.PoseLandmark.RIGHT_WRIST.value]
+                    hip = lms[mp_pose.PoseLandmark.RIGHT_HIP.value]
+                    ankle = lms[mp_pose.PoseLandmark.RIGHT_ANKLE.value]
+
+                vis = float(np.mean([shoulder.visibility, elbow.visibility, wrist.visibility, hip.visibility, ankle.visibility]))
+                ema_vis = ema(ema_vis, vis, alpha)
+
+                sh = (shoulder.x, shoulder.y)
+                el = (elbow.x, elbow.y)
+                wr = (wrist.x, wrist.y)
+                hp = (hip.x, hip.y)
+                ak = (ankle.x, ankle.y)
+
+                elbow_deg = angle_deg(sh, el, wr)
+                straight_deg = angle_deg(sh, hp, ak)
+
+                ema_elbow = ema(ema_elbow, elbow_deg, alpha)
+                ema_straight = ema(ema_straight, straight_deg, alpha)
+
+                s_straight = score_from_range(ema_straight, 165, 185, 145, 195)
+                s_elbow = score_from_range(ema_elbow, 85, 175, 60, 190)
+                s_vis = clamp((ema_vis - MIN_VIS) / (0.85 - MIN_VIS))
+
+                frame_prob = clamp(0.15 + 0.45 * s_elbow + 0.30 * s_straight + 0.10 * s_vis)
+
+                trusted = (ema_vis is not None and ema_vis >= MIN_VIS)
+
+                if trusted:
+                    if state in ["UNKNOWN", "UP"]:
+                        if ema_elbow <= DOWN_ELBOW_DEG and frame_prob >= 0.45:
+                            state = "DOWN"
+                            if current_rep is None:
+                                current_rep = {
+                                    "start_f": frame_idx,
+                                    "frame_probs": [],
+                                    "min_elbow": float(ema_elbow),
+                                    "min_straight": float(ema_straight),
+                                }
+
+                    elif state == "DOWN":
+                        if ema_elbow >= UP_ELBOW_DEG and frame_prob >= 0.35:
+                            end_f = frame_idx
+                            if current_rep is not None:
+                                duration_s = (end_f - current_rep["start_f"]) / fps
+                                if duration_s >= MIN_REP_TIME_S:
+                                    rep_count += 1
+                                    probs = current_rep["frame_probs"] if current_rep["frame_probs"] else [frame_prob]
+                                    rep_prob = float(np.mean(probs))
+
+                                    rep_events.append({
+                                        "rep": rep_count,
+                                        "start_f": int(current_rep["start_f"]),
+                                        "end_f": int(end_f),
+                                        "start_t": float(current_rep["start_f"] / fps),
+                                        "end_t": float(end_f / fps),
+                                        "prob": float(rep_prob),
+                                        "min_elbow": float(current_rep["min_elbow"]),
+                                        "min_straight": float(current_rep["min_straight"]),
+                                    })
+                            current_rep = None
+                            state = "UP"
+
+                if current_rep is not None:
+                    current_rep["frame_probs"].append(float(frame_prob))
+                    current_rep["min_elbow"] = float(min(current_rep["min_elbow"], ema_elbow))
+                    current_rep["min_straight"] = float(min(current_rep["min_straight"], ema_straight))
+
+                debug_txt = f"{'L' if left_side else 'R'} vis={ema_vis:.2f} elbow={ema_elbow:.0f} straight={ema_straight:.0f} p={frame_prob:.2f} state={state}"
+
+                # Simple overlay
+                cv2.putText(frame, f"Reps: {rep_count}", (20, 40),
+                            cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 2, cv2.LINE_AA)
+                cv2.putText(frame, debug_txt[:90], (20, 75),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
+
+            else:
+                cv2.putText(frame, f"Reps: {rep_count}", (20, 40),
+                            cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 2, cv2.LINE_AA)
+                cv2.putText(frame, debug_txt[:90], (20, 75),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
+
+            if writer is not None:
+                writer.write(frame)
+
+    except Exception as e:
+        cap.release()
+        if writer is not None:
+            writer.release()
+        pose.close()
+        return {
+            "ok": False,
+            "error": f"Runtime error: {type(e).__name__}: {e}",
+            "rep_count": rep_count,
+            "rep_events": rep_events,
+            "annotated_video_path": annotated_path,
+        }
+
+    cap.release()
+    if writer is not None:
+        writer.release()
+    pose.close()
+
+    return {
+        "ok": True,
+        "error": None,
+        "rep_count": rep_count,
+        "rep_events": rep_events,
+        "annotated_video_path": annotated_path,
+    }
+
+
+# -----------------------
+# Gradio wrapper
+# -----------------------
+def gradio_run(video_file):
+    # video_file is a temp path provided by Gradio
+    workdir = tempfile.mkdtemp()
+    in_path = os.path.join(workdir, "input.mp4")
+    shutil.copy(video_file, in_path)
+
+    out_path = os.path.join(workdir, "annotated.mp4")
+    result = analyze_pushup_video(in_path, save_annotated=True, annotated_out_path=out_path)
+
+    if not result["ok"]:
+        return "Error: " + str(result["error"]), None, []
+
+    summary = f"Rep count: {result['rep_count']}\n"
+    if result["rep_events"]:
+        avg_prob = sum(r["prob"] for r in result["rep_events"]) / len(result["rep_events"])
+        summary += f"Avg rep probability: {avg_prob:.2f}\n"
+
+    return summary, result["annotated_video_path"], result["rep_events"]
+
+
+demo = gr.Interface(
+    fn=gradio_run,
+    inputs=gr.Video(label="Upload pushup video"),
+    outputs=[
+        gr.Textbox(label="Results"),
+        gr.Video(label="Annotated output"),
+        gr.JSON(label="Per-rep details"),
+    ],
+    title="Pushup Prototype",
+    description="Uploads a video, counts reps, and gives per-rep likelihood.",
+)
+
+if __name__ == "__main__":
+    demo.launch()