Update app.py

app.py

@@ -2,10 +2,10 @@ import cv2
 import mediapipe as mp
 import numpy as np
 import gradio as gr
+import os
 
 mp_pose = mp.solutions.pose
 mp_drawing = mp.solutions.drawing_utils
-
 pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5, min_tracking_confidence=0.5)
 
 def calculate_angle(a, b, c):
@@ -15,18 +15,40 @@ def calculate_angle(a, b, c):
     ba = a - b
     bc = c - b
     cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
-    angle = np.arccos(cosine_angle)
+    angle = np.arccos(np.clip(cosine_angle, -1.0, 1.0))
     return np.degrees(angle)
 
-def detect_pose_video(video_path):
+def detect_pose_video(video_path, max_duration=20):
+    if not os.path.exists(video_path):
+        return None
+
     cap = cv2.VideoCapture(video_path)
-    output_frames = []
+    if not cap.isOpened():
+        raise ValueError("Cannot open video file")
+
+    # Video properties
+    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)
+    max_frames = int(fps * max_duration)  # Limit to 20 seconds
+
+    out_path = "annotated_video.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
 
-    while True:
+    frame_count = 0
+    while frame_count < max_frames:
         ret, frame = cap.read()
         if not ret:
             break
 
+        # Resize if too big
+        max_dim = 640
+        h, w, _ = frame.shape
+        if max(h, w) > max_dim:
+            scale = max_dim / max(h, w)
+            frame = cv2.resize(frame, (int(w*scale), int(h*scale)))
+
         frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         results = pose.process(frame_rgb)
 
@@ -40,31 +62,23 @@ def detect_pose_video(video_path):
             )
 
             # Example: left elbow angle
-            h, w, _ = frame.shape
             landmarks = results.pose_landmarks.landmark
-            shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x * w,
-                        landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y * h]
-            elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x * w,
-                     landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y * h]
-            wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x * w,
-                     landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y * h]
+            shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x * frame.shape[1],
+                        landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y * frame.shape[0]]
+            elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x * frame.shape[1],
+                     landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y * frame.shape[0]]
+            wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x * frame.shape[1],
+                     landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y * frame.shape[0]]
 
             angle = calculate_angle(shoulder, elbow, wrist)
             cv2.putText(frame, f"Left Elbow: {int(angle)} deg", (20,40),
                         cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,0), 2)
 
-        output_frames.append(frame)
+        out.write(frame)
+        frame_count += 1
 
     cap.release()
-
-    # Convert frames to video file
-    out_path = "annotated_video.mp4"
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(out_path, fourcc, 20.0, (frame.shape[1], frame.shape[0]))
-    for f in output_frames:
-        out.write(f)
     out.release()
-
     return out_path
 
 # Gradio interface
@@ -73,7 +87,7 @@ demo = gr.Interface(
     inputs=gr.Video(label="Upload Video"),
     outputs=gr.Video(label="Annotated Video"),
     title="Human Pose Estimation on Video",
-    description="Upload a video and see pose landmarks & joint angles detected in real-time."
+    description="Upload a video (max 20 seconds will be processed) and see pose landmarks & joint angles."
 )
 
 if __name__ == "__main__":