Spaces:
Sleeping
Sleeping
File size: 5,370 Bytes
abc9746 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
import cv2
import mediapipe as mp
import math
class PoseDetectorModified():
def _init_(self, mode=False, complexity=1, smooth_landmarks=True,
enable_segmentation=False, smooth_segmentation=True,
detectionCon=0.5, trackCon=0.5):
"""
Initializes a new instance of the PoseDetectorModified class.
Args:
mode (bool): Whether to use the upper-body-only pose landmark model or the full-body pose landmark model.
complexity (int): Complexity of the pose landmark model (must be between 0 and 2).
smooth_landmarks (bool): Whether to smooth the pose landmarks or not.
enable_segmentation (bool): Whether to enable person segmentation or not.
smooth_segmentation (bool): Whether to smooth the person segmentation or not.
detectionCon (float): Minimum confidence value (between 0 and 1) for the detection to be considered successful.
trackCon (float): Minimum confidence value (between 0 and 1) for the tracking to be considered successful.
"""
self.mode = mode
self.complexity = complexity
self.smooth_landmarks = smooth_landmarks
self.enable_segmentation = enable_segmentation
self.smooth_segmentation = smooth_segmentation
self.detectionCon = detectionCon
self.trackCon = trackCon
self.mpDraw = mp.solutions.drawing_utils
self.mpPose = mp.solutions.pose
self.pose = self.mpPose.Pose(self.mode, self.complexity, self.smooth_landmarks,
self.enable_segmentation, self.smooth_segmentation,
self.detectionCon, self.trackCon)
def find_pose(self, img, draw=True):
"""
Finds the pose landmarks and connections in an image or a video frame.
Args:
img (numpy.ndarray): The input image or video frame.
draw (bool): Whether to draw the pose landmarks and connections on the image or not.
Returns:
The input image with the pose landmarks and connections drawn on it (if draw is True).
"""
imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
self.results = self.pose.process(imgRGB)
if self.results.pose_landmarks:
if draw:
self.mpDraw.draw_landmarks(img, self.results.pose_landmarks,
self.mpPose.POSE_CONNECTIONS)
return img
def find_position(self, img, draw=True):
"""
Finds the pose landmark positions in an image or a video frame.
Args:
img (numpy.ndarray): The input image or video frame.
draw (bool): Whether to draw the pose landmark positions on the image or not.
Returns:
A list containing the landmark ID, X and Y positions for each landmark in the pose.
"""
lm_list = []
if self.results.pose_landmarks:
for id, lm in enumerate(self.results.pose_landmarks.landmark):
h, w, c = img.shape
cx, cy = int(lm.x * w), int(lm.y * h)
lm_list.append([id, cx, cy])
if draw:
cv2.circle(img, (cx, cy), 5, (255, 0, 0), cv2.FILLED)
return lm_list
def find_angle(self, img, p1, p2, p3, lm_list, draw=True):
"""
Calculates the angle between three landmarks.
Args:
img (numpy.ndarray): The input image or video frame.
p1 (int): ID of the first landmark.
p2 (int): ID of the second landmark (vertex of the angle).
p3 (int): ID of the third landmark.
lm_list (list): List of landmarks with their coordinates.
draw (bool): Whether to draw the angle on the image or not.
Returns:
float: The calculated angle in degrees.
"""
x1, y1 = lm_list[p1][1:]
x2, y2 = lm_list[p2][1:]
x3, y3 = lm_list[p3][1:]
# Calculate the angle
angle = math.degrees(math.atan2(y3 - y2, x3 - x2) - math.atan2(y1 - y2, x1 - x2))
angle = abs(angle)
if angle > 180:
angle = 360 - angle
# Draw the angle on the image
if draw:
cv2.line(img, (x1, y1), (x2, y2), (255, 255, 255), 3)
cv2.line(img, (x3, y3), (x2, y2), (255, 255, 255), 3)
cv2.circle(img, (x1, y1), 10, (0, 0, 255), cv2.FILLED)
cv2.circle(img, (x2, y2), 10, (0, 0, 255), cv2.FILLED)
cv2.circle(img, (x3, y3), 10, (0, 0, 255), cv2.FILLED)
cv2.putText(img, str(int(angle)), (x2 - 50, y2 + 50),
cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 255), 2)
return angle
def main():
detector = PoseDetectorModified()
cap = cv2.VideoCapture(0)
while cap.isOpened():
ret, img = cap.read()
if ret:
img = detector.find_pose(img)
lm_list = detector.find_position(img, False)
if len(lm_list) != 0:
angle = detector.find_angle(img, 11, 13, 15, lm_list)
print(angle)
cv2.imshow('Pose Detection', img)
if cv2.waitKey(10) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
|