nicolasbuitragob commited on
Commit
5504846
·
1 Parent(s): d5bf444

save initial vitpose

Browse files
Files changed (3) hide show
  1. app.py +4 -2
  2. tasks.py +147 -86
  3. vitpose.py +1 -1
app.py CHANGED
@@ -1,4 +1,4 @@
1
- from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,Body
2
  from fastapi.staticfiles import StaticFiles
3
  from vitpose import VitPose
4
  from dotenv import load_dotenv
@@ -65,7 +65,8 @@ async def upload(background_tasks: BackgroundTasks,
65
  player_data = json.loads(player_data)
66
 
67
  if token != AI_API_TOKEN:
68
- return JSONResponse(content={"message": "Unauthorized", "status": 401})
 
69
 
70
  logger.info("reading contents")
71
  contents = await file.read()
@@ -86,4 +87,5 @@ async def upload(background_tasks: BackgroundTasks,
86
  exercise_id)
87
 
88
  # Return the file as a response
 
89
  return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})
 
1
+ from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,Body,HTTPException
2
  from fastapi.staticfiles import StaticFiles
3
  from vitpose import VitPose
4
  from dotenv import load_dotenv
 
65
  player_data = json.loads(player_data)
66
 
67
  if token != AI_API_TOKEN:
68
+
69
+ raise HTTPException(status_code=401, detail="Unauthorized")
70
 
71
  logger.info("reading contents")
72
  contents = await file.read()
 
87
  exercise_id)
88
 
89
  # Return the file as a response
90
+ print(f"returning response")
91
  return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})
tasks.py CHANGED
@@ -64,6 +64,8 @@ def process_salto_alto(file_name: str, vitpose: VitPose, player_data: dict, repe
64
  exercise_id: ID of the exercise
65
  """
66
  # Use the provided VitPose instance
 
 
67
  model = vitpose.pipeline
68
 
69
  # 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,
152
  initial_right_shoulder_x = None
153
 
154
  # Process first frame to calibrate
155
- results_first_frame = model(frame) # Detect pose in first frame
156
- if results_first_frame and results_first_frame[0].keypoints and len(results_first_frame[0].keypoints.xy[0]) > 0:
157
- kpts_first = results_first_frame[0].keypoints.xy[0].cpu().numpy()
158
- if kpts_first[0][1] > 0 and kpts_first[15][1] > 0 and kpts_first[16][1] > 0: # Nose and ankles
159
- initial_person_height_px = min(kpts_first[15][1], kpts_first[16][1]) - kpts_first[0][1]
160
- PX_PER_METER = initial_person_height_px / reference_height
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  print(f"Escala calculada: {PX_PER_METER:.2f} px/m")
162
- if kpts_first[5][0] > 0 and kpts_first[6][0] > 0: # Left (5) and right (6) shoulders
163
- initial_left_shoulder_x = int(kpts_first[5][0])
164
- initial_right_shoulder_x = int(kpts_first[6][0])
165
 
166
  if PX_PER_METER is None or initial_left_shoulder_x is None or initial_right_shoulder_x is None:
167
  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,
200
  break
201
 
202
  annotated_frame = frame.copy()
203
- results = model(annotated_frame)
204
 
205
- if results and results[0].keypoints and len(results[0].keypoints.xy[0]) > 0:
206
- person_detected = True
207
- kpts = results[0].keypoints.xy[0].cpu().numpy()
208
- nose = kpts[0]
209
- ankles = [kpts[15], kpts[16]]
210
- left_shoulder = kpts[5]
211
- right_shoulder = kpts[6]
212
 
213
- if nose[1] > 0 and all(a[1] > 0 for a in ankles) and left_shoulder[0] > 0 and right_shoulder[0] > 0:
214
- current_ankle_y = min(a[1] for a in ankles)
215
- last_detected_ankles_y = current_ankle_y # Save current ankle position
216
- current_head_y = nose[1]
217
- current_left_shoulder_x = int(left_shoulder[0])
218
- current_right_shoulder_x = int(right_shoulder[0])
219
-
220
- # Smooth ankle and head positions
221
- ankle_y_history.append(current_ankle_y)
222
- if len(ankle_y_history) > SMOOTHING_WINDOW:
223
- ankle_y_history.pop(0)
224
- smoothed_ankle_y = np.mean(ankle_y_history)
225
-
226
- head_y_history.append(current_head_y)
227
- if len(head_y_history) > SMOOTHING_WINDOW:
228
- head_y_history.pop(0)
229
- smoothed_head_y = np.mean(head_y_history)
230
 
231
- # Calculate vertical velocity (using head position)
232
- head_y_buffer.append(smoothed_head_y)
233
- if len(head_y_buffer) > VELOCITY_WINDOW:
234
- head_y_buffer.pop(0)
235
- if PX_PER_METER is not None and fps > 0:
236
- delta_y_pixels = head_y_buffer[0] - head_y_buffer[-1]
237
- delta_y_meters = delta_y_pixels / PX_PER_METER
238
- delta_t = VELOCITY_WINDOW / fps
239
- velocity_vertical = delta_y_meters / delta_t
240
 
241
- # Set ground level in first frame where ankles are detected
242
- if ground_level is None:
243
- ground_level = smoothed_ankle_y
244
- takeoff_head_y = smoothed_head_y
245
-
246
- relative_ankle_change = (ground_level - smoothed_ankle_y) / ground_level if ground_level > 0 else 0
247
-
248
- # Detect jump start
249
- if not jump_started and relative_ankle_change > JUMP_THRESHOLD_PERCENT:
250
- jump_started = True
251
- takeoff_head_y = smoothed_head_y
252
- max_jump_height = 0
253
- max_head_height_px = smoothed_head_y
254
-
255
- # Detect jump end
256
- if jump_started and relative_ankle_change <= JUMP_THRESHOLD_PERCENT:
257
- # Add to repetition data
258
- salto_alto = calculate_absolute_jump_height(reference_height, max_jump_height)
259
- repetition_data.append({
260
- "repetition": repetition_count + 1,
261
- "relative_jump_m": round(max_jump_height, 3),
262
- "absolute_jump_m": round(salto_alto, 3),
263
- "peak_power_watts": round(current_power, 1)
264
- })
265
- repetition_count += 1
266
- jump_started = False
267
-
268
- # Update jump metrics while in air
269
- if jump_started:
270
- relative_jump = (takeoff_head_y - smoothed_head_y) / PX_PER_METER
271
- if relative_jump > max_jump_height:
272
- max_jump_height = relative_jump
273
- if smoothed_head_y < max_head_height_px:
274
- max_head_height_px = smoothed_head_y
275
- if relative_jump:
276
- current_power = calculate_peak_power_sayer(relative_jump, body_mass_kg)
277
- if current_power > peak_power_sayer:
278
- peak_power_sayer = current_power
279
- else:
280
- last_detected_ankles_y = None # Reset position if ankles not detected
281
- velocity_vertical = 0.0 # Reset velocity if no reliable detection
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
282
 
283
- # Calculate absolute jump height
 
284
  salto_alto = calculate_absolute_jump_height(reference_height, max_jump_height)
285
 
286
  # Draw floating metric boxes
@@ -370,7 +431,7 @@ def calculate_absolute_jump_height(reference_height, relative_jump):
370
  Returns:
371
  Absolute jump height in meters
372
  """
373
- absolute_jump = reference_height + relative_jump
374
  # Apply validation rule
375
  if absolute_jump > 1.72:
376
  return absolute_jump
 
64
  exercise_id: ID of the exercise
65
  """
66
  # Use the provided VitPose instance
67
+
68
+ print(f"start processing")
69
  model = vitpose.pipeline
70
 
71
  # Get player parameters from player_data or use defaults
 
154
  initial_right_shoulder_x = None
155
 
156
  # Process first frame to calibrate
157
+ output = model(frame) # Detect pose in first frame
158
+ keypoints = output.keypoints_xy.float().cpu().numpy()
159
+ print(f"keypoints {keypoints}")
160
+ labels = model.pose_estimator_config.label2id
161
+ print(labels)
162
+ nose_keypoint = labels["Nose"]
163
+ L_ankle_keypoint = labels["L_Ankle"]
164
+ R_ankle_keypoint = labels["R_Ankle"]
165
+ L_shoulder_keypoint = labels["L_Shoulder"]
166
+ R_shoulder_keypoint = labels["R_Shoulder"]
167
+ print(f"nose_keypoint {nose_keypoint}")
168
+ print(f"L_ankle_keypoint {L_ankle_keypoint}")
169
+ print(f"R_ankle_keypoint {R_ankle_keypoint}")
170
+ print(f"L_shoulder_keypoint {L_shoulder_keypoint}")
171
+ print(f"R_shoulder_keypoint {R_shoulder_keypoint}")
172
+
173
+ if (
174
+ keypoints is not None
175
+ and len(keypoints) > 0
176
+ and len(keypoints[0]) > 0):
177
+
178
+ kpts_first = keypoints[0]
179
+ if len(kpts_first[nose_keypoint]) > 0 and len(kpts_first[L_ankle_keypoint]) > 0: # Nose and ankles
180
+ initial_person_height_px = min(kpts_first[L_ankle_keypoint][1], kpts_first[R_ankle_keypoint][1]) - kpts_first[nose_keypoint][1]
181
+ print(f"initial_person_height_px {initial_person_height_px}")
182
+ PX_PER_METER = float(initial_person_height_px) / float(reference_height)
183
  print(f"Escala calculada: {PX_PER_METER:.2f} px/m")
184
+ if len(kpts_first[L_shoulder_keypoint]) > 0 and len(kpts_first[R_shoulder_keypoint]) > 0: # Left (5) and right (6) shoulders
185
+ initial_left_shoulder_x = int(kpts_first[L_shoulder_keypoint][0])
186
+ initial_right_shoulder_x = int(kpts_first[R_shoulder_keypoint][0])
187
 
188
  if PX_PER_METER is None or initial_left_shoulder_x is None or initial_right_shoulder_x is None:
189
  print("No se pudo calibrar la escala o detectar los hombros en el primer frame.")
 
222
  break
223
 
224
  annotated_frame = frame.copy()
 
225
 
226
+ # Add try-except block around the model inference to catch any model errors
227
+ try:
228
+ output = model(annotated_frame)
229
+ keypoints = output.keypoints_xy.float().cpu().numpy()
 
 
 
230
 
231
+ # Verify that keypoints array has valid data before processing
232
+ if (keypoints is not None and
233
+ len(keypoints) > 0 and
234
+ len(keypoints[0]) > 0 and
235
+ keypoints.size > 0): # Check if array is not empty
 
 
 
 
 
 
 
 
 
 
 
 
236
 
237
+ person_detected = True
238
+ kpts = keypoints[0]
 
 
 
 
 
 
 
239
 
240
+ # Make sure all required keypoints are detected
241
+ if (nose_keypoint < len(kpts) and L_ankle_keypoint < len(kpts) and
242
+ R_ankle_keypoint < len(kpts) and L_shoulder_keypoint < len(kpts) and
243
+ R_shoulder_keypoint < len(kpts)):
244
+
245
+ nose = kpts[nose_keypoint]
246
+ ankles = [kpts[L_ankle_keypoint], kpts[R_ankle_keypoint]]
247
+ left_shoulder = kpts[L_shoulder_keypoint]
248
+ right_shoulder = kpts[R_shoulder_keypoint]
249
+
250
+ # Check if keypoints have valid coordinates
251
+ if (nose[0] > 0 and nose[1] > 0 and
252
+ all(a[0] > 0 and a[1] > 0 for a in ankles) and
253
+ left_shoulder[0] > 0 and left_shoulder[1] > 0 and
254
+ right_shoulder[0] > 0 and right_shoulder[1] > 0):
255
+
256
+ # Continue with existing processing
257
+ current_ankle_y = min(a[1] for a in ankles)
258
+ last_detected_ankles_y = current_ankle_y
259
+ current_head_y = nose[1]
260
+ current_left_shoulder_x = int(left_shoulder[0])
261
+ current_right_shoulder_x = int(right_shoulder[0])
262
+
263
+ # Smooth ankle and head positions
264
+ ankle_y_history.append(current_ankle_y)
265
+ if len(ankle_y_history) > SMOOTHING_WINDOW:
266
+ ankle_y_history.pop(0)
267
+ smoothed_ankle_y = np.mean(ankle_y_history)
268
+
269
+ head_y_history.append(current_head_y)
270
+ if len(head_y_history) > SMOOTHING_WINDOW:
271
+ head_y_history.pop(0)
272
+ smoothed_head_y = np.mean(head_y_history)
273
+
274
+ # Calculate vertical velocity (using head position)
275
+ head_y_buffer.append(smoothed_head_y)
276
+ if len(head_y_buffer) > VELOCITY_WINDOW:
277
+ head_y_buffer.pop(0)
278
+ if PX_PER_METER is not None and fps > 0:
279
+ delta_y_pixels = head_y_buffer[0] - head_y_buffer[-1]
280
+ delta_y_meters = delta_y_pixels / PX_PER_METER
281
+ delta_t = VELOCITY_WINDOW / fps
282
+ velocity_vertical = delta_y_meters / delta_t
283
+
284
+ # Set ground level in first frame where ankles are detected
285
+ if ground_level is None:
286
+ ground_level = smoothed_ankle_y
287
+ takeoff_head_y = smoothed_head_y
288
+
289
+ relative_ankle_change = (ground_level - smoothed_ankle_y) / ground_level if ground_level > 0 else 0
290
+
291
+ # Detect jump start
292
+ if not jump_started and relative_ankle_change > JUMP_THRESHOLD_PERCENT:
293
+ jump_started = True
294
+ takeoff_head_y = smoothed_head_y
295
+ max_jump_height = 0
296
+ max_head_height_px = smoothed_head_y
297
+
298
+ # Detect jump end
299
+ if jump_started and relative_ankle_change <= JUMP_THRESHOLD_PERCENT:
300
+ # Add to repetition data
301
+ salto_alto = calculate_absolute_jump_height(reference_height, max_jump_height)
302
+ repetition_data.append({
303
+ "repetition": repetition_count + 1,
304
+ "relative_jump_m": round(max_jump_height, 3),
305
+ "absolute_jump_m": round(salto_alto, 3),
306
+ "peak_power_watts": round(current_power, 1)
307
+ })
308
+ repetition_count += 1
309
+ jump_started = False
310
+
311
+ # Update jump metrics while in air
312
+ if jump_started:
313
+ relative_jump = (takeoff_head_y - smoothed_head_y) / PX_PER_METER
314
+ if relative_jump > max_jump_height:
315
+ max_jump_height = relative_jump
316
+ if smoothed_head_y < max_head_height_px:
317
+ max_head_height_px = smoothed_head_y
318
+ if relative_jump:
319
+ current_power = calculate_peak_power_sayer(relative_jump, body_mass_kg)
320
+ if current_power > peak_power_sayer:
321
+ peak_power_sayer = current_power
322
+ else:
323
+ # Skip processing for this frame - invalid coordinates
324
+ print("Skipping frame - invalid keypoint coordinates")
325
+ print(f"keypoints {keypoints}")
326
+ else:
327
+ # Skip processing for this frame - missing required keypoints
328
+ print("Skipping frame - missing required keypoints")
329
+ print(f"keypoints {keypoints}")
330
+ else:
331
+ # Skip processing for this frame - no valid keypoints detected
332
+ print("Skipping frame - no valid keypoints detected")
333
+ print(f"keypoints {keypoints}")
334
+ last_detected_ankles_y = None
335
+ velocity_vertical = 0.0
336
+ except Exception as e:
337
+ # Handle any other exceptions that might occur during model inference
338
+ print(f"Error processing frame: {e}")
339
+ print(f"keypoints {keypoints}")
340
+ last_detected_ankles_y = None
341
+ velocity_vertical = 0.0
342
 
343
+ # Calculate metrics and draw overlay even if keypoints weren't detected
344
+ # This ensures video continues to show previous metrics
345
  salto_alto = calculate_absolute_jump_height(reference_height, max_jump_height)
346
 
347
  # Draw floating metric boxes
 
431
  Returns:
432
  Absolute jump height in meters
433
  """
434
+ absolute_jump = float(reference_height) + float(relative_jump)
435
  # Apply validation rule
436
  if absolute_jump > 1.72:
437
  return absolute_jump
vitpose.py CHANGED
@@ -17,7 +17,7 @@ class VitPose:
17
  object_detection_checkpoint="PekingU/rtdetr_r50vd_coco_o365",
18
  pose_estimation_checkpoint="usyd-community/vitpose-plus-small",
19
  device="cuda" if torch.cuda.is_available() else "cpu",
20
- dtype=torch.float16,
21
  compile=True, # or True to get more speedup
22
  )
23
  self.output_video_path = None
 
17
  object_detection_checkpoint="PekingU/rtdetr_r50vd_coco_o365",
18
  pose_estimation_checkpoint="usyd-community/vitpose-plus-small",
19
  device="cuda" if torch.cuda.is_available() else "cpu",
20
+ dtype=torch.bfloat16,
21
  compile=True, # or True to get more speedup
22
  )
23
  self.output_video_path = None