def test_invalid_connection(self):
landmark_list = text_format.Parse(
'landmark {x: 0.5 y: 0.5} landmark {x: 0.2 y: 0.2}',
landmark_pb2.NormalizedLandmarkList())
image = np.arange(27, dtype=np.uint8).reshape(3, 3, 3)
with self.assertRaisesRegex(ValueError, 'Landmark index is out of range.'):
drawing_utils.draw_landmarks(image, landmark_list, [(0, 2)])
def test_drawing_spec(self):
landmark_list = text_format.Parse(
'landmark {x: 0.1 y: 0.1}'
'landmark {x: 0.8 y: 0.8}', landmark_pb2.NormalizedLandmarkList())
image = np.zeros((100, 100, 3), np.uint8)
landmark_drawing_spec = drawing_utils.DrawingSpec(
color=(0, 0, 255), thickness=5)
connection_drawing_spec = drawing_utils.DrawingSpec(
color=(255, 0, 0), thickness=3)
expected_result = np.copy(image)
start_point = (10, 10)
end_point = (80, 80)
cv2.line(expected_result, start_point, end_point,
connection_drawing_spec.color, connection_drawing_spec.thickness)
cv2.circle(expected_result, start_point,
landmark_drawing_spec.circle_radius, landmark_drawing_spec.color,
landmark_drawing_spec.thickness)
cv2.circle(expected_result, end_point, landmark_drawing_spec.circle_radius,
landmark_drawing_spec.color, landmark_drawing_spec.thickness)
drawing_utils.draw_landmarks(
image=image,
landmark_list=landmark_list,
connections=[(0, 1)],
landmark_drawing_spec=landmark_drawing_spec,
connection_drawing_spec=connection_drawing_spec)
np.testing.assert_array_equal(image, expected_result)
def test_min_and_max_coordinate_values(self):
landmark_list = text_format.Parse(
'landmark {x: 0.0 y: 1.0}'
'landmark {x: 1.0 y: 0.0}', landmark_pb2.NormalizedLandmarkList())
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
start_point = (0, 99)
end_point = (99, 0)
cv2.line(expected_result, start_point, end_point,
DEFAULT_CONNECTION_DRAWING_SPEC.color,
DEFAULT_CONNECTION_DRAWING_SPEC.thickness)
cv2.circle(expected_result, start_point,
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
DEFAULT_CYCLE_BORDER_COLOR,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, end_point,
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
DEFAULT_CYCLE_BORDER_COLOR,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, start_point,
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, end_point,
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
drawing_utils.draw_landmarks(image=image,
landmark_list=landmark_list,
connections=[(0, 1)])
np.testing.assert_array_equal(image, expected_result)
def draw_pose_video(filename=0):
cap = cv2.VideoCapture(filename)
with pose.Pose(
min_detection_confidence=0.5, min_tracking_confidence=0.5,
) as extracter:
while cap.isOpened():
success, image = cap.read()
if not success:
print("Ignoring empty camera frame.")
continue
image = cv2.cvtColor(cv2.flip(image, 1), cv2.COLOR_BGR2RGB)
image.flags.writeable = False
results = extracter.process(image)
image.flags.writeable = True
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
drawing_utils.draw_landmarks(
image=image,
landmark_list=results.pose_landmarks,
connections=pose.POSE_CONNECTIONS,
)
cv2.imshow("MediaPipe Pose", image)
if cv2.waitKey(5) & 0xFF == 27:
break
cap.release()
def test_draw_landmarks_and_connections(self, landmark_list_text):
landmark_list = text_format.Parse(
landmark_list_text, landmark_pb2.NormalizedLandmarkList())
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
start_point = (10, 50)
end_point = (50, 10)
cv2.line(expected_result, start_point, end_point,
DEFAULT_CONNECTION_DRAWING_SPEC.color,
DEFAULT_CONNECTION_DRAWING_SPEC.thickness)
cv2.circle(expected_result, start_point,
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
DEFAULT_CYCLE_BORDER_COLOR,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, end_point,
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
DEFAULT_CYCLE_BORDER_COLOR,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, start_point,
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, end_point,
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
drawing_utils.draw_landmarks(image=image,
landmark_list=landmark_list,
connections=[(0, 1)])
np.testing.assert_array_equal(image, expected_result)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
for hand_landmarks in results.multi_hand_landmarks:
mp_drawing.draw_landmarks(frame, hand_landmarks,
mp_hands.HAND_CONNECTIONS)
path = os.path.join(tempfile.gettempdir(), self.id().split('.')[-1] +
'_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
mp_drawing.draw_landmarks(frame,
results.pose_landmarks,
mp_pose.POSE_CONNECTIONS,
landmark_drawing_spec=drawing_styles.
get_default_pose_landmarks_style())
path = self._get_output_path('_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1)
for face_landmarks in results.multi_face_landmarks:
mp_drawing.draw_landmarks(image=frame,
landmark_list=face_landmarks,
landmark_drawing_spec=drawing_spec)
path = os.path.join(
tempfile.gettempdir(),
self.id().split('.')[-1] + '_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def test_invalid_input_image(self):
image = np.arange(18, dtype=np.uint8).reshape(3, 3, 2)
with self.assertRaisesRegex(
ValueError,
'Input image must contain three channel rgb data.'):
drawing_utils.draw_landmarks(image,
landmark_pb2.NormalizedLandmarkList())
with self.assertRaisesRegex(
ValueError,
'Input image must contain three channel rgb data.'):
drawing_utils.draw_detection(image, detection_pb2.Detection())
def test_draw_single_landmark_point(self, landmark_list_text):
landmark_list = text_format.Parse(landmark_list_text,
landmark_pb2.NormalizedLandmarkList())
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
cv2.circle(expected_result, (10, 10),
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
drawing_utils.draw_landmarks(image, landmark_list)
np.testing.assert_array_equal(image, expected_result)
def draw_pose_img(filename: str):
with pose.Pose(static_image_mode=True, min_detection_confidence=0.5) as extracter:
image = cv2.imread(filename)
results = extracter.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
img_annot = image.copy()
drawing_utils.draw_landmarks(
image=img_annot,
landmark_list=results.pose_landmarks,
connections=pose.POSE_CONNECTIONS,
)
cv2.imwrite(filename.split(".")[0] + "-pose" + ".png", img_annot)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
mp_drawing.draw_landmarks(frame,
results.face_landmarks,
mp_holistic.FACEMESH_TESSELATION,
landmark_drawing_spec=None,
connection_drawing_spec=drawing_styles.
get_default_face_mesh_tesselation_style())
mp_drawing.draw_landmarks(frame,
results.pose_landmarks,
mp_holistic.POSE_CONNECTIONS,
landmark_drawing_spec=drawing_styles.
get_default_pose_landmarks_style())
path = os.path.join(
tempfile.gettempdir(),
self.id().split('.')[-1] + '_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int,
draw_iris: bool):
for face_landmarks in results.multi_face_landmarks:
mp_drawing.draw_landmarks(
frame,
face_landmarks,
mp_faces.FACEMESH_TESSELATION,
landmark_drawing_spec=None,
connection_drawing_spec=drawing_styles.
get_default_face_mesh_tesselation_style())
mp_drawing.draw_landmarks(frame,
face_landmarks,
mp_faces.FACEMESH_CONTOURS,
landmark_drawing_spec=None,
connection_drawing_spec=drawing_styles.
get_default_face_mesh_contours_style())
if draw_iris:
mp_drawing.draw_landmarks(
frame,
face_landmarks,
mp_faces.FACEMESH_IRISES,
landmark_drawing_spec=None,
connection_drawing_spec=drawing_styles.
get_default_face_mesh_iris_connections_style())
path = os.path.join(
tempfile.gettempdir(),
self.id().split('.')[-1] + '_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1)
mp_drawing.draw_landmarks(image=frame,
landmark_list=results.face_landmarks,
landmark_drawing_spec=drawing_spec)
mp_drawing.draw_landmarks(frame, results.left_hand_landmarks,
mp_holistic.HAND_CONNECTIONS)
mp_drawing.draw_landmarks(frame, results.right_hand_landmarks,
mp_holistic.HAND_CONNECTIONS)
mp_drawing.draw_landmarks(frame, results.pose_landmarks,
mp_holistic.POSE_CONNECTIONS)
path = os.path.join(
tempfile.gettempdir(),
self.id().split('.')[-1] + '_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def bootstrap(self, per_pose_class_limit=None):
"""Bootstraps images in a given folder.
Required image in folder (same use for image out folder):
pushups_up/
image_001.jpg
image_002.jpg
...
pushups_down/
image_001.jpg
image_002.jpg
...
...
Produced CSVs out folder:
pushups_up.csv
pushups_down.csv
Produced CSV structure with pose 3D landmarks:
sample_00001,x1,y1,z1,x2,y2,z2,....
sample_00002,x1,y1,z1,x2,y2,z2,....
"""
# Create output folder for CVSs.
if not os.path.exists(self._csvs_out_folder):
os.makedirs(self._csvs_out_folder)
for pose_class_name in self._pose_class_names:
print('Bootstrapping ', pose_class_name, file=sys.stderr)
# Paths for the pose class.
images_in_folder = os.path.join(self._images_in_folder,
pose_class_name)
images_out_folder = os.path.join(self._images_out_folder,
pose_class_name)
csv_out_path = os.path.join(self._csvs_out_folder,
pose_class_name + '.csv')
if not os.path.exists(images_out_folder):
os.makedirs(images_out_folder)
with open(csv_out_path, 'w') as csv_out_file:
csv_out_writer = csv.writer(csv_out_file,
delimiter=',',
quoting=csv.QUOTE_MINIMAL)
# Get list of images.
image_names = sorted([
n for n in os.listdir(images_in_folder)
if not n.startswith('.')
])
if per_pose_class_limit is not None:
image_names = image_names[:per_pose_class_limit]
# Bootstrap every image.
for image_name in tqdm.tqdm(image_names):
# Load image.
input_frame = cv2.imread(
os.path.join(images_in_folder, image_name))
input_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2RGB)
# Initialize fresh pose tracker and run it.
with mp_pose.Pose(upper_body_only=False) as pose_tracker:
result = pose_tracker.process(image=input_frame)
pose_landmarks = result.pose_landmarks
# Save image with pose prediction (if pose was detected).
output_frame = input_frame.copy()
if pose_landmarks is not None:
mp_drawing.draw_landmarks(
image=output_frame,
landmark_list=pose_landmarks,
connections=mp_pose.POSE_CONNECTIONS)
output_frame = cv2.cvtColor(output_frame,
cv2.COLOR_RGB2BGR)
cv2.imwrite(os.path.join(images_out_folder, image_name),
output_frame)
# Save landmarks if pose was detected.
if pose_landmarks is not None:
# Get landmarks.
frame_height, frame_width = output_frame.shape[
0], output_frame.shape[1]
pose_landmarks = np.array([[
lmk.x * frame_width, lmk.y * frame_height,
lmk.z * frame_width
] for lmk in pose_landmarks.landmark],
dtype=np.float32)
assert pose_landmarks.shape == (
33, 3), 'Unexpected landmarks shape: {}'.format(
pose_landmarks.shape)
csv_out_writer.writerow(
[image_name] +
pose_landmarks.flatten().astype(np.str).tolist())
# Draw XZ projection and concatenate with the image.
projection_xz = self._draw_xz_projection(
output_frame=output_frame,
pose_landmarks=pose_landmarks)
output_frame = np.concatenate(
(output_frame, projection_xz), axis=1)
while True:
# Get next frame of the video.
success, input_frame = video_cap.read()
if not success:
break
# Run pose tracker.
input_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2RGB)
result = pose_tracker.process(image=input_frame)
pose_landmarks = result.pose_landmarks
# Draw pose prediction.
output_frame = input_frame.copy()
if pose_landmarks is not None:
mp_drawing.draw_landmarks(image=output_frame,
landmark_list=pose_landmarks,
connections=mp_pose.POSE_CONNECTIONS)
if pose_landmarks is not None:
# Get landmarks.
frame_height, frame_width = output_frame.shape[
0], output_frame.shape[1]
pose_landmarks = np.array([[
lmk.x * frame_width, lmk.y * frame_height,
lmk.z * frame_width
] for lmk in pose_landmarks.landmark],
dtype=np.float32)
assert pose_landmarks.shape == (
33, 3), 'Unexpected landmarks shape: {}'.format(
pose_landmarks.shape)
# Draw the hand annotations on the image.
image.flags.writeable = True
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
########################################
# Stage 05 - Check result
if not results.multi_hand_landmarks:
socket.send(b"error3")
continue
coordenates = results.multi_hand_landmarks
landmarks_count = 0
for hand_landmarks in coordenates:
mp_drawing.draw_landmarks(image, hand_landmarks, mp_hands.HAND_CONNECTIONS)
landmarks_count += 1
if landmarks_count != 21:
socket.send(b"error4")
continue
#relative_coordenates = converter.convert_to_relative(absolute_coordenates=coordenates)
#new_coordenates = converter.convert_to_absolute(relative_coordenates)
socket.send(b"ack")
landmarks_count = 0
for hand_landmarks in coordenates:
mp_hands.HAND_CONNECTIONS
def classify(filename):
# Specify your video name and target pose class to count the repetitions.
video_path = filename
class_names = ['pushups_down', 'lunges_down', 'squats_down']
pose_samples_folders = [
'fitness_poses_csvs_out_pushups', 'fitness_poses_csvs_out_lunges',
'fitness_poses_csvs_out_squats'
]
all_repetition_counts = []
# out_video_path = f'{filename[:-3]}mov'
# Open the video.
for i in range(len(class_names)):
class_name = class_names[i]
video_cap = cv2.VideoCapture(video_path)
# Get some video parameters to generate output video with classificaiton.
video_n_frames = video_cap.get(cv2.CAP_PROP_FRAME_COUNT)
# video_fps = video_cap.get(cv2.CAP_PROP_FPS)
# video_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
# video_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# Initilize tracker, classifier and counter.
# Do that before every video as all of them have state.
# Folder with pose class CSVs. That should be the same folder you using while
# building classifier to output CSVs.
pose_samples_folder = pose_samples_folders[i]
# Initialize tracker.
pose_tracker = mp_pose.Pose(upper_body_only=False)
# Initialize embedder.
pose_embedder = FullBodyPoseEmbedder()
# Initialize classifier.
# Ceck that you are using the same parameters as during bootstrapping.
pose_classifier = PoseClassifier(
pose_samples_folder=pose_samples_folder,
pose_embedder=pose_embedder,
top_n_by_max_distance=30,
top_n_by_mean_distance=10)
# # Uncomment to validate target poses used by classifier and find outliers.
# outliers = pose_classifier.find_pose_sample_outliers()
# print('Number of pose sample outliers (consider removing them): ', len(outliers))
# Initialize EMA smoothing.
pose_classification_filter = EMADictSmoothing(window_size=10,
alpha=0.2)
# Initialize counter.
repetition_counter = RepetitionCounter(class_name=class_name,
enter_threshold=6,
exit_threshold=4)
# Initialize renderer.
# pose_classification_visualizer = PoseClassificationVisualizer(
# class_name=class_name,
# plot_x_max=video_n_frames,
# # Graphic looks nicer if it's the same as `top_n_by_mean_distance`.
# plot_y_max=10)
# Run classification on a video.
import os
# Open output video.
# out_video = cv2.VideoWriter(out_video_path, cv2.VideoWriter_fourcc(*'mp4v'), video_fps, (video_width, video_height))
frame_idx = 0
repetitions_count = -1
# output_frame = None
with tqdm.tqdm(total=video_n_frames, position=0, leave=True) as pbar:
while True:
# Get next frame of the video.
success, input_frame = video_cap.read()
if not success:
break
# Run pose tracker.
input_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2RGB)
result = pose_tracker.process(image=input_frame)
pose_landmarks = result.pose_landmarks
# Draw pose prediction.
output_frame = input_frame.copy()
if pose_landmarks is not None:
mp_drawing.draw_landmarks(
image=output_frame,
landmark_list=pose_landmarks,
connections=mp_pose.POSE_CONNECTIONS)
if pose_landmarks is not None:
# Get landmarks.
frame_height, frame_width = output_frame.shape[
0], output_frame.shape[1]
pose_landmarks = np.array([[
lmk.x * frame_width, lmk.y * frame_height,
lmk.z * frame_width
] for lmk in pose_landmarks.landmark],
dtype=np.float32)
assert pose_landmarks.shape == (
33, 3), 'Unexpected landmarks shape: {}'.format(
pose_landmarks.shape)
# Classify the pose on the current frame.
pose_classification = pose_classifier(pose_landmarks)
# Smooth classification using EMA.
pose_classification_filtered = pose_classification_filter(
pose_classification)
# Count repetitions.
repetitions_count = repetition_counter(
pose_classification_filtered)
else:
# No pose => no classification on current frame.
# pose_classification = None
# Still add empty classification to the filter to maintaing correct
# smoothing for future frames.
# pose_classification_filtered = pose_classification_filter(dict())
# pose_classification_filtered = None
# Don't update the counter presuming that person is 'frozen'. Just
# take the latest repetitions count.
repetitions_count = repetition_counter.n_repeats
# Draw classification plot and repetition counter.
# output_frame = pose_classification_visualizer(
# frame=output_frame,
# pose_classification=pose_classification,
# pose_classification_filtered=pose_classification_filtered,
# repetitions_count=repetitions_count)
# Save the output frame.
# out_video.write(cv2.cvtColor(np.array(output_frame), cv2.COLOR_RGB2BGR))
# Show intermediate frames of the video to track progress.
# if frame_idx % 50 == 0:
# show_image(output_frame)
frame_idx += 1
pbar.update()
all_repetition_counts.append(repetitions_count)
pose_tracker.close()
# Close output video.
# out_video.release()
# Release MediaPipe resources.
print(all_repetition_counts)
pushup_counts, lunge_counts, squat_counts = all_repetition_counts[
0], all_repetition_counts[1], all_repetition_counts[2]
return pushup_counts, lunge_counts, squat_counts
