import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import cv2
import torch
import facer
from typing import Tuple

def process_image(input_image: np.ndarray) -> np.ndarray:
    """
    Process the input image to apply face smoothing effect.
    
    Args:
        input_image (np.ndarray): Input image in numpy array format
        
    Returns:
        np.ndarray: Processed image with smoothing effect applied to face
    """
    device = 'cpu'
    
    # Convert numpy array to PIL Image and back to ensure correct format
    input_pil = Image.fromarray(input_image)
    
    # Convert image to format expected by facer
    image = facer.hwc2bchw(np.array(input_pil)).to(device=device)
    
    # Initialize face detector
    face_detector = facer.face_detector('retinaface/mobilenet', device=device)
    
    # Detect faces
    with torch.inference_mode():
        faces = face_detector(image)
        
    if len(faces['bbox']) == 0:
        raise ValueError("No faces detected in the image!")
    
    # Initialize face parser
    face_parser = facer.face_parser('farl/lapa/448', device=device)
    
    # Parse face features
    with torch.inference_mode():
        faces = face_parser(image, faces)
    
    # Process nose segment
    nose_array = np.array(faces['seg']['logits'][0][6])
    nose_array = np.where(nose_array > 0, 1, 0)
    
    # Process face segment
    face_array = np.array(faces['seg']['logits'][0][1])
    face_array = np.where(face_array > 0, 1, 0)
    
    # Combine face and nose arrays
    face_array = np.clip(face_array + nose_array, 0, 1)
    
    # Apply bilateral filter for smoothing
    smooth_img = cv2.bilateralFilter(input_image, 30, 75, 75)
    
    # Apply smoothing only to face region
    smooth_img[face_array == 0] = input_image[face_array == 0]
    
    return smooth_img

def smooth_face(input_img) -> Tuple[np.ndarray, str]:
    """
    Gradio interface function to process the image and handle errors.
    
    Args:
        input_img: Input image from Gradio interface
        
    Returns:
        Tuple[np.ndarray, str]: Processed image and status message
    """
    try:
        processed_img = process_image(input_img)
        return processed_img, "Face smoothing applied successfully!"
    except ValueError as e:
        return input_img, str(e)
    except Exception as e:
        return input_img, f"Error processing image: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=smooth_face,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Image(type="numpy", label="Processed Image"),
        gr.Textbox(label="Status")
    ],
    title="Face Smoothing App",
    description="Upload an image to apply face smoothing effect. The app will detect faces and apply smoothing only to the face region.",
    examples=["face-4.jpg"]  # Add example images here if you have any
)

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