save initial vitpose

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,Body
+from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,Body,HTTPException
 from fastapi.staticfiles import StaticFiles
 from vitpose import VitPose
 from dotenv import load_dotenv
@@ -65,7 +65,8 @@ async def upload(background_tasks: BackgroundTasks,
     player_data = json.loads(player_data)
 
     if token != AI_API_TOKEN:
-        return JSONResponse(content={"message": "Unauthorized", "status": 401})
+        
+        raise HTTPException(status_code=401, detail="Unauthorized")
     
     logger.info("reading contents")
     contents = await file.read()
@@ -86,4 +87,5 @@ async def upload(background_tasks: BackgroundTasks,
                               exercise_id)  
 
     # Return the file as a response
+    print(f"returning response")
     return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})

tasks.py

@@ -64,6 +64,8 @@ def process_salto_alto(file_name: str, vitpose: VitPose, player_data: dict, repe
         exercise_id: ID of the exercise
     """
     # Use the provided VitPose instance
+    
+    print(f"start processing")
     model = vitpose.pipeline
     
     # Get player parameters from player_data or use defaults
@@ -152,16 +154,36 @@ def analyze_jump_video(model, input_video, output_video, reference_height=1.68,
     initial_right_shoulder_x = None
     
     # Process first frame to calibrate
-    results_first_frame = model(frame)  # Detect pose in first frame
-    if results_first_frame and results_first_frame[0].keypoints and len(results_first_frame[0].keypoints.xy[0]) > 0:
-        kpts_first = results_first_frame[0].keypoints.xy[0].cpu().numpy()
-        if kpts_first[0][1] > 0 and kpts_first[15][1] > 0 and kpts_first[16][1] > 0:  # Nose and ankles
-            initial_person_height_px = min(kpts_first[15][1], kpts_first[16][1]) - kpts_first[0][1]
-            PX_PER_METER = initial_person_height_px / reference_height
+    output = model(frame)  # Detect pose in first frame
+    keypoints = output.keypoints_xy.float().cpu().numpy()
+    print(f"keypoints {keypoints}")
+    labels = model.pose_estimator_config.label2id
+    print(labels)
+    nose_keypoint = labels["Nose"]
+    L_ankle_keypoint = labels["L_Ankle"]
+    R_ankle_keypoint = labels["R_Ankle"]
+    L_shoulder_keypoint = labels["L_Shoulder"]
+    R_shoulder_keypoint = labels["R_Shoulder"]
+    print(f"nose_keypoint {nose_keypoint}")
+    print(f"L_ankle_keypoint {L_ankle_keypoint}")
+    print(f"R_ankle_keypoint {R_ankle_keypoint}")
+    print(f"L_shoulder_keypoint {L_shoulder_keypoint}")
+    print(f"R_shoulder_keypoint {R_shoulder_keypoint}")
+
+    if (
+        keypoints is not None
+        and len(keypoints) > 0
+        and len(keypoints[0]) > 0):
+        
+        kpts_first = keypoints[0]
+        if len(kpts_first[nose_keypoint]) > 0 and len(kpts_first[L_ankle_keypoint]) > 0:  # Nose and ankles
+            initial_person_height_px = min(kpts_first[L_ankle_keypoint][1], kpts_first[R_ankle_keypoint][1]) - kpts_first[nose_keypoint][1]
+            print(f"initial_person_height_px {initial_person_height_px}")
+            PX_PER_METER = float(initial_person_height_px) / float(reference_height)
             print(f"Escala calculada: {PX_PER_METER:.2f} px/m")
-        if kpts_first[5][0] > 0 and kpts_first[6][0] > 0:  # Left (5) and right (6) shoulders
-            initial_left_shoulder_x = int(kpts_first[5][0])
-            initial_right_shoulder_x = int(kpts_first[6][0])
+        if len(kpts_first[L_shoulder_keypoint]) > 0 and len(kpts_first[R_shoulder_keypoint]) > 0:  # Left (5) and right (6) shoulders
+            initial_left_shoulder_x = int(kpts_first[L_shoulder_keypoint][0])
+            initial_right_shoulder_x = int(kpts_first[R_shoulder_keypoint][0])
     
     if PX_PER_METER is None or initial_left_shoulder_x is None or initial_right_shoulder_x is None:
         print("No se pudo calibrar la escala o detectar los hombros en el primer frame.")
@@ -200,87 +222,126 @@ def analyze_jump_video(model, input_video, output_video, reference_height=1.68,
             break
         
         annotated_frame = frame.copy()
-        results = model(annotated_frame)
         
-        if results and results[0].keypoints and len(results[0].keypoints.xy[0]) > 0:
-            person_detected = True
-            kpts = results[0].keypoints.xy[0].cpu().numpy()
-            nose = kpts[0]
-            ankles = [kpts[15], kpts[16]]
-            left_shoulder = kpts[5]
-            right_shoulder = kpts[6]
+        # Add try-except block around the model inference to catch any model errors
+        try:
+            output = model(annotated_frame)
+            keypoints = output.keypoints_xy.float().cpu().numpy()
             
-            if nose[1] > 0 and all(a[1] > 0 for a in ankles) and left_shoulder[0] > 0 and right_shoulder[0] > 0:
-                current_ankle_y = min(a[1] for a in ankles)
-                last_detected_ankles_y = current_ankle_y  # Save current ankle position
-                current_head_y = nose[1]
-                current_left_shoulder_x = int(left_shoulder[0])
-                current_right_shoulder_x = int(right_shoulder[0])
-                
-                # Smooth ankle and head positions
-                ankle_y_history.append(current_ankle_y)
-                if len(ankle_y_history) > SMOOTHING_WINDOW:
-                    ankle_y_history.pop(0)
-                smoothed_ankle_y = np.mean(ankle_y_history)
-                
-                head_y_history.append(current_head_y)
-                if len(head_y_history) > SMOOTHING_WINDOW:
-                    head_y_history.pop(0)
-                smoothed_head_y = np.mean(head_y_history)
+            # Verify that keypoints array has valid data before processing
+            if (keypoints is not None and 
+                len(keypoints) > 0 and 
+                len(keypoints[0]) > 0 and
+                keypoints.size > 0):  # Check if array is not empty
                 
-                # Calculate vertical velocity (using head position)
-                head_y_buffer.append(smoothed_head_y)
-                if len(head_y_buffer) > VELOCITY_WINDOW:
-                    head_y_buffer.pop(0)
-                    if PX_PER_METER is not None and fps > 0:
-                        delta_y_pixels = head_y_buffer[0] - head_y_buffer[-1]
-                        delta_y_meters = delta_y_pixels / PX_PER_METER
-                        delta_t = VELOCITY_WINDOW / fps
-                        velocity_vertical = delta_y_meters / delta_t
+                person_detected = True
+                kpts = keypoints[0]
                 
-                # Set ground level in first frame where ankles are detected
-                if ground_level is None:
-                    ground_level = smoothed_ankle_y
-                    takeoff_head_y = smoothed_head_y
-                
-                relative_ankle_change = (ground_level - smoothed_ankle_y) / ground_level if ground_level > 0 else 0
-                
-                # Detect jump start
-                if not jump_started and relative_ankle_change > JUMP_THRESHOLD_PERCENT:
-                    jump_started = True
-                    takeoff_head_y = smoothed_head_y
-                    max_jump_height = 0
-                    max_head_height_px = smoothed_head_y
-                
-                # Detect jump end
-                if jump_started and relative_ankle_change <= JUMP_THRESHOLD_PERCENT:
-                    # Add to repetition data
-                    salto_alto = calculate_absolute_jump_height(reference_height, max_jump_height)
-                    repetition_data.append({
-                        "repetition": repetition_count + 1,
-                        "relative_jump_m": round(max_jump_height, 3),
-                        "absolute_jump_m": round(salto_alto, 3),
-                        "peak_power_watts": round(current_power, 1)
-                    })
-                    repetition_count += 1
-                    jump_started = False
-                
-                # Update jump metrics while in air
-                if jump_started:
-                    relative_jump = (takeoff_head_y - smoothed_head_y) / PX_PER_METER
-                    if relative_jump > max_jump_height:
-                        max_jump_height = relative_jump
-                    if smoothed_head_y < max_head_height_px:
-                        max_head_height_px = smoothed_head_y
-                    if relative_jump:
-                        current_power = calculate_peak_power_sayer(relative_jump, body_mass_kg)
-                        if current_power > peak_power_sayer:
-                            peak_power_sayer = current_power
-        else:
-            last_detected_ankles_y = None  # Reset position if ankles not detected
-            velocity_vertical = 0.0  # Reset velocity if no reliable detection
+                # Make sure all required keypoints are detected
+                if (nose_keypoint < len(kpts) and L_ankle_keypoint < len(kpts) and 
+                    R_ankle_keypoint < len(kpts) and L_shoulder_keypoint < len(kpts) and 
+                    R_shoulder_keypoint < len(kpts)):
+                    
+                    nose = kpts[nose_keypoint]
+                    ankles = [kpts[L_ankle_keypoint], kpts[R_ankle_keypoint]]
+                    left_shoulder = kpts[L_shoulder_keypoint]
+                    right_shoulder = kpts[R_shoulder_keypoint]
+                    
+                    # Check if keypoints have valid coordinates
+                    if (nose[0] > 0 and nose[1] > 0 and 
+                        all(a[0] > 0 and a[1] > 0 for a in ankles) and 
+                        left_shoulder[0] > 0 and left_shoulder[1] > 0 and 
+                        right_shoulder[0] > 0 and right_shoulder[1] > 0):
+                        
+                        # Continue with existing processing
+                        current_ankle_y = min(a[1] for a in ankles)
+                        last_detected_ankles_y = current_ankle_y
+                        current_head_y = nose[1]
+                        current_left_shoulder_x = int(left_shoulder[0])
+                        current_right_shoulder_x = int(right_shoulder[0])
+                        
+                        # Smooth ankle and head positions
+                        ankle_y_history.append(current_ankle_y)
+                        if len(ankle_y_history) > SMOOTHING_WINDOW:
+                            ankle_y_history.pop(0)
+                        smoothed_ankle_y = np.mean(ankle_y_history)
+                        
+                        head_y_history.append(current_head_y)
+                        if len(head_y_history) > SMOOTHING_WINDOW:
+                            head_y_history.pop(0)
+                        smoothed_head_y = np.mean(head_y_history)
+                        
+                        # Calculate vertical velocity (using head position)
+                        head_y_buffer.append(smoothed_head_y)
+                        if len(head_y_buffer) > VELOCITY_WINDOW:
+                            head_y_buffer.pop(0)
+                            if PX_PER_METER is not None and fps > 0:
+                                delta_y_pixels = head_y_buffer[0] - head_y_buffer[-1]
+                                delta_y_meters = delta_y_pixels / PX_PER_METER
+                                delta_t = VELOCITY_WINDOW / fps
+                                velocity_vertical = delta_y_meters / delta_t
+                        
+                        # Set ground level in first frame where ankles are detected
+                        if ground_level is None:
+                            ground_level = smoothed_ankle_y
+                            takeoff_head_y = smoothed_head_y
+                        
+                        relative_ankle_change = (ground_level - smoothed_ankle_y) / ground_level if ground_level > 0 else 0
+                        
+                        # Detect jump start
+                        if not jump_started and relative_ankle_change > JUMP_THRESHOLD_PERCENT:
+                            jump_started = True
+                            takeoff_head_y = smoothed_head_y
+                            max_jump_height = 0
+                            max_head_height_px = smoothed_head_y
+                        
+                        # Detect jump end
+                        if jump_started and relative_ankle_change <= JUMP_THRESHOLD_PERCENT:
+                            # Add to repetition data
+                            salto_alto = calculate_absolute_jump_height(reference_height, max_jump_height)
+                            repetition_data.append({
+                                "repetition": repetition_count + 1,
+                                "relative_jump_m": round(max_jump_height, 3),
+                                "absolute_jump_m": round(salto_alto, 3),
+                                "peak_power_watts": round(current_power, 1)
+                            })
+                            repetition_count += 1
+                            jump_started = False
+                        
+                        # Update jump metrics while in air
+                        if jump_started:
+                            relative_jump = (takeoff_head_y - smoothed_head_y) / PX_PER_METER
+                            if relative_jump > max_jump_height:
+                                max_jump_height = relative_jump
+                            if smoothed_head_y < max_head_height_px:
+                                max_head_height_px = smoothed_head_y
+                            if relative_jump:
+                                current_power = calculate_peak_power_sayer(relative_jump, body_mass_kg)
+                                if current_power > peak_power_sayer:
+                                    peak_power_sayer = current_power
+                    else:
+                        # Skip processing for this frame - invalid coordinates
+                        print("Skipping frame - invalid keypoint coordinates")
+                        print(f"keypoints {keypoints}")
+                else:
+                    # Skip processing for this frame - missing required keypoints
+                    print("Skipping frame - missing required keypoints")
+                    print(f"keypoints {keypoints}")
+            else:
+                # Skip processing for this frame - no valid keypoints detected
+                print("Skipping frame - no valid keypoints detected")
+                print(f"keypoints {keypoints}")
+                last_detected_ankles_y = None
+                velocity_vertical = 0.0
+        except Exception as e:
+            # Handle any other exceptions that might occur during model inference
+            print(f"Error processing frame: {e}")
+            print(f"keypoints {keypoints}")
+            last_detected_ankles_y = None
+            velocity_vertical = 0.0
         
-        # Calculate absolute jump height
+        # Calculate metrics and draw overlay even if keypoints weren't detected
+        # This ensures video continues to show previous metrics
         salto_alto = calculate_absolute_jump_height(reference_height, max_jump_height)
         
         # Draw floating metric boxes
@@ -370,7 +431,7 @@ def calculate_absolute_jump_height(reference_height, relative_jump):
     Returns:
         Absolute jump height in meters
     """
-    absolute_jump = reference_height + relative_jump
+    absolute_jump = float(reference_height) + float(relative_jump)
     # Apply validation rule
     if absolute_jump > 1.72:
         return absolute_jump

vitpose.py

@@ -17,7 +17,7 @@ class VitPose:
             object_detection_checkpoint="PekingU/rtdetr_r50vd_coco_o365",
             pose_estimation_checkpoint="usyd-community/vitpose-plus-small",
             device="cuda" if torch.cuda.is_available() else "cpu",
-            dtype=torch.float16,
+            dtype=torch.bfloat16,
             compile=True,  # or True to get more speedup
         )
         self.output_video_path = None