Update app.py

app.py

@@ -1,22 +1,37 @@
 import os
-import shutil
-import tempfile
-from typing import Optional, Dict, Any, List
-
 import cv2
 import numpy as np
 import mediapipe as mp
 
-from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import JSONResponse
-
-app = FastAPI()
-
-
-def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotated_out_path: Optional[str] = None) -> Dict[str, Any]:
+import gradio as gr
+import tempfile
+import shutil
+from typing import Optional, Any, Dict
+
+
+# -----------------------
+# Core pipeline function
+# -----------------------
+def analyze_pushup_video(
+    video_path: str,
+    save_annotated: bool = True,
+    annotated_out_path: Optional[str] = None,
+):
+    """
+    Returns dict:
+    {
+      "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}
+        return {"ok": False, "error": f"Could not find input video: {video_path}",
+                "rep_count": 0, "rep_events": [], "annotated_video_path": None}
 
+    # ---------- helpers ----------
     def clamp(x, lo=0.0, hi=1.0):
         return max(lo, min(hi, x))
 
@@ -41,6 +56,7 @@ def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotate
     def ema(prev, x, a=0.25):
         return x if prev is None else (a * x + (1 - a) * prev)
 
+    # ---------- pose ----------
     mp_pose = mp.solutions.pose
     pose = mp_pose.Pose(
         static_image_mode=False,
@@ -54,11 +70,22 @@ def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotate
     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}
+        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 = float(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)
+
+    if W <= 0 or H <= 0:
+        cap.release()
+        pose.close()
+        return {"ok": False, "error": f"Video decode failed (W={W}, H={H}). Try re-encoding the video.",
+                "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
+    # Some encoders hate odd dimensions
+    if W % 2 == 1: W -= 1
+    if H % 2 == 1: H -= 1
 
     annotated_path = None
     writer = None
@@ -68,9 +95,14 @@ def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotate
         annotated_path = annotated_out_path
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
         writer = cv2.VideoWriter(annotated_path, fourcc, fps, (W, H))
+        if not writer.isOpened():
+            # Don’t crash the whole run if writing fails
+            writer = None
+            annotated_path = None
 
+    # ---------- detection ----------
     state = "UNKNOWN"
-    rep_events: List[Dict[str, Any]] = []
+    rep_events = []
     current_rep = None
     rep_count = 0
 
@@ -93,10 +125,14 @@ def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotate
                 break
             frame_idx += 1
 
+            # Resize to the exact writer size if we adjusted odd dims
+            frame = frame[:H, :W]
+
             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
@@ -135,7 +171,7 @@ def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotate
 
                 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))
+                s_vis = clamp((ema_vis - MIN_VIS) / (0.85 - MIN_VIS)) if ema_vis is not None else 0.0
 
                 frame_prob = clamp(0.15 + 0.45 * s_elbow + 0.30 * s_straight + 0.10 * s_vis)
 
@@ -146,8 +182,12 @@ def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotate
                         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)}
-
+                                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
@@ -173,6 +213,13 @@ def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotate
                     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)
 
@@ -181,30 +228,64 @@ def analyze_pushup_video(video_path: str, save_annotated: bool = False, annotate
         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}
+        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}
+    return {"ok": True, "error": None, "rep_count": rep_count,
+            "rep_events": rep_events, "annotated_video_path": annotated_path}
+
+
+# -----------------------
+# Gradio wrapper
+# -----------------------
+def _extract_video_path(video_input: Any) -> str:
+    # Gradio sometimes gives a string path, sometimes a dict with "path"
+    if isinstance(video_input, str):
+        return video_input
+    if isinstance(video_input, dict) and "path" in video_input:
+        return video_input["path"]
+    # Some versions use "name"
+    if isinstance(video_input, dict) and "name" in video_input:
+        return video_input["name"]
+    raise ValueError(f"Unexpected video input type: {type(video_input)} value={video_input}")
 
 
-@app.post("/analyze")
-async def analyze(video: UploadFile = File(...)):
+def gradio_run(video_file):
     workdir = tempfile.mkdtemp()
     in_path = os.path.join(workdir, "input.mp4")
 
-    with open(in_path, "wb") as f:
-        shutil.copyfileobj(video.file, f)
+    src_path = _extract_video_path(video_file)
+    shutil.copy(src_path, in_path)
 
-    result = analyze_pushup_video(in_path, save_annotated=False)
+    out_path = os.path.join(workdir, "annotated.mp4")
+    result = analyze_pushup_video(in_path, save_annotated=True, annotated_out_path=out_path)
 
-    # clean temp (optional)
-    try:
-        shutil.rmtree(workdir)
-    except:
-        pass
+    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"),
+    ],
+    allow_flagging="never",
+)
 
-    return JSONResponse(result)
+if __name__ == "__main__":
+    demo.launch()