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/situp/2.jpg")[0]  # 윗몸일으키기 이미지 예시

keypoints = results.keypoints

# 왼쪽 어깨, 힙, 무릎의 keypoint 인덱스: 어깨-5, 힙-11, 무릎-13

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

hip, h_prob = keypoints.xy[0, 11].to('cpu').numpy(), keypoints.conf[0, 11].item()

knee, k_prob = keypoints.xy[0, 13].to('cpu').numpy(), keypoints.conf[0, 13].item()

print(shoulder, s_prob)

# 각도 계산

situp_angle = calc_degree(shoulder, hip, knee)

situp_threshold = 45  # 예시 임계값, 각도가 이 값 이하이면 윗몸일으키기 성공으로 판단

# 성공 여부 판단

situp_success = situp_angle <= situp_threshold

status_text = "성공" if situp_success else "좀 더 열심히"

print(f"Sit-Up Angle: {situp_angle:.2f}°, {situp_success} {status_text}")

# 이미지에 각도와 성공 여부 표시

img = results.orig_img

img = cv2.putText(img, f"Angle: {round(situp_angle)}'", hip.astype('int') + 20, cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 2)

img = cv2.putText(img, status_text, (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0) if situp_success else (0, 0, 255), 2)

# 어깨, 힙, 무릎을 연결한 뼈대 표현

pts = np.array([shoulder.astype('int'), hip.astype('int'), knee.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')

plt.show()