150 lines
5.3 KiB
Python
150 lines
5.3 KiB
Python
import cv2
|
|
import mediapipe as mp
|
|
import math
|
|
|
|
|
|
class poseDetector():
|
|
|
|
def __init__(self, mode=False, complexity=1, smooth_landmarks=True,
|
|
enable_segmentation=False, smooth_segmentation=True,
|
|
detectionCon=0.5, trackCon=0.5):
|
|
|
|
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 findPose(self, img, draw=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 drawPoint(self, img, p1, p2, p3):
|
|
x1, y1 = self.lmList[p1][1], self.lmList[p1][2]
|
|
x2, y2 = self.lmList[p2][1], self.lmList[p2][2]
|
|
x3, y3 = self.lmList[p3][1], self.lmList[p3][2]
|
|
cv2.circle(img, (x1, y1), 5, (102, 106, 233), cv2.FILLED)
|
|
cv2.circle(img, (x2, y2), 5, (50, 80, 4), cv2.FILLED)
|
|
cv2.circle(img, (x3, y3), 5, (20, 95, 104), cv2.FILLED)
|
|
|
|
def Screen_occupation_left(self):
|
|
if self.results.pose_landmarks :
|
|
for id, lm in enumerate(self.results.pose_landmarks.landmark):
|
|
if id == 23:
|
|
x1 = lm.x
|
|
if id == 29:
|
|
x2 = lm.x
|
|
ocp = x2 - x1
|
|
return ocp
|
|
|
|
def Screen_occupation_right(self):
|
|
if self.results.pose_landmarks:
|
|
for id, lm in enumerate(self.results.pose_landmarks.landmark):
|
|
if id == 24:
|
|
x1 = lm.x
|
|
if id == 30:
|
|
x2 = lm.x
|
|
ocp = x2 - x1
|
|
return ocp
|
|
|
|
def findPosition(self, img, draw=True):
|
|
self.lmList = []
|
|
if self.results.pose_landmarks:
|
|
for id, lm in enumerate(self.results.pose_landmarks.landmark):
|
|
# finding height, width of the image printed
|
|
h, w, c = img.shape
|
|
# Determining the pixels of the landmarks
|
|
cx, cy, cz = int(lm.x * w), int(lm.y * h), lm.z
|
|
vis = lm.visibility
|
|
self.lmList.append([id, cx, cy, cz, vis])
|
|
if draw:
|
|
cv2.circle(img, (cx, cy), 5, (255, 0, 0), cv2.FILLED)
|
|
return self.lmList
|
|
|
|
def findConcave(self, p1, p2, p3):
|
|
x1, y1 = self.lmList[p1][1], self.lmList[p1][2]
|
|
x2, y2 = self.lmList[p2][1], self.lmList[p2][2]
|
|
x3, y3 = self.lmList[p3][1], self.lmList[p3][2]
|
|
|
|
# Calculate Angle
|
|
Concave_angle = math.degrees(math.atan2(y3 - y2, x3 - x2) - math.atan2(y1 - y2, x1 - x2))
|
|
if Concave_angle < 0:
|
|
Concave_angle += 360
|
|
|
|
return Concave_angle
|
|
|
|
def findAngle(self, img, p1, p2, p3, draw=True):
|
|
# Get the landmarks
|
|
x1, y1 = self.lmList[p1][1], self.lmList[p1][2]
|
|
x2, y2 = self.lmList[p2][1], self.lmList[p2][2]
|
|
x3, y3 = self.lmList[p3][1], self.lmList[p3][2]
|
|
|
|
# Calculate Angle
|
|
angle = math.degrees(math.atan2(y3 - y2, x3 - x2) -
|
|
math.atan2(y1 - y2, x1 - x2))
|
|
if angle < 0:
|
|
angle += 360
|
|
if angle > 180:
|
|
angle = 360 - angle
|
|
elif angle > 180:
|
|
angle = 360 - angle
|
|
# print(angle)
|
|
|
|
# Draw
|
|
if draw:
|
|
cv2.line(img, (x1, y1), (x2, y2), (255, 255, 255), 2)
|
|
cv2.line(img, (x3, y3), (x2, y2), (255, 255, 255), 2)
|
|
|
|
cv2.circle(img, (x1, y1), 3, (0, 0, 255), cv2.FILLED)
|
|
cv2.circle(img, (x1, y1), 10, (0, 0, 255), 2)
|
|
cv2.circle(img, (x2, y2), 3, (0, 0, 255), cv2.FILLED)
|
|
cv2.circle(img, (x2, y2), 10, (0, 0, 255), 2)
|
|
cv2.circle(img, (x3, y3), 3, (0, 0, 255), cv2.FILLED)
|
|
cv2.circle(img, (x3, y3), 10, (0, 0, 255), 2)
|
|
|
|
cv2.putText(img, str(int(angle)), (x2 - 20, y2 + 20),
|
|
cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 2)
|
|
return angle
|
|
|
|
def findIncludedAngle(self, x1, y1, x2, y2):
|
|
IncludedAngle = math.asin((y1 - y2) / math.hypot((y1 - y2), (x1 - x2))) * 180 / math.pi
|
|
return IncludedAngle
|
|
|
|
def findRange(self, x1, y1, x2, y2):
|
|
Range = math.sqrt(math.pow(x1 - x2, 2) + math.pow(y1 - y2, 2))
|
|
return Range
|
|
|
|
|
|
def main():
|
|
detector = poseDetector()
|
|
cap = cv2.VideoCapture(0)
|
|
while cap.isOpened():
|
|
ret, img = cap.read() # ret is just the return variable, not much in there that we will use.
|
|
if ret:
|
|
img = detector.findPose(img)
|
|
cv2.imshow('Pose Detection', img)
|
|
if cv2.waitKey(10) & 0xFF == ord('q'):
|
|
break
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|