Spaces:
Sleeping
Sleeping
| import cv2 | |
| import mediapipe as mp | |
| import numpy as np | |
| # Load the correct and incorrect posture images as BGR colors | |
| correct = cv2.imread('right.png') | |
| correct = cv2.cvtColor(correct, cv2.COLOR_BGR2RGB) | |
| incorrect = cv2.imread('wrong.png') | |
| incorrect = cv2.cvtColor(incorrect, cv2.COLOR_BGR2RGB) | |
| def draw_rounded_rect(img, rect_start, rect_end, corner_width, box_color): | |
| # This function draws a rectangle with rounded corners on an image. | |
| # Args: | |
| #img: The image to draw on. | |
| #rect_start: The top-left corner of the rectangle as a tuple (x1, y1). | |
| #rect_end: The bottom-right corner of the rectangle as a tuple (x2, y2). | |
| #corner_width: The width of the rounded corners. | |
| #box_color: The color of the rectangle in BGR format. | |
| x1, y1 = rect_start | |
| x2, y2 = rect_end | |
| w = corner_width | |
| # Draw filled rectangles for each side of the box | |
| cv2.rectangle(img, (x1 + w, y1), (x2 - w, y1 + w), box_color, -1) | |
| cv2.rectangle(img, (x1 + w, y2 - w), (x2 - w, y2), box_color, -1) | |
| cv2.rectangle(img, (x1, y1 + w), (x1 + w, y2 - w), box_color, -1) | |
| cv2.rectangle(img, (x2 - w, y1 + w), (x2, y2 - w), box_color, -1) | |
| cv2.rectangle(img, (x1 + w, y1 + w), (x2 - w, y2 - w), box_color, -1) | |
| # Draw filled ellipses for the corners | |
| cv2.ellipse(img, (x1 + w, y1 + w), (w, w), | |
| angle = 0, startAngle = -90, endAngle = -180, color = box_color, thickness = -1) | |
| cv2.ellipse(img, (x2 - w, y1 + w), (w, w), | |
| angle = 0, startAngle = 0, endAngle = -90, color = box_color, thickness = -1) | |
| cv2.ellipse(img, (x1 + w, y2 - w), (w, w), | |
| angle = 0, startAngle = 90, endAngle = 180, color = box_color, thickness = -1) | |
| cv2.ellipse(img, (x2 - w, y2 - w), (w, w), | |
| angle = 0, startAngle = 0, endAngle = 90, color = box_color, thickness = -1) | |
| return img | |
| def draw_dotted_line(frame, lm_coord, start, end, line_color): | |
| #This function draws a dotted line on a frame based on landmark coordinates. | |
| #Args: | |
| #frame: The image to draw on. | |
| #lm_coord: The landmark coordinates as a NumPy array. | |
| #start: The index of the starting landmark in the lm_coord array. | |
| #end: The index of the ending landmark in the lm_coord array. | |
| #line_color: The color of the line in BGR format. | |
| pix_step = 0 | |
| # Draw circles at every 8th element between the start and end landmarks | |
| for i in range(start, end+1, 8): | |
| cv2.circle(frame, (lm_coord[0], i+pix_step), 2, line_color, -1, lineType=cv2.LINE_AA) | |
| return frame | |
| def draw_text( | |
| img, | |
| msg, | |
| width = 7, | |
| font=cv2.FONT_HERSHEY_SIMPLEX, | |
| pos=(0, 0), | |
| font_scale=1, | |
| font_thickness=2, | |
| text_color=(0, 255, 0), | |
| text_color_bg=(0, 0, 0), | |
| box_offset=(20, 10), | |
| overlay_image = False, | |
| overlay_type = None | |
| ): | |
| #This function draws text with a customizable background box on an image. | |
| #Args: | |
| #img: The image to draw on. | |
| #msg: The message to display as a string. | |
| #width: The thickness of the background box border (default: 7). | |
| #font: The font style for the text (default: cv2.FONT_HERSHEY_SIMPLEX). | |
| #pos: The top-left corner coordinates of the text box (default: (0, 0)). | |
| #font_scale: The scaling factor for the font size (default: 1). | |
| #font_thickness: The thickness of the text (default: 2). | |
| #text_color: The color of the text in BGR format (default: green - (0, 255, 0)). | |
| #text_color_bg: The color of the background box in BGR format (default: black - (0, 0, 0)). | |
| #box_offset: The offset for the background box relative to the text (default: (20, 10)). | |
| #overlay_image: Flag to display an overlay image inside the box (default: False). | |
| #overlay_type: Type of overlay image ("correct" or "incorrect") - used when overlay_image is True. | |
| #Returns: | |
| #The size of the drawn text (width, height) as a NumPy array. | |
| offset = box_offset | |
| x, y = pos | |
| # Get the size of the text with the specified font and scale | |
| text_size, _ = cv2.getTextSize(msg, font, font_scale, font_thickness) | |
| text_w, text_h = text_size | |
| # Calculate the top-left and bottom-right corners of the text box with padding | |
| rec_start = tuple(p - o for p, o in zip(pos, offset)) | |
| rec_end = tuple(m + n - o for m, n, o in zip((x + text_w, y + text_h), offset, (25, 0))) | |
| resize_height = 0 | |
| # Handle overlay image logic | |
| if overlay_image: | |
| resize_height = rec_end[1] - rec_start[1] | |
| # Draw a rounded rectangle box with the background color | |
| img = draw_rounded_rect(img, rec_start, (rec_end[0]+resize_height, rec_end[1]), width, text_color_bg) | |
| # Resize the overlay image based on the box height | |
| if overlay_type == "correct": | |
| overlay_res = cv2.resize(correct, (resize_height, resize_height), interpolation = cv2.INTER_AREA) | |
| elif overlay_type == "incorrect": | |
| overlay_res = cv2.resize(incorrect, (resize_height, resize_height), interpolation = cv2.INTER_AREA) | |
| # Overlay the resized image onto the background box | |
| img[rec_start[1]:rec_start[1]+resize_height, rec_start[0]+width:rec_start[0]+width+resize_height] = overlay_res | |
| else: | |
| img = draw_rounded_rect(img, rec_start, rec_end, width, text_color_bg) | |
| # Draw the text onto the image with specified parameters | |
| cv2.putText( | |
| img, | |
| msg, | |
| (int(rec_start[0]+resize_height + 8), int(y + text_h + font_scale - 1)), | |
| font, | |
| font_scale, | |
| text_color, | |
| font_thickness, | |
| cv2.LINE_AA, | |
| ) | |
| return text_size | |
| def find_angle(p1, p2, ref_pt = np.array([0,0])): | |
| #This function calculates the angle between two points relative to a reference point. | |
| #Args: | |
| #p1: The first point coordinates as a NumPy array (x, y). | |
| #p2: The second point coordinates as a NumPy array (x, y). | |
| #ref_pt: The reference point coordinates as a NumPy array (default: [0, 0]). | |
| #Returns: | |
| #The angle between the two points in degrees (int). | |
| # Subtract the reference point from both points for normalization | |
| p1_ref = p1 - ref_pt | |
| p2_ref = p2 - ref_pt | |
| # Calculate the cosine of the angle using the dot product | |
| cos_theta = (np.dot(p1_ref,p2_ref)) / (1.0 * np.linalg.norm(p1_ref) * np.linalg.norm(p2_ref)) | |
| # Clip the cosine value to avoid potential errors | |
| theta = np.arccos(np.clip(cos_theta, -1.0, 1.0)) | |
| # Convert the angle from radians to degrees and cast to integer | |
| degree = int(180 / np.pi) * theta | |
| return int(degree) | |
| def get_landmark_array(pose_landmark, key, frame_width, frame_height): | |
| #This function extracts the normalized image coordinates for a landmark. | |
| #Args: | |
| #pose_landmark: A MediaPipe pose landmark object. | |
| #key: The key name of the landmark to extract (e.g., 'nose', 'shoulder.x'). | |
| #frame_width: The width of the image frame. | |
| #frame_height: The height of the image frame. | |
| #Returns: | |
| #A NumPy array containing the normalized x and y coordinates of the landmark. | |
| denorm_x = int(pose_landmark[key].x * frame_width) | |
| denorm_y = int(pose_landmark[key].y * frame_height) | |
| return np.array([denorm_x, denorm_y]) | |
| def get_landmark_features(kp_results, dict_features, feature, frame_width, frame_height): | |
| #This function extracts landmark coordinates for various body parts based on a feature name. | |
| #Args: | |
| #kp_results: The MediaPipe pose landmark results object. | |
| #dict_features: A dictionary containing landmark key names for different body parts. | |
| #feature: The name of the body part feature to extract (e.g., 'nose', 'left', 'right'). | |
| #frame_width: The width of the image frame. | |
| #frame_height: The height of the image frame. | |
| #Returns: | |
| #A list containing the landmark coordinates (as NumPy arrays) or raises an error if the feature is invalid. | |
| if feature == 'nose': | |
| return get_landmark_array(kp_results, dict_features[feature], frame_width, frame_height) | |
| elif feature == 'left' or 'right': | |
| shldr_coord = get_landmark_array(kp_results, dict_features[feature]['shoulder'], frame_width, frame_height) | |
| elbow_coord = get_landmark_array(kp_results, dict_features[feature]['elbow'], frame_width, frame_height) | |
| wrist_coord = get_landmark_array(kp_results, dict_features[feature]['wrist'], frame_width, frame_height) | |
| hip_coord = get_landmark_array(kp_results, dict_features[feature]['hip'], frame_width, frame_height) | |
| knee_coord = get_landmark_array(kp_results, dict_features[feature]['knee'], frame_width, frame_height) | |
| ankle_coord = get_landmark_array(kp_results, dict_features[feature]['ankle'], frame_width, frame_height) | |
| foot_coord = get_landmark_array(kp_results, dict_features[feature]['foot'], frame_width, frame_height) | |
| return shldr_coord, elbow_coord, wrist_coord, hip_coord, knee_coord, ankle_coord, foot_coord | |
| else: | |
| raise ValueError("feature needs to be either 'nose', 'left' or 'right") | |
| def get_mediapipe_pose( | |
| static_image_mode = False, | |
| model_complexity = 1, | |
| smooth_landmarks = True, | |
| min_detection_confidence = 0.5, | |
| min_tracking_confidence = 0.5 | |
| ): | |
| pose = mp.solutions.pose.Pose( | |
| static_image_mode = static_image_mode, | |
| model_complexity = model_complexity, | |
| smooth_landmarks = smooth_landmarks, | |
| min_detection_confidence = min_detection_confidence, | |
| min_tracking_confidence = min_tracking_confidence | |
| ) | |
| return pose |