import gradio as gr
import tensorflow as tf
import numpy as np

# 허깅페이스에서 미리 학습된 객체 탐지 모델 불러오기
model = tf.keras.applications.MobileNetV2(weights='imagenet')

# 운동 종류와 해당 운동 자세의 정상 범위 딕셔너리
exercise_ranges = {
    '스쿼트': {'nose': (0.4, 0.6), 'left_shoulder': (0.35, 0.55), 'right_shoulder': (0.35, 0.55),
              'left_hip': (0.45, 0.65), 'right_hip': (0.45, 0.65)},
    '푸시업': {'nose': (0.3, 0.5), 'left_shoulder': (0.25, 0.45), 'right_shoulder': (0.25, 0.45),
              'left_hip': (0.35, 0.55), 'right_hip': (0.35, 0.55)}
}

def detect_and_correct_pose(image):
    # 입력 이미지 전처리
    img = tf.image.resize(image, (224, 224))
    img = tf.keras.applications.mobilenet_v2.preprocess_input(img)
    img = np.expand_dims(img, axis=0)
    
    # 객체 탐지 수행
    predictions = model.predict(img)
    predicted_class = tf.keras.applications.mobilenet_v2.decode_predictions(predictions, top=1)[0][0][1]
    
    # 운동 종류 추출
    exercise_type = predicted_class.lower()
    
    # 운동 자세 검사 및 교정
    if exercise_type in exercise_ranges:
        pose_keypoints = {}  # 여기에 키포인트 위치 정보가 있어야 하지만 실제로는 예시입니다.
        correct_pose = True
        for keypoint, (min_range, max_range) in exercise_ranges[exercise_type].items():
            if keypoint not in pose_keypoints or pose_keypoints[keypoint] < min_range or pose_keypoints[keypoint] > max_range:
                correct_pose = False
                break
        if correct_pose:
            return f"입력된 {exercise_type} 운동 자세가 올바릅니다."
        else:
            return f"입력된 {exercise_type} 운동 자세가 올바르지 않습니다. 자세를 교정하세요."
    else:
        return "운동 종류를 인식할 수 없습니다."

# Gradio 인터페이스 정의
input_image = gr.inputs.Image(shape=(224, 224))
gr.Interface(detect_and_correct_pose, input_image, "text").launch()