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