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

# Initialize Mediapipe solutions
mp_hands = mp.solutions.hands
mp_face_mesh = mp.solutions.face_mesh
mp_drawing = mp.solutions.drawing_utils

# Function to process and draw landmarks on the input image
def process_image(input_image):
    # Convert input image to RGB for Mediapipe
    rgb_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)

    # Initialize Hands and Face Mesh models
    with mp_hands.Hands(static_image_mode=True, min_detection_confidence=0.5) as hands, \
         mp_face_mesh.FaceMesh(static_image_mode=True, min_detection_confidence=0.5) as face_mesh:
        
        # Process the image for hands
        hand_results = hands.process(rgb_image)
        # Process the image for face mesh
        face_results = face_mesh.process(rgb_image)

        # Draw Hand Landmarks
        if hand_results.multi_hand_landmarks:
            for hand_landmarks in hand_results.multi_hand_landmarks:
                mp_drawing.draw_landmarks(
                    input_image, hand_landmarks, mp_hands.HAND_CONNECTIONS,
                    mp_drawing.DrawingSpec(color=(121, 22, 76), thickness=2, circle_radius=4),
                    mp_drawing.DrawingSpec(color=(250, 44, 250), thickness=2, circle_radius=2)
                )

        # Draw Face Mesh Landmarks
        if face_results.multi_face_landmarks:
            for face_landmarks in face_results.multi_face_landmarks:
                mp_drawing.draw_landmarks(
                    input_image, face_landmarks, mp_face_mesh.FACEMESH_TESSELATION,
                    mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1),
                    mp_drawing.DrawingSpec(color=(80, 256, 121), thickness=1, circle_radius=1)
                )
    
    return input_image

# Gradio interface
def gradio_interface(image):
    # Convert Gradio PIL image to OpenCV format
    image = np.array(image)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # Process the image to detect landmarks
    processed_image = process_image(image)

    # Convert the processed image back to RGB for display in Gradio
    processed_image = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
    return processed_image

# Define Gradio app
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="numpy"),
    title="Face and Hand Landmarks Detection",
    description="Upload an image or take a photo to detect face and hand landmarks using Mediapipe and OpenCV."
)

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