aubome/temp.py
2024-06-13 20:07:15 +05:30

211 lines
7.6 KiB
Python

import argparse
import math
import cv2
import mediapipe as mp
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)
mp_pose = mp.solutions.pose
self.pose = mp_pose.Pose(
static_image_mode=True,
min_detection_confidence=args.pose_detection_confidence,
min_tracking_confidence=args.pose_tracking_confidence,
)
self.resized_height = 256
self.resized_width = int(
self.front_image.shape[1] *
(self.resized_height / self.front_image.shape[0]))
self.front_image_resized = cv2.resize(
self.front_image, (self.resized_width, self.resized_height))
self.side_image_resized = cv2.resize(
self.side_image, (self.resized_width, self.resized_height))
self.landmarks_indices = [
11, 12, 13, 14, 15, 16, 23, 24, 25, 26, 27, 28
]
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",
deafult=0.5,
dest="Pose detection confidence",
type=float,
help="Confidence score for pose detection")
parser.add_argument("--pose_tracking_confidence",
default=0.5,
dest="Pose detection tracking",
type=float,
help="Confidence score for pose tracking")
return parser.parse_args()
def run(self):
front_results = self.pose.process(
cv2.cvtColor(self.front_image_resized, cv2.COLOR_BGR2RGB))
side_results = self.pose.process(
cv2.cvtColor(self.side_image_resized, cv2.COLOR_BGR2RGB))
self.side_image_keypoints = self.side_image_resized.copy()
self.front_image_keypoints = self.front_image_resized.copy()
if front_results.pose_landmarks:
self.draw_landmarks(
self.front_image_keypoints,
front_results.pose_landmarks,
self.landmarks_indices,
)
if side_results.pose_landmarks:
self.draw_landmarks(
self.side_image_keypoints,
side_results.pose_landmarks,
self.landmarks_indices,
)
self.get_center_top_point(front_results)
self.print_distance(front_results)
self.display_images()
self.pose.close()
def calculate_distance(self, landmark1, landmark2):
l1 = landmark1.x, landmark1.y
l1 = int(l1[0] * self.resized_width), int(l1[1] * self.resized_height)
x1, y1 = l1
l2 = landmark2.x, landmark2.y
l2 = int(l2[0] * self.resized_width), int(l2[1] * self.resized_height)
x2, y2 = l2
pixel_distance = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
real_distance = pixel_distance * self.pixel_to_metric_ratio()
return real_distance
def pixel_to_metric_ratio(self):
height_person = 153
pixel_distance_height = 255
pixel_to_metric_ratio = height_person / pixel_distance_height
pixel_to_metric_ratio = round(pixel_to_metric_ratio)
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)
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[11]
elbow_left = landmarks[13]
wrist_left = landmarks[15]
hip_left = landmarks[23]
knee_left = landmarks[25]
ankle_left = landmarks[27]
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 destroy():
cv2.destroyAllWindows()
def display_images(self):
cv2.imshow("front_image_keypoints", self.front_image_keypoints)
cv2.imshow("side_image_keypoints", self.side_image_keypoints)
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.resized_height / 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(topmost_point[0]*self.resized_width), int(topmost_point[1] * self.resized_height))
topmost_point = (self.topmost_point[0], self.topmost_point[1])
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.resized_width),
int(center_point[1] * self.resized_height),
)
cv2.circle(self.front_image_keypoints, center_point, 2,
(255, 0, 0), -1)
xc, yc = center_point
xt, yt = topmost_point
distance = math.sqrt((xc - xt)**2 + (yc - yt)**2)
#return distance
return center_point
l = Landmarker()
try:
l.run()
except:
pass
finally:
l.destroy()