nicolasbuitragob commited on
Commit
d5bf444
·
1 Parent(s): 2a37377
Files changed (3) hide show
  1. app.py +38 -2
  2. tasks.py +517 -2
  3. vitpose.py +5 -2
app.py CHANGED
@@ -2,9 +2,10 @@ from fastapi import FastAPI, UploadFile, File, Response,Header, BackgroundTasks,
2
  from fastapi.staticfiles import StaticFiles
3
  from vitpose import VitPose
4
  from dotenv import load_dotenv
5
- from tasks import process_video
6
  from fastapi.responses import JSONResponse
7
  from config import AI_API_TOKEN
 
8
  import logging
9
  logging.basicConfig(level=logging.INFO)
10
  logger = logging.getLogger(__name__)
@@ -50,4 +51,39 @@ async def upload(background_tasks: BackgroundTasks,
50
 
51
  # Return the file as a response
52
  return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})
53
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
  from fastapi.staticfiles import StaticFiles
3
  from vitpose import VitPose
4
  from dotenv import load_dotenv
5
+ from tasks import process_video,process_salto_alto
6
  from fastapi.responses import JSONResponse
7
  from config import AI_API_TOKEN
8
+ from typing import Any
9
  import logging
10
  logging.basicConfig(level=logging.INFO)
11
  logger = logging.getLogger(__name__)
 
51
 
52
  # Return the file as a response
53
  return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})
54
+
55
+ @app.post("/exercise/salto_alto")
56
+ async def upload(background_tasks: BackgroundTasks,
57
+ file: UploadFile = File(...),
58
+ token: str = Header(...),
59
+ player_data: str = Body(...),
60
+ repetitions: int|str = Body(...),
61
+ exercise_id: str = Body(...)
62
+ ):
63
+
64
+ import json
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()
72
+
73
+ # Save the file to the local directory
74
+ logger.info("saving file")
75
+ with open(file.filename, "wb") as f:
76
+ f.write(contents)
77
+
78
+ logger.info(f"file saved {file.filename}")
79
+
80
+ # Create a clone of the file with content already read
81
+ background_tasks.add_task(process_salto_alto,
82
+ file.filename,
83
+ vitpose,
84
+ player_data,
85
+ repetitions,
86
+ exercise_id)
87
+
88
+ # Return the file as a response
89
+ return JSONResponse(content={"message": "Video uploaded successfully", "status": 200})
tasks.py CHANGED
@@ -4,9 +4,14 @@ import os
4
  from config import API_URL,API_KEY
5
  from fastapi import UploadFile
6
  import logging
 
 
 
 
 
 
7
  logging.basicConfig(level=logging.INFO)
8
  logger = logging.getLogger(__name__)
9
-
10
  def process_video(file_name: str,vitpose: VitPose,user_id: str,player_id: str):
11
 
12
  video_path = file_name
@@ -45,4 +50,514 @@ def process_video(file_name: str,vitpose: VitPose,user_id: str,player_id: str):
45
  logger.info(f"Response: {response.status_code}")
46
  logger.info(f"Response: {response.text}")
47
  logger.info(f"Video sent to {url}")
48
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
  from config import API_URL,API_KEY
5
  from fastapi import UploadFile
6
  import logging
7
+ import cv2
8
+ import numpy as np
9
+
10
+ import time
11
+ import json
12
+
13
  logging.basicConfig(level=logging.INFO)
14
  logger = logging.getLogger(__name__)
 
15
  def process_video(file_name: str,vitpose: VitPose,user_id: str,player_id: str):
16
 
17
  video_path = file_name
 
50
  logger.info(f"Response: {response.status_code}")
51
  logger.info(f"Response: {response.text}")
52
  logger.info(f"Video sent to {url}")
53
+
54
+
55
+ def process_salto_alto(file_name: str, vitpose: VitPose, player_data: dict, repetitions: int, exercise_id: str):
56
+ """
57
+ Process a high jump exercise video using VitPose for pose estimation.
58
+
59
+ Args:
60
+ file_name: Path to the input video
61
+ vitpose: VitPose instance for pose estimation
62
+ player_data: Dictionary containing player information
63
+ repetitions: Expected number of repetitions
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
70
+ reference_height = player_data.get('height', 1.68) # Altura aproximada de la persona en metros
71
+ body_mass_kg = player_data.get('weight', 64) # Peso corporal en kg
72
+
73
+ # Generate output paths
74
+ output_video = file_name.replace('.mp4', '_analyzed.mp4')
75
+ output_json = output_video.replace('.mp4', '.json')
76
+
77
+ # Process the video and get the jump metrics
78
+ results_dict = analyze_jump_video(
79
+ model=model,
80
+ input_video=file_name,
81
+ output_video=output_video,
82
+ reference_height=reference_height,
83
+ body_mass_kg=body_mass_kg
84
+ )
85
+
86
+ # Save results to JSON
87
+ with open(output_json, 'w') as f:
88
+ json.dumps(results_dict, indent=4)
89
+
90
+ # Print summary
91
+ print("\nResultados finales:")
92
+ print(f"Salto Relativo máximo: {results_dict['jump_metrics']['max_relative_jump']:.2f}m")
93
+ print(f"Salto Alto máximo: {results_dict['jump_metrics']['max_high_jump']:.2f}m")
94
+ print(f"Potencia Sayer (estimada): {results_dict['jump_metrics']['peak_power_sayer']:.2f} W")
95
+
96
+ # Return results dictionary
97
+ return {
98
+ "output_video": output_video,
99
+ "output_json": output_json,
100
+ "metrics": results_dict
101
+ }
102
+
103
+
104
+ def analyze_jump_video(model, input_video, output_video, reference_height=1.68, body_mass_kg=64):
105
+ """
106
+ Analyze a jump video to calculate various jump metrics.
107
+
108
+ Args:
109
+ model: VitPose model instance
110
+ input_video: Path to input video
111
+ output_video: Path to output video
112
+ reference_height: Height of the person in meters
113
+ body_mass_kg: Weight of the person in kg
114
+
115
+ Returns:
116
+ Dictionary containing jump metrics and video analysis data
117
+ """
118
+ # Configuration parameters
119
+ JUMP_THRESHOLD_PERCENT = 0.05 # Porcentaje de cambio en la altura del tobillo para detectar el inicio del salto
120
+ SMOOTHING_WINDOW = 5 # Ventana para suavizar la altura de los tobillos
121
+ HORIZONTAL_OFFSET_FACTOR = 0.75 # Factor para ubicar el cuadro entre el hombro y el borde
122
+ VELOCITY_WINDOW = 3 # Número de frames para calcular la velocidad
123
+ METRICS_BELOW_FEET_OFFSET = 20 # Offset en píxeles para colocar los cuadros debajo de los pies
124
+
125
+ # Color palette
126
+ BLUE = (255, 0, 0)
127
+ GREEN = (0, 255, 0)
128
+ YELLOW = (0, 255, 255)
129
+ WHITE = (255, 255, 255)
130
+ BLACK = (0, 0, 0)
131
+ GRAY = (128, 128, 128)
132
+ LIGHT_GRAY = (200, 200, 200)
133
+
134
+ repetition_data = []
135
+
136
+ # Open the video
137
+ cap = cv2.VideoCapture(input_video)
138
+ if not cap.isOpened():
139
+ print("Error al abrir el video")
140
+ return {}
141
+
142
+ # Get first frame to calibrate and get initial shoulder positions
143
+ ret, frame = cap.read()
144
+ if not ret:
145
+ print("Error al leer el video")
146
+ return {}
147
+
148
+ # Initialize calibration variables
149
+ PX_PER_METER = None
150
+ initial_person_height_px = None
151
+ initial_left_shoulder_x = None
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.")
168
+ cap.release()
169
+ return {}
170
+
171
+ # Reset video for processing
172
+ cap.release()
173
+ cap = cv2.VideoCapture(input_video)
174
+ fps = cap.get(cv2.CAP_PROP_FPS)
175
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
176
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
177
+ out = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
178
+
179
+ # Variables for metrics and visualization
180
+ ground_level = None
181
+ takeoff_head_y = None
182
+ max_jump_height = 0 # Maximum relative jump
183
+ max_head_height_px = None # Maximum head height in pixels (lowest in y coordinates)
184
+ jump_started = False
185
+ head_y_history = []
186
+ ankle_y_history = []
187
+ last_detected_ankles_y = None
188
+ head_y_buffer = []
189
+ velocity_vertical = 0.0
190
+ peak_power_sayer = 0.0 # Initialize Sayer power
191
+ person_detected = False # Flag to indicate if person was detected in any frame
192
+ current_power = 0.0
193
+ repetition_count = 0
194
+ jump_in_air = False
195
+
196
+ # Process each frame
197
+ while cap.isOpened():
198
+ ret, frame = cap.read()
199
+ if not ret:
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
287
+ annotated_frame = draw_metrics_overlay(
288
+ frame=annotated_frame,
289
+ max_jump_height=max_jump_height,
290
+ salto_alto=salto_alto,
291
+ velocity_vertical=velocity_vertical,
292
+ peak_power_sayer=peak_power_sayer,
293
+ repetition_count=repetition_count,
294
+ last_detected_ankles_y=last_detected_ankles_y,
295
+ initial_left_shoulder_x=initial_left_shoulder_x,
296
+ initial_right_shoulder_x=initial_right_shoulder_x,
297
+ width=width,
298
+ height=height,
299
+ colors={
300
+ "blue": BLUE,
301
+ "green": GREEN,
302
+ "yellow": YELLOW,
303
+ "white": WHITE,
304
+ "black": BLACK,
305
+ "gray": GRAY,
306
+ "light_gray": LIGHT_GRAY
307
+ },
308
+ metrics_below_feet_offset=METRICS_BELOW_FEET_OFFSET,
309
+ horizontal_offset_factor=HORIZONTAL_OFFSET_FACTOR
310
+ )
311
+
312
+ out.write(annotated_frame)
313
+
314
+ # Prepare results dictionary
315
+ results_dict = {
316
+ "jump_metrics": {
317
+ "max_relative_jump": float(max(0, max_jump_height)),
318
+ "max_high_jump": float(max(0, salto_alto)),
319
+ "peak_power_sayer": float(peak_power_sayer),
320
+ "repetitions": int(repetition_count),
321
+ "reference_height": float(reference_height),
322
+ "body_mass_kg": float(body_mass_kg),
323
+ "px_per_meter": float(PX_PER_METER) if PX_PER_METER is not None else 0.0
324
+ },
325
+ "video_analysis": {
326
+ "input_video": str(input_video),
327
+ "output_video": str(output_video),
328
+ "fps": float(fps),
329
+ "resolution": f"{int(width)}x{int(height)}"
330
+ },
331
+ "repetition_data": [
332
+ {
333
+ "repetition": int(rep["repetition"]),
334
+ "relative_jump_m": float(rep["relative_jump_m"]),
335
+ "absolute_jump_m": float(rep["absolute_jump_m"]),
336
+ "peak_power_watts": float(rep["peak_power_watts"])
337
+ } for rep in repetition_data
338
+ ]
339
+ }
340
+
341
+ cap.release()
342
+ out.release()
343
+
344
+ return results_dict
345
+
346
+
347
+ def calculate_peak_power_sayer(jump_height_m, body_mass_kg):
348
+ """
349
+ Estimates peak anaerobic power using Sayer's equation.
350
+
351
+ Args:
352
+ jump_height_m: Jump height in meters
353
+ body_mass_kg: Body mass in kg
354
+
355
+ Returns:
356
+ Estimated peak power in watts
357
+ """
358
+ jump_height_cm = jump_height_m * 100
359
+ return (60.7 * jump_height_cm) + (45.3 * body_mass_kg) - 2055
360
+
361
+
362
+ def calculate_absolute_jump_height(reference_height, relative_jump):
363
+ """
364
+ Calculate absolute jump height based on reference height and relative jump.
365
+
366
+ Args:
367
+ reference_height: Reference height in meters
368
+ relative_jump: Relative jump height in meters
369
+
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
377
+ else:
378
+ return 0
379
+
380
+
381
+ def draw_metrics_overlay(frame, max_jump_height, salto_alto, velocity_vertical, peak_power_sayer,
382
+ repetition_count, last_detected_ankles_y, initial_left_shoulder_x,
383
+ initial_right_shoulder_x, width, height, colors, metrics_below_feet_offset=20,
384
+ horizontal_offset_factor=0.75):
385
+ """
386
+ Draw metrics overlay on the frame.
387
+
388
+ Args:
389
+ frame: Input frame
390
+ max_jump_height: Maximum jump height in meters
391
+ salto_alto: Absolute jump height in meters
392
+ velocity_vertical: Vertical velocity in m/s
393
+ peak_power_sayer: Peak power in watts
394
+ repetition_count: Number of repetitions
395
+ last_detected_ankles_y: Y-coordinate of last detected ankles
396
+ initial_left_shoulder_x: X-coordinate of left shoulder
397
+ initial_right_shoulder_x: X-coordinate of right shoulder
398
+ width: Frame width
399
+ height: Frame height
400
+ colors: Dictionary with color values
401
+ metrics_below_feet_offset: Offset for metrics below feet
402
+ horizontal_offset_factor: Factor for horizontal offset
403
+
404
+ Returns:
405
+ Frame with metrics overlay
406
+ """
407
+ overlay = frame.copy()
408
+ alpha = 0.7
409
+ font = cv2.FONT_HERSHEY_SIMPLEX
410
+ font_scale_title_metric = 0.5
411
+ font_scale_value = 0.7
412
+ font_scale_title_main = 1.2 # Scale for main title (larger)
413
+ font_thickness_metric = 1
414
+ font_thickness_title_main = 1 # Thickness for main title
415
+ line_height_title_metric = int(20 * 1.2)
416
+ line_height_value = int(25 * 1.2)
417
+ padding_vertical = int(15 * 1.2)
418
+ padding_horizontal = int(15 * 1.2)
419
+ text_color_title = colors["light_gray"]
420
+ text_color_value = colors["white"]
421
+ text_color_title_main = colors["white"]
422
+ bg_color = colors["gray"]
423
+ border_color = colors["white"]
424
+ border_thickness = 1
425
+ corner_radius = 10
426
+ spacing_horizontal = 30
427
+ title_y_offset = 50 # Lower vertical position of title
428
+ metrics_y_offset_alto = 80 # Adjust Salto Alto position to leave space below
429
+ metrics_y_offset_relativo = None # Will be calculated dynamically
430
+ metrics_y_offset_velocidad = None # Will be calculated dynamically
431
+ metrics_y_offset_potencia = None # Will be calculated dynamically
432
+
433
+ # Helper function to draw rounded rectangles
434
+ def draw_rounded_rect(img, pt1, pt2, color, thickness=-1, lineType=cv2.LINE_AA, radius=10):
435
+ x1, y1 = pt1
436
+ x2, y2 = pt2
437
+ w = x2 - x1
438
+ h = y2 - y1
439
+ if radius > 0:
440
+ img = cv2.ellipse(img, (x1 + radius, y1 + radius), (radius, radius), 0, 0, 90, color, thickness, lineType)
441
+ img = cv2.ellipse(img, (x2 - radius, y1 + radius), (radius, radius), 0, 90, 180, color, thickness, lineType)
442
+ img = cv2.ellipse(img, (x2 - radius, y2 - radius), (radius, radius), 0, 180, 270, color, thickness, lineType)
443
+ img = cv2.ellipse(img, (x1 + radius, y2 - radius), (radius, radius), 0, 270, 360, color, thickness, lineType)
444
+
445
+ img = cv2.rectangle(img, (x1, y1 + radius), (x2, y2 - radius), color, thickness, lineType)
446
+ img = cv2.rectangle(img, (x1 + radius, y1), (x2 - radius, y2), color, thickness, lineType)
447
+ else:
448
+ img = cv2.rectangle(img, pt1, pt2, color, thickness, lineType)
449
+ return img
450
+
451
+ # --- Main Title ---
452
+ title_text = "Ejercicio de Salto"
453
+ title_text_size = cv2.getTextSize(title_text, font, font_scale_title_main, font_thickness_title_main)[0]
454
+ title_x = (width - title_text_size[0]) // 2
455
+ title_y = title_y_offset
456
+ cv2.putText(overlay, title_text, (title_x, title_y), font, font_scale_title_main, text_color_title_main, font_thickness_title_main, cv2.LINE_AA)
457
+
458
+ # --- Relative Jump Box (dynamically positioned) ---
459
+ relativo_text = "SALTO RELATIVO"
460
+ relativo_value = f"{max(0, max_jump_height):.2f} m"
461
+ relativo_text_size = cv2.getTextSize(relativo_text, font, font_scale_title_metric, font_thickness_metric)[0]
462
+ relativo_value_size = cv2.getTextSize(relativo_value, font, font_scale_value, font_thickness_metric)[0]
463
+ bg_width_relativo = max(relativo_text_size[0], relativo_value_size[0]) + 2 * padding_horizontal
464
+ bg_height_relativo = line_height_title_metric + line_height_value + 2 * padding_vertical
465
+ x_relativo = 20
466
+
467
+ if last_detected_ankles_y is not None and bg_height_relativo is not None:
468
+ metrics_y_offset_relativo = int(last_detected_ankles_y - bg_height_relativo - 10) # 10 pixels above ankle
469
+ # Make sure box doesn't go off top
470
+ if metrics_y_offset_relativo < title_y_offset + 50:
471
+ metrics_y_offset_relativo = int(last_detected_ankles_y + metrics_below_feet_offset) # Show below
472
+ else:
473
+ metrics_y_offset_relativo = height - 150 # Default position if ankles not detected
474
+
475
+ if metrics_y_offset_relativo is not None:
476
+ y_relativo = metrics_y_offset_relativo
477
+ pt1_relativo = (x_relativo, y_relativo)
478
+ pt2_relativo = (x_relativo + bg_width_relativo, y_relativo + bg_height_relativo)
479
+ overlay = draw_rounded_rect(overlay, pt1_relativo, pt2_relativo, bg_color, cv2.FILLED, cv2.LINE_AA, corner_radius)
480
+ cv2.rectangle(overlay, pt1_relativo, pt2_relativo, border_color, border_thickness, cv2.LINE_AA)
481
+ cv2.putText(overlay, relativo_text, (x_relativo + (bg_width_relativo - relativo_text_size[0]) // 2, y_relativo + padding_vertical + line_height_title_metric // 2 + 2), font, font_scale_title_metric, text_color_title, font_thickness_metric, cv2.LINE_AA)
482
+ cv2.putText(overlay, relativo_value, (x_relativo + (bg_width_relativo - relativo_value_size[0]) // 2, y_relativo + padding_vertical + line_height_title_metric + line_height_value // 2 + 5), font, font_scale_value, text_color_value, font_thickness_metric, cv2.LINE_AA)
483
+
484
+ # --- High Jump Box (stays in top right) ---
485
+ alto_text = "SALTO ALTO"
486
+ alto_value = f"{max(0, salto_alto):.2f} m"
487
+ alto_text_size = cv2.getTextSize(alto_text, font, font_scale_title_metric, font_thickness_metric)[0]
488
+ alto_value_size = cv2.getTextSize(alto_value, font, font_scale_value, font_thickness_metric)[0]
489
+ bg_width_alto = max(alto_text_size[0], alto_value_size[0]) + 2 * padding_horizontal
490
+ bg_height_alto = line_height_title_metric + line_height_value + 2 * padding_vertical
491
+ x_alto = width - bg_width_alto - 20 # Default position near right edge
492
+
493
+ if initial_right_shoulder_x is not None:
494
+ available_space = width - initial_right_shoulder_x
495
+ x_alto_calculated = initial_right_shoulder_x + int(available_space * (1 - horizontal_offset_factor)) - bg_width_alto
496
+ # Make sure doesn't go off left edge and there's space from first box
497
+ if x_alto_calculated > x_relativo + bg_width_relativo + spacing_horizontal + 10 and x_alto_calculated + bg_width_alto < width - 10:
498
+ x_alto = x_alto_calculated
499
+ y_alto = metrics_y_offset_alto
500
+ pt1_alto = (x_alto, y_alto)
501
+ pt2_alto = (x_alto + bg_width_alto, y_alto + bg_height_alto)
502
+ overlay = draw_rounded_rect(overlay, pt1_alto, pt2_alto, bg_color, cv2.FILLED, cv2.LINE_AA, corner_radius)
503
+ cv2.rectangle(overlay, pt1_alto, pt2_alto, border_color, border_thickness, cv2.LINE_AA)
504
+ cv2.putText(overlay, alto_text, (x_alto + (bg_width_alto - alto_text_size[0]) // 2, y_alto + padding_vertical + line_height_title_metric // 2 + 2), font, font_scale_title_metric, text_color_title, font_thickness_metric, cv2.LINE_AA)
505
+ cv2.putText(overlay, alto_value, (x_alto + (bg_width_alto - alto_value_size[0]) // 2, y_alto + padding_vertical + line_height_title_metric + line_height_value // 2 + 5), font, font_scale_value, text_color_value, font_thickness_metric, cv2.LINE_AA)
506
+
507
+ # --- Repetitions Box ---
508
+ reps_text = "REPETICIONES"
509
+ reps_value = f"{repetition_count}"
510
+ reps_text_size = cv2.getTextSize(reps_text, font, font_scale_title_metric, font_thickness_metric)[0]
511
+ reps_value_size = cv2.getTextSize(reps_value, font, font_scale_value, font_thickness_metric)[0]
512
+ bg_width_reps = max(reps_text_size[0], reps_value_size[0]) + 2 * padding_horizontal
513
+ bg_height_reps = line_height_title_metric + line_height_value + 2 * padding_vertical
514
+ x_reps = x_relativo
515
+ y_reps = y_relativo + bg_height_relativo + 10
516
+
517
+ pt1_reps = (x_reps, y_reps)
518
+ pt2_reps = (x_reps + bg_width_reps, y_reps + bg_height_reps)
519
+ overlay = draw_rounded_rect(overlay, pt1_reps, pt2_reps, bg_color, cv2.FILLED, cv2.LINE_AA, corner_radius)
520
+ cv2.rectangle(overlay, pt1_reps, pt2_reps, border_color, border_thickness, cv2.LINE_AA)
521
+ cv2.putText(overlay, reps_text, (x_reps + (bg_width_reps - reps_text_size[0]) // 2, y_reps + padding_vertical + line_height_title_metric // 2 + 2), font, font_scale_title_metric, text_color_title, font_thickness_metric, cv2.LINE_AA)
522
+ cv2.putText(overlay, reps_value, (x_reps + (bg_width_reps - reps_value_size[0]) // 2, y_reps + padding_vertical + line_height_title_metric + line_height_value // 2 + 5), font, font_scale_value, text_color_value, font_thickness_metric, cv2.LINE_AA)
523
+
524
+ # --- Vertical Velocity Box (below feet) ---
525
+ if last_detected_ankles_y is not None:
526
+ velocidad_text = "VELOCIDAD VERTICAL"
527
+ velocidad_value = f"{abs(velocity_vertical):.2f} m/s" # Show absolute value
528
+ velocidad_text_size = cv2.getTextSize(velocidad_text, font, font_scale_title_metric, font_thickness_metric)[0]
529
+ velocidad_value_size = cv2.getTextSize(velocidad_value, font, font_scale_value, font_thickness_metric)[0]
530
+ bg_width_velocidad = max(velocidad_text_size[0], velocidad_value_size[0]) + 2 * padding_horizontal
531
+ bg_height_velocidad = line_height_title_metric + line_height_value + 2 * padding_vertical
532
+
533
+ x_velocidad = int(width / 2 - bg_width_velocidad / 2) # Horizontally centered
534
+ y_velocidad = int(last_detected_ankles_y + metrics_below_feet_offset + bg_height_velocidad)
535
+
536
+ pt1_velocidad = (int(x_velocidad), int(y_velocidad - bg_height_velocidad))
537
+ pt2_velocidad = (int(x_velocidad + bg_width_velocidad), int(y_velocidad))
538
+ overlay = draw_rounded_rect(overlay, pt1_velocidad, pt2_velocidad, bg_color, cv2.FILLED, cv2.LINE_AA, corner_radius)
539
+ cv2.rectangle(overlay, pt1_velocidad, pt2_velocidad, border_color, border_thickness, cv2.LINE_AA)
540
+ cv2.putText(overlay, velocidad_text, (int(x_velocidad + (bg_width_velocidad - velocidad_text_size[0]) // 2), int(y_velocidad - bg_height_velocidad + padding_vertical + line_height_title_metric // 2 + 2)), font, font_scale_title_metric, text_color_title, font_thickness_metric, cv2.LINE_AA)
541
+ cv2.putText(overlay, velocidad_value, (int(x_velocidad + (bg_width_velocidad - velocidad_value_size[0]) // 2), int(y_velocidad - bg_height_velocidad + padding_vertical + line_height_title_metric + line_height_value // 2 + 5)), font, font_scale_value, text_color_value, font_thickness_metric, cv2.LINE_AA)
542
+
543
+ # --- Sayer Power Box (below velocity box) ---
544
+ potencia_text = "POTENCIA SAYER"
545
+ potencia_value = f"{peak_power_sayer:.2f} W"
546
+ potencia_text_size = cv2.getTextSize(potencia_text, font, font_scale_title_metric, font_thickness_metric)[0]
547
+ potencia_value_size = cv2.getTextSize(potencia_value, font, font_scale_value, font_thickness_metric)[0]
548
+ bg_width_potencia = max(potencia_text_size[0], potencia_value_size[0]) + 2 * padding_horizontal
549
+ bg_height_potencia = line_height_title_metric + line_height_value + 2 * padding_vertical
550
+
551
+ x_potencia = x_velocidad # Same horizontal position as velocity
552
+ y_potencia = y_velocidad + 5 # Below velocity box
553
+
554
+ pt1_potencia = (int(x_potencia), int(y_potencia))
555
+ pt2_potencia = (int(x_potencia + bg_width_potencia), int(y_potencia + bg_height_potencia))
556
+ overlay = draw_rounded_rect(overlay, pt1_potencia, pt2_potencia, bg_color, cv2.FILLED, cv2.LINE_AA, corner_radius)
557
+ cv2.rectangle(overlay, pt1_potencia, pt2_potencia, border_color, border_thickness, cv2.LINE_AA)
558
+ cv2.putText(overlay, potencia_text, (int(x_potencia + (bg_width_potencia - potencia_text_size[0]) // 2), int(y_potencia + padding_vertical + line_height_title_metric // 2 + 2)), font, font_scale_title_metric, text_color_title, font_thickness_metric, cv2.LINE_AA)
559
+ cv2.putText(overlay, potencia_value, (int(x_potencia + (bg_width_potencia - potencia_value_size[0]) // 2), int(y_potencia + padding_vertical + line_height_title_metric + line_height_value // 2 + 5)), font, font_scale_value, text_color_value, font_thickness_metric, cv2.LINE_AA)
560
+
561
+ # Blend overlay with original frame
562
+ result = cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
563
+ return result
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.bfloat16,
21
  compile=True, # or True to get more speedup
22
  )
23
  self.output_video_path = None
@@ -104,4 +104,7 @@ class VitPose:
104
  else:
105
  # Already vertical, no rotation needed
106
  out.write(frame)
107
- out.release()
 
 
 
 
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
 
104
  else:
105
  # Already vertical, no rotation needed
106
  out.write(frame)
107
+ out.release()
108
+
109
+
110
+