Yolo Yolo-Pose 팔 각도 구하기

[주인장]

import math

from ultralytics import YOLO

import matplotlib.pyplot as plt

import cv2

import numpy as np

# 세 점 사이 각도를 구하는 함수

def calc_degree(p1,p2,p3):

    result=math.atan((p1[1]-p2[1])/(p1[0]-p2[0])) - math.atan((p3[1]-p2[1])/(p3[0]-p2[0]))

    result= result*180 / math.pi

    return abs(result) #절대값

# 모델 불러오기

model = YOLO("../Models/yolov8l-pose.pt")

results = model("../Test_Images/sports/sport3.jpg")[0] #사람이 한명밖에 없으므로 0을 적는다.

keypoints = results.keypoints

# 이미지 확인 

plt.imshow(results.plot()[:,:,::-1])

plt.axis('off')

# keypoint index: 오른쪽 어깨-6, 팔꿈치-8, 손목-10

shoulder, s_prob = keypoints.xy[0,6].to('cpu').numpy(),keypoints.conf[0,6].item()

elbow, e_prob =keypoints.xy[0,8].to('cpu').numpy(),keypoints.conf[0,6].item()

wrist,w_prob = keypoints.xy[0,10].to('cpu').numpy(),keypoints.conf[0,10].item()

print(shoulder,s_prob)

print(elbow,e_prob)

print(wrist,w_prob )

r = []

for idx,xy in enumerate(keypoints.xy[0]):

    if idx in [6,8,10]:

        r.append(xy)

# 사잇값

degree = calc_degree(shoulder,elbow,wrist)

# 원본이미지에 표시

img = results.orig_img

# elbow에서 50을 더해서, elbow에서 50픽셀씩 이동한 거리에서 텍스트를 띄운다.

img = cv2.putText(img,str(round(degree))+"'",elbow.astype('int')+50,cv2.FONT_HERSHEY_COMPLEX,1,(255,255,255),2)

# 어께, 팔꿈치,손목을 연결한 뼈대 표현

pts = np.array([shoulder.astype('int'),elbow.astype('int'),wrist.astype('int')])

img = cv2.polylines(img,[pts],isClosed=False,color=(255,255,255),thickness=2,lineType=cv2.LINE_AA)

# 최종 이미지 출력.

plt.imshow(img)

plt.axis('off')