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
import torchvision.transforms as transforms

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)

    transform = transforms.Compose([transforms.PILToTensor()])
    
    # Convert image to format expected by facer
    image = facer.hwc2bchw(transform(input_pil).permute(1, 2, 0)).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)
    
    # 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 | Video tutorial - https://youtu.be/tY1u3XErmfg?si=cRWmA7iyQsNEdIBo",
    examples=["face-4.jpg"]  # Add example images here if you have any
)

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