Update app.py

app.py

@@ -1,254 +1,282 @@
 import os
+import math
+import shutil
+import tempfile
+
 import cv2
 import numpy as np
-import mediapipe as mp
-import tempfile
-import shutil
+import pandas as pd
+from scipy.signal import savgol_filter
 import gradio as gr
 
 
-# -----------------------
-# Core pipeline function
-# -----------------------
-def analyze_pushup_video(video_path: str, save_annotated: bool = True, annotated_out_path: str | None = None):
-    if not os.path.exists(video_path):
-        return {
-            "ok": False,
-            "error": f"Could not find input video: {video_path}",
-            "rep_count": 0,
-            "avg_rep_prob": 0.0,
-            "rep_events": [],
-            "annotated_video_path": None,
-        }
-
-    def clamp(x, lo=0.0, hi=1.0):
-        return max(lo, min(hi, x))
-
-    def angle_deg(a, b, c):
-        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):
-        if good_lo <= val <= good_hi:
-            return 1.0
-        if val < good_lo:
-            return clamp((val - ok_lo) / (good_lo - ok_lo))
-        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)
-
-    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,
-    )
+# ----------------------------
+# Settings (same as Colab)
+# ----------------------------
+UP_ANGLE = 155
+DOWN_ANGLE = 105
+MIN_REP_FRAMES = 8
+FRAME_STRIDE = 1
+
+
+# ----------------------------
+# Load YOLO pose model (lazy)
+# ----------------------------
+_MODEL = None
+
+def load_pose_model():
+    global _MODEL
+    if _MODEL is not None:
+        return _MODEL
+
+    from ultralytics import YOLO
+    last_err = None
+    for w in ["yolo11n-pose.pt", "yolov8n-pose.pt"]:
+        try:
+            _MODEL = YOLO(w)
+            print("Loaded model:", w)
+            return _MODEL
+        except Exception as e:
+            last_err = e
+
+    raise RuntimeError(f"Could not load YOLO pose model. Last error: {last_err}")
+
+
+# ----------------------------
+# Helpers (from your script)
+# ----------------------------
+def angle_deg(a, b, c):
+    a = np.asarray(a, dtype=np.float32)
+    b = np.asarray(b, dtype=np.float32)
+    c = np.asarray(c, dtype=np.float32)
+    ba = a - b
+    bc = c - b
+    denom = (np.linalg.norm(ba) * np.linalg.norm(bc)) + 1e-9
+    cosv = np.clip(np.dot(ba, bc) / denom, -1.0, 1.0)
+    return float(math.degrees(math.acos(cosv)))
+
+def pick_best_side(kxy, kconf):
+    left = [5, 7, 9]   # L shoulder, L elbow, L wrist (YOLO COCO indices)
+    right = [6, 8, 10] # R shoulder, R elbow, R wrist
+    if float(np.mean(kconf[right])) >= float(np.mean(kconf[left])):
+        return right, float(np.mean(kconf[right]))
+    return left, float(np.mean(kconf[left]))
+
+def sigmoid(x):
+    return 1.0 / (1.0 + math.exp(-x))
+
+def rep_likelihood(min_ang, max_ang, mean_conf):
+    ang_range = max_ang - min_ang
+    range_score   = sigmoid((ang_range - 45) / 10)
+    depth_score   = sigmoid((DOWN_ANGLE - min_ang) / 8)
+    lockout_score = sigmoid((max_ang - UP_ANGLE) / 8)
+    conf_score    = float(np.clip(mean_conf, 0.0, 1.0))
+    return float(np.clip(range_score * depth_score * lockout_score * conf_score, 0.0, 1.0))
+
+def likelihood_to_score(p):
+    p = float(np.clip(p, 0.0, 1.0))
+    buckets = [
+        (0.50, 1.00, 90, 100),
+        (0.45, 0.50, 80, 89),
+        (0.40, 0.45, 70, 79),
+        (0.35, 0.40, 60, 69),
+        (0.30, 0.35, 50, 59),
+        (0.25, 0.30, 40, 49),
+        (0.20, 0.25, 30, 39),
+        (0.15, 0.20, 20, 29),
+        (0.10, 0.15, 10, 19),
+        (0.00, 0.10, 0, 9),
+    ]
+    for lo, hi, s_lo, s_hi in buckets:
+        if (lo <= p < hi) or (p == 1.0 and hi == 1.0):
+            t = (p - lo) / max(hi - lo, 1e-6)
+            return int(round(s_lo + t * (s_hi - s_lo)))
+    return 0
+
+
+# ----------------------------
+# Core pipeline
+# ----------------------------
+def analyze_pushup_video_yolo(video_path: str, out_dir: str):
+    model = load_pose_model()
 
     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,
-            "avg_rep_prob": 0.0,
-            "rep_events": [],
-            "annotated_video_path": None,
-        }
+        raise RuntimeError("OpenCV could not open the video. Try a different mp4 encoding.")
 
     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
-
-    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))
-
-    state = "UNKNOWN"
-    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
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) or 0
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 0
+
+    # 1) First pass: compute angles + confs per sampled frame
+    angles, confs, frame_ids = [], [], []
+    frame_i = 0
+
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+
+        if frame_i % FRAME_STRIDE != 0:
+            frame_i += 1
+            continue
+
+        res = model(frame, verbose=False)[0]
+        if res.keypoints is None or len(res.keypoints.xy) == 0:
+            angles.append(np.nan)
+            confs.append(0.0)
+            frame_ids.append(frame_i)
+            frame_i += 1
+            continue
+
+        kxy_all = res.keypoints.xy.cpu().numpy()
+        kconf_all = res.keypoints.conf.cpu().numpy()
+
+        # choose best person by mean confidence
+        pidx = int(np.argmax(np.mean(kconf_all, axis=1)))
+        kxy = kxy_all[pidx]
+        kconf = kconf_all[pidx]
+
+        ids, side_conf = pick_best_side(kxy, kconf)
+        if side_conf < 0.2:
+            angles.append(np.nan)
+            confs.append(float(side_conf))
+            frame_ids.append(frame_i)
+            frame_i += 1
+            continue
+
+        a, b, c = kxy[ids[0]], kxy[ids[1]], kxy[ids[2]]
+        angles.append(angle_deg(a, b, c))
+        confs.append(float(side_conf))
+        frame_ids.append(frame_i)
+        frame_i += 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
-
-                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_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}"
-
-            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)
+    cap.release()
 
-    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,
-            "avg_rep_prob": float(np.mean([r["prob"] for r in rep_events])) if rep_events else 0.0,
-            "rep_events": rep_events,
-            "annotated_video_path": annotated_path,
-        }
+    angles = np.array(angles, dtype=np.float32)
+    confs = np.array(confs, dtype=np.float32)
+    frame_ids = np.array(frame_ids, dtype=np.int32)
+
+    if len(angles) < 5:
+        raise RuntimeError("Video too short or no usable frames detected.")
+
+    mask = np.isfinite(angles)
+    if np.any(mask) and not np.all(mask):
+        angles[~mask] = np.interp(frame_ids[~mask], frame_ids[mask], angles[mask])
+    elif not np.any(mask):
+        raise RuntimeError("No valid pose angles detected.")
+
+    win = min(31, (len(angles) // 2) * 2 + 1)
+    win = max(win, 5)  # savgol requires >= 5 for polyorder=2 comfortably
+    angles_smooth = savgol_filter(angles, win, 2)
+
+    # 2) Rep detection on smoothed angles
+    reps = []
+    state = "WAIT_DOWN"
+    rep_min = rep_max = rep_conf_sum = rep_len = rep_start = None
+
+    for i, ang in enumerate(angles_smooth):
+        cf = float(confs[i])
+
+        if state == "WAIT_DOWN":
+            if ang <= DOWN_ANGLE:
+                state = "IN_DOWN"
+                rep_min = rep_max = float(ang)
+                rep_conf_sum = cf
+                rep_len = 1
+                rep_start = i
+        else:
+            rep_min = min(rep_min, float(ang))
+            rep_max = max(rep_max, float(ang))
+            rep_conf_sum += cf
+            rep_len += 1
+
+            if ang >= UP_ANGLE:
+                if rep_len >= MIN_REP_FRAMES:
+                    mean_cf = float(rep_conf_sum / rep_len)
+                    likelihood = rep_likelihood(rep_min, rep_max, mean_cf)
+                    score = likelihood_to_score(likelihood)
+
+                    sf = int(frame_ids[rep_start])
+                    ef = int(frame_ids[i])
+
+                    reps.append({
+                        "rep": len(reps) + 1,
+                        "start_frame": sf,
+                        "end_frame": ef,
+                        "start_time_s": float(sf / fps),
+                        "end_time_s": float(ef / fps),
+                        "min_elbow_angle": float(rep_min),
+                        "max_elbow_angle": float(rep_max),
+                        "mean_kpt_conf": float(mean_cf),
+                        "pushup_likelihood": float(likelihood),
+                        "pushup_score": int(score),
+                    })
+
+                state = "WAIT_DOWN"
+
+    # 3) Save CSV
+    csv_path = os.path.join(out_dir, "pushup_reps.csv")
+    df = pd.DataFrame(reps)
+    df.to_csv(csv_path, index=False)
+
+    # 4) Annotated video
+    annotated_path = os.path.join(out_dir, "pushup_annotated.mp4")
+    cap = cv2.VideoCapture(video_path)
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    writer = cv2.VideoWriter(annotated_path, fourcc, fps, (w, h))
 
-    cap.release()
-    if writer is not None:
-        writer.release()
-    pose.close()
+    rep_windows = [(r["start_frame"], r["end_frame"], r["pushup_score"]) for r in reps]
+
+    frame_i = 0
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
 
-    avg_prob = float(np.mean([r["prob"] for r in rep_events])) if rep_events else 0.0
+        active = next((s for sf, ef, s in rep_windows if sf <= frame_i <= ef), None)
+        count = sum(1 for _, ef, _ in rep_windows if ef < frame_i)
 
-    return {
+        j = int(min(np.searchsorted(frame_ids, frame_i), len(angles_smooth) - 1))
+        ang_disp = float(angles_smooth[j])
+
+        cv2.putText(frame, f"Reps: {count}/{len(reps)}", (20, 40),
+                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255,255,255), 2)
+        cv2.putText(frame, f"Elbow angle: {ang_disp:.1f}", (20, 80),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,255,255), 2)
+        cv2.putText(frame, f"Rep score: {active if active is not None else '-'}", (20, 120),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,255,255), 2)
+
+        writer.write(frame)
+        frame_i += 1
+
+    cap.release()
+    writer.release()
+
+    summary = {
         "ok": True,
         "error": None,
-        "rep_count": rep_count,
-        "avg_rep_prob": avg_prob,
-        "rep_events": rep_events,
-        "annotated_video_path": annotated_path,
+        "rep_count": int(len(reps)),
+        "avg_score": int(round(float(np.mean([r["pushup_score"] for r in reps])))) if reps else 0,
+        "avg_likelihood": float(np.mean([r["pushup_likelihood"] for r in reps])) if reps else 0.0,
+        "rep_events": reps,
     }
 
+    return summary, annotated_path, csv_path
 
-# -----------------------
-# API function (Space endpoint)
-# -----------------------
+
+# ----------------------------
+# API wrapper (robust file handling like your old one)
+# ----------------------------
 def api_analyze(uploaded_file):
-    """
-    Gradio can pass:
-    - a FileData object with .path
-    - a dict with "path"
-    - a file-like with .name
-    """
+    if uploaded_file is None:
+        return {"ok": False, "error": "No file received.", "rep_count": 0, "rep_events": []}, None, None
+
     workdir = tempfile.mkdtemp()
     in_path = os.path.join(workdir, "input.mp4")
 
     # Resolve source path robustly
     src_path = None
-    if uploaded_file is None:
-        return {"ok": False, "error": "No file received.", "rep_count": 0, "avg_rep_prob": 0.0, "rep_events": []}, None
-
     if hasattr(uploaded_file, "path") and uploaded_file.path:
         src_path = uploaded_file.path
     elif isinstance(uploaded_file, dict) and uploaded_file.get("path"):
@@ -258,46 +286,41 @@ def api_analyze(uploaded_file):
     else:
         src_path = str(uploaded_file)
 
-    # Optional: enforce extension yourself (Gradio can mis-detect MIME)
+    # Optional extension check (same idea as your old code)
     ext = os.path.splitext(src_path)[1].lower()
     allowed = {".mp4", ".mov", ".webm", ".mkv"}
     if ext and ext not in allowed:
-        return {"ok": False, "error": f"Unsupported extension: {ext}. Use mp4/mov/webm/mkv.", "rep_count": 0, "avg_rep_prob": 0.0, "rep_events": []}, None
+        return {"ok": False, "error": f"Unsupported extension: {ext}. Use mp4/mov/webm/mkv.", "rep_count": 0, "rep_events": []}, None, None
 
     shutil.copy(src_path, 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 {"ok": False, "error": result["error"], "rep_count": 0, "avg_rep_prob": 0.0, "rep_events": []}, None
-
-    summary = {
-        "ok": True,
-        "error": None,
-        "rep_count": result["rep_count"],
-        "avg_rep_prob": result["avg_rep_prob"],
-        "rep_events": result["rep_events"],
-    }
-    return summary, result["annotated_video_path"]
+    try:
+        summary, annotated_path, csv_path = analyze_pushup_video_yolo(in_path, out_dir=workdir)
+        return summary, annotated_path, csv_path
+    except Exception as e:
+        return {"ok": False, "error": f"{type(e).__name__}: {e}", "rep_count": 0, "rep_events": []}, None, None
 
 
-with gr.Blocks(title="Pushup API") as demo:
-    gr.Markdown("# Pushup Analyzer API\nUpload a video, get rep count + confidence.\n")
+# ----------------------------
+# Gradio UI + API endpoint
+# ----------------------------
+with gr.Blocks(title="Pushup API (YOLO)") as demo:
+    gr.Markdown("# Pushup Analyzer API (YOLO)\nUpload a video, get rep scores + CSV + annotated video.\n")
 
-    # KEY CHANGE: accept any video to avoid Gradio rejecting it
+    # IMPORTANT: keep this as gr.File to avoid the “Invalid file type: ['video']” problem you hit before
     video_file = gr.File(label="Upload video")
 
     btn = gr.Button("Analyze")
     out_json = gr.JSON(label="Results JSON")
     out_video = gr.Video(label="Annotated Output")
+    out_csv = gr.File(label="CSV Output")
 
     btn.click(
         fn=api_analyze,
         inputs=[video_file],
-        outputs=[out_json, out_video],
+        outputs=[out_json, out_video, out_csv],
         api_name="analyze",
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()