244 lines
8.4 KiB
Python
244 lines
8.4 KiB
Python
import argparse
|
|
import math
|
|
import cv2
|
|
from mediapipe.python.solutions import pose
|
|
|
|
|
|
class Landmarker:
|
|
|
|
def __init__(self) -> None:
|
|
args = self.parse_args()
|
|
if args.front_image == None:
|
|
raise Exception("front image needs to be passed")
|
|
if args.side_image == None:
|
|
raise Exception("side image needs to be passed")
|
|
|
|
self.front_image = cv2.imread(args.front_image)
|
|
self.side_image = cv2.imread(args.side_image)
|
|
|
|
self.person_height = args.person_height
|
|
self.pixel_height = args.pixel_height
|
|
|
|
self.pose = pose.Pose(
|
|
static_image_mode=True,
|
|
min_detection_confidence=args.pose_detection_confidence,
|
|
min_tracking_confidence=args.pose_tracking_confidence,
|
|
)
|
|
|
|
self.landmarks_indices = [
|
|
pose.PoseLandmark.LEFT_SHOULDER.value,
|
|
pose.PoseLandmark.RIGHT_SHOULDER.value,
|
|
pose.PoseLandmark.LEFT_ELBOW.value,
|
|
pose.PoseLandmark.RIGHT_ELBOW.value,
|
|
pose.PoseLandmark.LEFT_WRIST.value,
|
|
pose.PoseLandmark.RIGHT_WRIST.value,
|
|
pose.PoseLandmark.LEFT_HIP.value,
|
|
pose.PoseLandmark.RIGHT_HIP.value,
|
|
pose.PoseLandmark.LEFT_KNEE.value,
|
|
pose.PoseLandmark.RIGHT_KNEE.value,
|
|
pose.PoseLandmark.LEFT_ANKLE.value,
|
|
pose.PoseLandmark.RIGHT_ANKLE.value,
|
|
]
|
|
|
|
def parse_args(self):
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--front",
|
|
dest="front_image",
|
|
type=str,
|
|
help="Front image")
|
|
parser.add_argument("--side",
|
|
dest="side_image",
|
|
type=str,
|
|
help="Side image")
|
|
parser.add_argument(
|
|
"--pose_detection_confidence",
|
|
dest="pose_detection_confidence",
|
|
default=0.5,
|
|
type=float,
|
|
help="Confidence score for pose detection",
|
|
)
|
|
parser.add_argument(
|
|
"--pose_tracking_confidence",
|
|
dest="pose_tracking_confidence",
|
|
default=0.5,
|
|
type=float,
|
|
help="Confidence score for pose tracking",
|
|
)
|
|
parser.add_argument(
|
|
"--person_height",
|
|
default=153,
|
|
dest="person_height",
|
|
type=int,
|
|
help="person height of person",
|
|
)
|
|
parser.add_argument(
|
|
"--pixel_height",
|
|
default=255,
|
|
dest="pixel_height",
|
|
type=int,
|
|
help="pixel height of person",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
def run(self):
|
|
|
|
front_results = self.process_images()
|
|
|
|
self.get_center_top_point(front_results)
|
|
|
|
self.print_distance(front_results)
|
|
|
|
self.display_images()
|
|
|
|
self.pose.close()
|
|
|
|
def process_images(self):
|
|
front_results = self.pose.process(
|
|
cv2.cvtColor(self.front_image, cv2.COLOR_BGR2RGB))
|
|
side_results = self.pose.process(
|
|
cv2.cvtColor(self.side_image, cv2.COLOR_BGR2RGB))
|
|
|
|
self.side_image_keypoints = self.side_image.copy()
|
|
self.front_image_keypoints = self.front_image.copy()
|
|
|
|
if front_results.pose_landmarks: # type: ignore
|
|
self.draw_landmarks(
|
|
self.front_image_keypoints,
|
|
front_results.pose_landmarks, # type: ignore
|
|
self.landmarks_indices,
|
|
)
|
|
if side_results.pose_landmarks: # type: ignore
|
|
self.draw_landmarks(
|
|
self.side_image_keypoints,
|
|
side_results.pose_landmarks, # type: ignore# type: ignore
|
|
self.landmarks_indices,
|
|
)
|
|
return front_results
|
|
|
|
def calculate_distance(self, landmark1, landmark2):
|
|
l1 = landmark1.x, landmark1.y
|
|
l1 = int(l1[0]), int(l1[1])
|
|
x1, y1 = l1
|
|
l2 = landmark2.x, landmark2.y
|
|
l2 = int(l2[0]), int(l2[1])
|
|
x2, y2 = l2
|
|
pixel_distance = self.euclidean_distance(x1, x2, y1, y2)
|
|
real_distance = pixel_distance * self.pixel_to_metric_ratio()
|
|
return real_distance
|
|
|
|
def pixel_to_metric_ratio(self):
|
|
pixel_to_metric_ratio = self.person_height / self.pixel_height
|
|
return pixel_to_metric_ratio
|
|
|
|
def draw_landmarks(self, image, landmarks, indices):
|
|
for idx in indices:
|
|
landmark = landmarks.landmark[idx]
|
|
h, w, _ = image.shape
|
|
cx, cy = int(landmark.x * w), int(landmark.y * h)
|
|
self.circle(image, cx, cy)
|
|
|
|
def circle(self, image, cx, cy):
|
|
return cv2.circle(image, (cx, cy), 2, (255, 0, 0), -1)
|
|
|
|
def print_distance(self, front_results):
|
|
if not front_results.pose_landmarks:
|
|
return
|
|
|
|
landmarks = front_results.pose_landmarks.landmark
|
|
|
|
shoulder_left = landmarks[pose.PoseLandmark.LEFT_SHOULDER.value]
|
|
elbow_left = landmarks[pose.PoseLandmark.LEFT_ELBOW.value]
|
|
wrist_left = landmarks[pose.PoseLandmark.LEFT_WRIST.value]
|
|
hip_left = landmarks[pose.PoseLandmark.LEFT_HIP.value]
|
|
knee_left = landmarks[pose.PoseLandmark.LEFT_KNEE.value]
|
|
ankle_left = landmarks[pose.PoseLandmark.LEFT_ANKLE.value]
|
|
|
|
distance_left_hand_up = self.calculate_distance(
|
|
shoulder_left, elbow_left)
|
|
print("Distance between left shoulder and left elbow:",
|
|
distance_left_hand_up)
|
|
|
|
distance_left_hand_down = self.calculate_distance(
|
|
elbow_left, wrist_left)
|
|
print("Distance between left elbow and left wrist:",
|
|
distance_left_hand_down)
|
|
|
|
distance_left_leg_up = self.calculate_distance(hip_left, knee_left)
|
|
print("Distance between left hip and left knee:", distance_left_leg_up)
|
|
|
|
distance_left_leg_down = self.calculate_distance(knee_left, ankle_left)
|
|
print("Distance between left knee and left ankle:",
|
|
distance_left_leg_down)
|
|
|
|
def euclidean_distance(self, x1, x2, y1, y2):
|
|
distance = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
|
|
return distance
|
|
|
|
def destroy(self):
|
|
cv2.destroyAllWindows()
|
|
|
|
def display_images(self):
|
|
cv2.imshow("front_image_keypoints", self.front_image_keypoints)
|
|
cv2.imshow("side_image_keypoints", self.side_image_keypoints)
|
|
cv2.waitKey(0)
|
|
|
|
def get_center_top_point(self, front_results):
|
|
gray_image = cv2.cvtColor(self.front_image_keypoints,
|
|
cv2.COLOR_BGR2GRAY)
|
|
blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
|
|
roi = blurred_image[0:int(self.front_image.shape[0] / 2), :]
|
|
edges = cv2.Canny(roi, 50, 150)
|
|
contours, hierarchy = cv2.findContours(edges, cv2.RETR_TREE,
|
|
cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
self.topmost_point = None
|
|
|
|
outermost_contours = [
|
|
contours[i] for i in range(len(contours))
|
|
if hierarchy[0][i][3] == -1
|
|
]
|
|
if outermost_contours:
|
|
largest_contour = max(outermost_contours, key=cv2.contourArea)
|
|
self.topmost_point = tuple(
|
|
largest_contour[largest_contour[:, :, 1].argmin()][0])
|
|
topmost_point = (
|
|
int(self.topmost_point[0] * self.front_image.shape[1]),
|
|
int(self.topmost_point[1] * self.front_image.shape[0]),
|
|
)
|
|
cv2.circle(self.front_image_keypoints, topmost_point, 2,
|
|
(255, 0, 0), -1)
|
|
|
|
self.left_hip_idx = 23
|
|
self.right_hip_idx = 24
|
|
|
|
self.center_point = None
|
|
|
|
if front_results.pose_landmarks:
|
|
left_hip = front_results.pose_landmarks.landmark[self.left_hip_idx]
|
|
right_hip = front_results.pose_landmarks.landmark[
|
|
self.right_hip_idx]
|
|
center_point = (
|
|
(left_hip.x + right_hip.x) / 2,
|
|
(left_hip.y + right_hip.y) / 2,
|
|
)
|
|
center_point = (
|
|
int(center_point[0] * self.front_image.shape[1]),
|
|
int(center_point[1] * self.front_image.shape[0]),
|
|
)
|
|
cv2.circle(self.front_image_keypoints, center_point, 2,
|
|
(255, 0, 0), -1)
|
|
#xc, yc = center_point
|
|
#xt, yt = topmost_point
|
|
#distance = self.euclidean_distance(xc, xt, yc, yt)
|
|
# return distance
|
|
return center_point
|
|
|
|
|
|
l = Landmarker()
|
|
try:
|
|
l.run()
|
|
except:
|
|
print("error")
|
|
finally:
|
|
l.destroy()
|