import cv2
import mediapipe as mp
import numpy as np
import gradio as gr

mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils

drawing_spec_landmarks = mp_drawing.DrawingSpec(color=(255, 0, 0), thickness=4, circle_radius=5)
drawing_spec_connections = mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=4)

pose = mp_pose.Pose(static_image_mode=False, min_detection_confidence=0.5)

def detect_pose(image):
    if image is None:
        return None

    # If the image is a NumPy array from webcam, no need to decode
    if isinstance(image, np.ndarray):
        image = image
    else:
        # Uploads might come in as PIL Images
        image = np.array(image)
    
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    if results.pose_landmarks:
        mp_drawing.draw_landmarks(
            image,
            results.pose_landmarks,
            mp_pose.POSE_CONNECTIONS,
            landmark_drawing_spec=drawing_spec_landmarks,
            connection_drawing_spec=drawing_spec_connections
        )
    
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

image_input = gr.Image()
image_input.source = "webcam"
image_input.streaming = True

gr.Interface(
    fn=detect_pose,
    inputs=image_input,
    outputs=gr.Image(),
    live=True
).launch()