iDrops/Exercise_Tracking_Demo
0
1import time2import cv23import numpy as np4from utils import find_angle, get_landmark_features, draw_text5 6 7class ProcessFrame:8 def __init__(self, thresholds, flip_frame = False):9 10 # Set if frame should be flipped or not.11 self.flip_frame = flip_frame12 13 # self.thresholds14 self.thresholds = thresholds15 # Font type.16 self.font = cv2.FONT_HERSHEY_SIMPLEX17 # line type18 self.linetype = cv2.LINE_AA19 # set radius to draw arc20 self.radius = 20 # Radius for drawing arcs21 # Colors in BGR format.22 self.COLORS = {23 'blue' : (0, 127, 255),24 'red' : (255, 50, 50),25 'green' : (0, 255, 127),26 'light_green': (100, 233, 127),27 'orange' : (252, 83, 5),28 'magenta' : (255, 0, 255),29 'white' : (255,255,255),30 'cyan' : (0, 255, 255),31 'off_white' : (221, 226, 237)32 }33 34 # Dictionary to maintain the various landmark features.35 self.dict_features = {} # Dictionary to store various landmark features based on body part names36 # Define landmark feature names for left and right body sides37 self.left_features = {38 'shoulder': 11,39 'elbow' : 13,40 'wrist' : 15, 41 'hip' : 23,42 'knee' : 25,43 'ankle' : 27,44 'foot' : 3145 }46 47 self.right_features = {48 'shoulder': 12,49 'elbow' : 14,50 'wrist' : 16,51 'hip' : 24,52 'knee' : 26,53 'ankle' : 28,54 'foot' : 3255 }56 57 # Assign landmark feature dictionaries for left, right body sides and nose to main dictionary58 self.dict_features['left'] = self.left_features59 self.dict_features['right'] = self.right_features60 self.dict_features['nose'] = 061 62 # For tracking counters and sharing states in and out of callbacks.63 self.state_tracker = {64 'state_seq': [], # Sequence of detected states65 66 'start_inactive_time': time.perf_counter(), # Timer for front side inactivity67 'start_inactive_time_front': time.perf_counter(), # Timer for side view inactivity68 'INACTIVE_TIME': 0.0, # Total front side inactivity time69 'INACTIVE_TIME_FRONT': 0.0, # Total side view inactivity time70 71 # 0 --> Bend Backwards, 1 --> Bend Forward, 2 --> Keep shin straight, 3 --> Deep squat72 'DISPLAY_TEXT' : np.full((4,), False),73 'COUNT_FRAMES' : np.zeros((4,), dtype=np.int64),74 75 'LOWER_HIPS': False,76 77 'INCORRECT_POSTURE': False,78 79 'prev_state': None,80 'curr_state':None,81 82 'SQUAT_COUNT': 0,83 'IMPROPER_SQUAT':084 85 }86 87 self.FEEDBACK_ID_MAP = {88 0: ('BEND BACKWARDS', 215, (0, 153, 255)),89 1: ('BEND FORWARD', 215, (0, 153, 255)),90 2: ('KNEE FALLING OVER TOE', 170, (255, 80, 80)),91 3: ('SQUAT TOO DEEP', 125, (255, 80, 80))92 }93 94 def _get_state(self, knee_angle):95 96 knee = None 97 98 if self.thresholds['HIP_KNEE_VERT']['NORMAL'][0] <= knee_angle <= self.thresholds['HIP_KNEE_VERT']['NORMAL'][1]:99 knee = 1100 elif self.thresholds['HIP_KNEE_VERT']['TRANS'][0] <= knee_angle <= self.thresholds['HIP_KNEE_VERT']['TRANS'][1]:101 knee = 2102 elif self.thresholds['HIP_KNEE_VERT']['PASS'][0] <= knee_angle <= self.thresholds['HIP_KNEE_VERT']['PASS'][1]:103 knee = 3104 105 return f's{knee}' if knee else None106 107 def _update_state_sequence(self, state):108 109 if state == 's2':110 if (('s3' not in self.state_tracker['state_seq']) and (self.state_tracker['state_seq'].count('s2'))==0) or \111 (('s3' in self.state_tracker['state_seq']) and (self.state_tracker['state_seq'].count('s2')==1)):112 self.state_tracker['state_seq'].append(state)113 114 115 elif state == 's3':116 if (state not in self.state_tracker['state_seq']) and 's2' in self.state_tracker['state_seq']: 117 self.state_tracker['state_seq'].append(state)118 119 def _show_feedback(self, frame, c_frame, dict_maps, lower_hips_disp):120 121 if lower_hips_disp:122 draw_text(123 frame, 124 'LOWER YOUR HIPS', 125 pos=(30, 80),126 text_color=(0, 0, 0),127 font_scale=0.6,128 text_color_bg=(255, 255, 0)129 ) 130 131 for idx in np.where(c_frame)[0]:132 draw_text(133 frame, 134 dict_maps[idx][0], 135 pos=(30, dict_maps[idx][1]),136 text_color=(255, 255, 230),137 font_scale=0.6,138 text_color_bg=dict_maps[idx][2]139 )140 141 return frame142 143 def process(self, frame: np.array, pose):144 play_sound = None145 146 147 frame_height, frame_width, _ = frame.shape148 149 # Process the image.150 keypoints = pose.process(frame)151 152 if keypoints.pose_landmarks:153 ps_lm = keypoints.pose_landmarks154 155 nose_coord = get_landmark_features(ps_lm.landmark, self.dict_features, 'nose', frame_width, frame_height)156 left_shldr_coord, left_elbow_coord, left_wrist_coord, left_hip_coord, left_knee_coord, left_ankle_coord, left_foot_coord = \157 get_landmark_features(ps_lm.landmark, self.dict_features, 'left', frame_width, frame_height)158 right_shldr_coord, right_elbow_coord, right_wrist_coord, right_hip_coord, right_knee_coord, right_ankle_coord, right_foot_coord = \159 get_landmark_features(ps_lm.landmark, self.dict_features, 'right', frame_width, frame_height)160 161 offset_angle = find_angle(left_shldr_coord, right_shldr_coord, nose_coord)162 163 if offset_angle > self.thresholds['OFFSET_THRESH']:164 165 display_inactivity = False166 167 end_time = time.perf_counter()168 self.state_tracker['INACTIVE_TIME_FRONT'] += end_time - self.state_tracker['start_inactive_time_front']169 self.state_tracker['start_inactive_time_front'] = end_time170 171 if self.state_tracker['INACTIVE_TIME_FRONT'] >= self.thresholds['INACTIVE_THRESH']:172 self.state_tracker['SQUAT_COUNT'] = 0173 self.state_tracker['IMPROPER_SQUAT'] = 0174 display_inactivity = True175 176 cv2.circle(frame, nose_coord, 7, self.COLORS['white'], -1)177 cv2.circle(frame, left_shldr_coord, 7, self.COLORS['orange'], -1)178 cv2.circle(frame, right_shldr_coord, 7, self.COLORS['magenta'], -1)179 180 if self.flip_frame:181 frame = cv2.flip(frame, 1)182 183 if display_inactivity:184 # cv2.putText(frame, 'Resetting SQUAT_COUNT due to inactivity!!!', (10, frame_height - 90), 185 # self.font, 0.5, self.COLORS['blue'], 2, lineType=self.linetype)186 play_sound = 'reset_counters'187 self.state_tracker['INACTIVE_TIME_FRONT'] = 0.0188 self.state_tracker['start_inactive_time_front'] = time.perf_counter()189 190 draw_text(191 frame, 192 "CORRECT: " + str(self.state_tracker['SQUAT_COUNT']), 193 pos=(int(frame_width*0.75), 30),194 text_color=(255, 255, 230),195 font_scale=0.7,196 text_color_bg=(18, 185, 0)197 ) 198 199 draw_text(200 frame, 201 "INCORRECT: " + str(self.state_tracker['IMPROPER_SQUAT']), 202 pos=(int(frame_width*0.75), 80),203 text_color=(255, 255, 230),204 font_scale=0.7,205 text_color_bg=(221, 0, 0),206 207 ) 208 209 draw_text(210 frame, 211 'CAMERA NOT ALIGNED PROPERLY!!!', 212 pos=(30, frame_height-60),213 text_color=(255, 255, 230),214 font_scale=0.65,215 text_color_bg=(237, 33, 37),216 ) 217 218 draw_text(219 frame, 220 'OFFSET ANGLE: '+str(offset_angle), 221 pos=(30, frame_height-30),222 text_color=(255, 255, 230),223 font_scale=0.65,224 text_color_bg=(255, 153, 0),225 ) 226 227 # Reset inactive times for side view.228 self.state_tracker['start_inactive_time'] = time.perf_counter()229 self.state_tracker['INACTIVE_TIME'] = 0.0230 self.state_tracker['prev_state'] = None231 self.state_tracker['curr_state'] = None232 233 # Camera is aligned properly.234 else:235 236 self.state_tracker['INACTIVE_TIME_FRONT'] = 0.0237 self.state_tracker['start_inactive_time_front'] = time.perf_counter()238 239 240 dist_l_sh_hip = abs(left_foot_coord[1]- left_shldr_coord[1])241 dist_r_sh_hip = abs(right_foot_coord[1] - right_shldr_coord)[1]242 243 shldr_coord = None244 elbow_coord = None245 wrist_coord = None246 hip_coord = None247 knee_coord = None248 ankle_coord = None249 foot_coord = None250 251 if dist_l_sh_hip > dist_r_sh_hip:252 shldr_coord = left_shldr_coord253 elbow_coord = left_elbow_coord254 wrist_coord = left_wrist_coord255 hip_coord = left_hip_coord256 knee_coord = left_knee_coord257 ankle_coord = left_ankle_coord258 foot_coord = left_foot_coord259 260 multiplier = -1261 262 263 else:264 shldr_coord = right_shldr_coord265 elbow_coord = right_elbow_coord266 wrist_coord = right_wrist_coord267 hip_coord = right_hip_coord268 knee_coord = right_knee_coord269 ankle_coord = right_ankle_coord270 foot_coord = right_foot_coord271 272 multiplier = 1273 274 275 # ------------------- Verical Angle calculation --------------276 277 hip_vertical_angle = find_angle(shldr_coord, np.array([hip_coord[0], 0]), hip_coord)278 cv2.ellipse(frame, hip_coord, (30, 30), 279 angle = 0, startAngle = -90, endAngle = -90+multiplier*hip_vertical_angle, 280 color = self.COLORS['white'], thickness = 3, lineType = self.linetype)281 282 cv2.line(frame, (hip_coord[0], hip_coord[1] + 20), (hip_coord[0], hip_coord[1] - 80), self.COLORS['blue'], 4, lineType=self.linetype)283 284 285 286 287 knee_vertical_angle = find_angle(hip_coord, np.array([knee_coord[0], 0]), knee_coord)288 cv2.ellipse(frame, knee_coord, (20, 20), 289 angle = 0, startAngle = -90, endAngle = -90-multiplier*knee_vertical_angle, 290 color = self.COLORS['white'], thickness = 3, lineType = self.linetype)291 292 cv2.line(frame, (knee_coord[0], knee_coord[1] + 20), (knee_coord[0], knee_coord[1] - 50), self.COLORS['blue'], 4, lineType=self.linetype)293 294 295 296 ankle_vertical_angle = find_angle(knee_coord, np.array([ankle_coord[0], 0]), ankle_coord)297 cv2.ellipse(frame, ankle_coord, (30, 30),298 angle = 0, startAngle = -90, endAngle = -90 + multiplier*ankle_vertical_angle,299 color = self.COLORS['white'], thickness = 3, lineType=self.linetype)300 301 cv2.line(frame, (ankle_coord[0], ankle_coord[1] + 20), (ankle_coord[0], ankle_coord[1] - 50), self.COLORS['blue'], 4, lineType=self.linetype)302 303 # ------------------------------------------------------------304 305 306 # Join landmarks.307 cv2.line(frame, shldr_coord, elbow_coord, self.COLORS['off_white'], 4, lineType=self.linetype)308 cv2.line(frame, wrist_coord, elbow_coord, self.COLORS['off_white'], 4, lineType=self.linetype)309 cv2.line(frame, shldr_coord, hip_coord, self.COLORS['off_white'], 4, lineType=self.linetype)310 cv2.line(frame, knee_coord, hip_coord, self.COLORS['off_white'], 4, lineType=self.linetype)311 cv2.line(frame, ankle_coord, knee_coord,self.COLORS['off_white'], 4, lineType=self.linetype)312 cv2.line(frame, ankle_coord, foot_coord, self.COLORS['off_white'], 4, lineType=self.linetype)313 314 # Plot landmark points315 cv2.circle(frame, shldr_coord, 7, self.COLORS['orange'], -1, lineType=self.linetype)316 cv2.circle(frame, elbow_coord, 7, self.COLORS['orange'], -1, lineType=self.linetype)317 cv2.circle(frame, wrist_coord, 7, self.COLORS['orange'], -1, lineType=self.linetype)318 cv2.circle(frame, hip_coord, 7, self.COLORS['orange'], -1, lineType=self.linetype)319 cv2.circle(frame, knee_coord, 7, self.COLORS['orange'], -1, lineType=self.linetype)320 cv2.circle(frame, ankle_coord, 7, self.COLORS['orange'], -1, lineType=self.linetype)321 cv2.circle(frame, foot_coord, 7, self.COLORS['orange'], -1, lineType=self.linetype)322 323 324 325 current_state = self._get_state(int(knee_vertical_angle))326 self.state_tracker['curr_state'] = current_state327 self._update_state_sequence(current_state)328 329 330 331 # -------------------------------------- COMPUTE COUNTERS --------------------------------------332 333 if current_state == 's1':334 335 if len(self.state_tracker['state_seq']) == 3 and not self.state_tracker['INCORRECT_POSTURE']:336 self.state_tracker['SQUAT_COUNT']+=1337 play_sound = str(self.state_tracker['SQUAT_COUNT'])338 339 elif 's2' in self.state_tracker['state_seq'] and len(self.state_tracker['state_seq'])==1:340 self.state_tracker['IMPROPER_SQUAT']+=1341 play_sound = 'incorrect'342 343 elif self.state_tracker['INCORRECT_POSTURE']:344 self.state_tracker['IMPROPER_SQUAT']+=1345 play_sound = 'incorrect'346 347 348 self.state_tracker['state_seq'] = []349 self.state_tracker['INCORRECT_POSTURE'] = False350 351 352 # ----------------------------------------------------------------------------------------------------353 354 355 356 357 # -------------------------------------- PERFORM FEEDBACK ACTIONS --------------------------------------358 359 else:360 if hip_vertical_angle > self.thresholds['HIP_THRESH'][1]:361 self.state_tracker['DISPLAY_TEXT'][0] = True362 363 364 elif hip_vertical_angle < self.thresholds['HIP_THRESH'][0] and \365 self.state_tracker['state_seq'].count('s2')==1:366 self.state_tracker['DISPLAY_TEXT'][1] = True367 368 369 370 if self.thresholds['KNEE_THRESH'][0] < knee_vertical_angle < self.thresholds['KNEE_THRESH'][1] and \371 self.state_tracker['state_seq'].count('s2')==1:372 self.state_tracker['LOWER_HIPS'] = True373 374 375 elif knee_vertical_angle > self.thresholds['KNEE_THRESH'][2]:376 self.state_tracker['DISPLAY_TEXT'][3] = True377 self.state_tracker['INCORRECT_POSTURE'] = True378 379 380 if (ankle_vertical_angle > self.thresholds['ANKLE_THRESH']):381 self.state_tracker['DISPLAY_TEXT'][2] = True382 self.state_tracker['INCORRECT_POSTURE'] = True383 384 385 # ----------------------------------------------------------------------------------------------------386 387 388 389 390 # ----------------------------------- COMPUTE INACTIVITY ---------------------------------------------391 392 display_inactivity = False393 394 if self.state_tracker['curr_state'] == self.state_tracker['prev_state']:395 396 end_time = time.perf_counter()397 self.state_tracker['INACTIVE_TIME'] += end_time - self.state_tracker['start_inactive_time']398 self.state_tracker['start_inactive_time'] = end_time399 400 if self.state_tracker['INACTIVE_TIME'] >= self.thresholds['INACTIVE_THRESH']:401 self.state_tracker['SQUAT_COUNT'] = 0402 self.state_tracker['IMPROPER_SQUAT'] = 0403 display_inactivity = True404 405 406 else:407 408 self.state_tracker['start_inactive_time'] = time.perf_counter()409 self.state_tracker['INACTIVE_TIME'] = 0.0410 411 # -------------------------------------------------------------------------------------------------------412 413 414 415 hip_text_coord_x = hip_coord[0] + 10416 knee_text_coord_x = knee_coord[0] + 15417 ankle_text_coord_x = ankle_coord[0] + 10418 419 if self.flip_frame:420 frame = cv2.flip(frame, 1)421 hip_text_coord_x = frame_width - hip_coord[0] + 10422 knee_text_coord_x = frame_width - knee_coord[0] + 15423 ankle_text_coord_x = frame_width - ankle_coord[0] + 10424 425 426 427 if 's3' in self.state_tracker['state_seq']:428 self.state_tracker['LOWER_HIPS'] = False429 430 self.state_tracker['COUNT_FRAMES'][self.state_tracker['DISPLAY_TEXT']]+=1431 432 frame = self._show_feedback(frame, self.state_tracker['COUNT_FRAMES'], self.FEEDBACK_ID_MAP, self.state_tracker['LOWER_HIPS'])433 434 435 436 if display_inactivity:437 # cv2.putText(frame, 'Resetting COUNTERS due to inactivity!!!', (10, frame_height - 20), self.font, 0.5, self.COLORS['blue'], 2, lineType=self.linetype)438 play_sound = 'reset_counters'439 self.state_tracker['start_inactive_time'] = time.perf_counter()440 self.state_tracker['INACTIVE_TIME'] = 0.0441 442 443 cv2.putText(frame, str(int(hip_vertical_angle)), (hip_text_coord_x, hip_coord[1]), self.font, 0.6, self.COLORS['light_green'], 2, lineType=self.linetype)444 cv2.putText(frame, str(int(knee_vertical_angle)), (knee_text_coord_x, knee_coord[1]+10), self.font, 0.6, self.COLORS['light_green'], 2, lineType=self.linetype)445 cv2.putText(frame, str(int(ankle_vertical_angle)), (ankle_text_coord_x, ankle_coord[1]), self.font, 0.6, self.COLORS['light_green'], 2, lineType=self.linetype)446 447 448 draw_text(449 frame, 450 "CORRECT: " + str(self.state_tracker['SQUAT_COUNT']), 451 pos=(int(frame_width*0.75), 30),452 text_color=(255, 255, 230),453 font_scale=0.7,454 text_color_bg=(18, 185, 0)455 ) 456 457 458 draw_text(459 frame, 460 "INCORRECT: " + str(self.state_tracker['IMPROPER_SQUAT']), 461 pos=(int(frame_width*0.75), 80),462 text_color=(255, 255, 230),463 font_scale=0.7,464 text_color_bg=(221, 0, 0),465 466 ) 467 468 469 self.state_tracker['DISPLAY_TEXT'][self.state_tracker['COUNT_FRAMES'] > self.thresholds['CNT_FRAME_THRESH']] = False470 self.state_tracker['COUNT_FRAMES'][self.state_tracker['COUNT_FRAMES'] > self.thresholds['CNT_FRAME_THRESH']] = 0 471 self.state_tracker['prev_state'] = current_state472 473 474 475 476 else:477 478 if self.flip_frame:479 frame = cv2.flip(frame, 1)480 481 end_time = time.perf_counter()482 self.state_tracker['INACTIVE_TIME'] += end_time - self.state_tracker['start_inactive_time']483 484 display_inactivity = False485 486 if self.state_tracker['INACTIVE_TIME'] >= self.thresholds['INACTIVE_THRESH']:487 self.state_tracker['SQUAT_COUNT'] = 0488 self.state_tracker['IMPROPER_SQUAT'] = 0489 # cv2.putText(frame, 'Resetting SQUAT_COUNT due to inactivity!!!', (10, frame_height - 25), self.font, 0.7, self.COLORS['blue'], 2)490 display_inactivity = True491 492 self.state_tracker['start_inactive_time'] = end_time493 494 draw_text(495 frame, 496 "CORRECT: " + str(self.state_tracker['SQUAT_COUNT']), 497 pos=(int(frame_width*0.75), 30),498 text_color=(255, 255, 230),499 font_scale=0.7,500 text_color_bg=(18, 185, 0)501 ) 502 503 504 draw_text(505 frame, 506 "INCORRECT: " + str(self.state_tracker['IMPROPER_SQUAT']), 507 pos=(int(frame_width*0.75), 80),508 text_color=(255, 255, 230),509 font_scale=0.7,510 text_color_bg=(221, 0, 0),511 512 ) 513 514 if display_inactivity:515 play_sound = 'reset_counters'516 self.state_tracker['start_inactive_time'] = time.perf_counter()517 self.state_tracker['INACTIVE_TIME'] = 0.0518 519 520 # Reset all other state variables521 522 self.state_tracker['prev_state'] = None523 self.state_tracker['curr_state'] = None524 self.state_tracker['INACTIVE_TIME_FRONT'] = 0.0525 self.state_tracker['INCORRECT_POSTURE'] = False526 self.state_tracker['DISPLAY_TEXT'] = np.full((5,), False)527 self.state_tracker['COUNT_FRAMES'] = np.zeros((5,), dtype=np.int64)528 self.state_tracker['start_inactive_time_front'] = time.perf_counter()529 530 531 532 return frame, play_sound