import gradio as gr
import cv2
import numpy as np
import mediapipe as mp
from rembg import remove
from PIL import Image, ImageEnhance

# Initialisation de Mediapipe pour la détection de visage
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh()

# Fonction pour la suppression d’arrière-plan
def remove_bg(image):
    return remove(image)

# Fonction pour la retouche beauté
def beauty_filter(image, smooth_factor=50):
    img = np.array(image)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    results = face_mesh.process(img_rgb)
    if results.multi_face_landmarks:
        mask = np.zeros_like(img_rgb)
        for face_landmarks in results.multi_face_landmarks:
            points = [(int(p.x * img.shape[1]), int(p.y * img.shape[0])) for p in face_landmarks.landmark]
            convex_hull = cv2.convexHull(np.array(points))
            cv2.fillConvexPoly(mask, convex_hull, (255, 255, 255))

        blur = cv2.GaussianBlur(img_rgb, (smooth_factor, smooth_factor), 0)
        img_rgb = np.where(mask == (255, 255, 255), blur, img_rgb)

    return Image.fromarray(cv2.cvtColor(img_rgb, cv2.COLOR_BGR2RGB))

# Fonction pour la correction automatique des couleurs
def auto_color_correction(image, brightness, contrast, saturation):
    img = ImageEnhance.Brightness(image).enhance(brightness)
    img = ImageEnhance.Contrast(img).enhance(contrast)
    img = ImageEnhance.Color(img).enhance(saturation)
    return img

# Fonction pour le maquillage virtuel
def virtual_makeup(image, lipstick_intensity, blush_intensity):
    img = np.array(image)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    results = face_mesh.process(img)
    if results.multi_face_landmarks:
        for face_landmarks in results.multi_face_landmarks:
            lips = [(int(p.x * img.shape[1]), int(p.y * img.shape[0])) for p in face_landmarks.landmark[61:68]]
            cv2.fillPoly(img, [np.array(lips)], (0, 0, int(255 * lipstick_intensity)))
            
            cheek = [(int(p.x * img.shape[1]), int(p.y * img.shape[0])) for p in face_landmarks.landmark[234:236]]
            cv2.fillPoly(img, [np.array(cheek)], (int(255 * blush_intensity), 0, 0))

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return Image.fromarray(img)

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("## Éditeur de Photos IA : Retouche Beauté, Maquillage, Amélioration des Images")

    with gr.Tab("Suppression d'Arrière-Plan"):
        img_input = gr.Image(type="pil", label="Image Originale")
        img_output = gr.Image(type="pil", label="Image sans Arrière-Plan")
        btn_remove_bg = gr.Button("Supprimer l'Arrière-Plan")
        btn_remove_bg.click(remove_bg, inputs=img_input, outputs=img_output)

    with gr.Tab("Retouche Beauté"):
        img_input2 = gr.Image(type="pil", label="Image Originale")
        smooth_slider = gr.Slider(5, 100, value=50, label="Lissage de la peau")
        img_output2 = gr.Image(type="pil", label="Image Retouchée")
        btn_beauty = gr.Button("Appliquer Retouche Beauté")
        btn_beauty.click(beauty_filter, inputs=[img_input2, smooth_slider], outputs=img_output2)

    with gr.Tab("Correction des Couleurs"):
        img_input3 = gr.Image(type="pil", label="Image Originale")
        brightness = gr.Slider(0.5, 2.0, value=1.0, label="Luminosité")
        contrast = gr.Slider(0.5, 2.0, value=1.0, label="Contraste")
        saturation = gr.Slider(0.5, 2.0, value=1.0, label="Saturation")
        img_output3 = gr.Image(type="pil", label="Image Corrigée")
        btn_correct_colors = gr.Button("Corriger les Couleurs")
        btn_correct_colors.click(auto_color_correction, inputs=[img_input3, brightness, contrast, saturation], outputs=img_output3)

    with gr.Tab("Maquillage Virtuel"):
        img_input4 = gr.Image(type="pil", label="Image Originale")
        lipstick_slider = gr.Slider(0.0, 1.0, value=0.5, label="Intensité Rouge à Lèvres")
        blush_slider = gr.Slider(0.0, 1.0, value=0.5, label="Intensité Blush")
        img_output4 = gr.Image(type="pil", label="Image avec Maquillage")
        btn_makeup = gr.Button("Appliquer Maquillage Virtuel")
        btn_makeup.click(virtual_makeup, inputs=[img_input4, lipstick_slider, blush_slider], outputs=img_output4)

    demo.launch()