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

# Initialize Mediapipe Pose Estimation
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=True, model_complexity=2)
mp_drawing = mp.solutions.drawing_utils

def estimate_pose(image):
    # Convert image from BGR (OpenCV) to RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # Perform pose detection
    results = pose.process(image_rgb)
    
    if not results.pose_landmarks:
        return image  # No pose found, return the original image

    # Draw pose landmarks on the image
    annotated_image = image.copy()
    mp_drawing.draw_landmarks(
        annotated_image,
        results.pose_landmarks,
        mp_pose.POSE_CONNECTIONS,
        landmark_drawing_spec=mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
        connection_drawing_spec=mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=2, circle_radius=2),
    )

    return annotated_image

# Gradio Interface
interface = gr.Interface(
    fn=estimate_pose,
    inputs=gr.Image(type="numpy", label="Upload an Image"),
    outputs=gr.Image(type="numpy", label="Pose Landmarks Image"),
    title="Human Pose Estimation",
    description="Upload an image to detect and visualize human pose landmarks.",
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch()