import gradio as gr
import numpy as np
import cv2
from PIL import Image

def content_aware_fill(image, mask=None, method="telea", radius=3):
    """Função simples de Content-Aware Fill"""
    
    if image is None:
        return None
    
    # Converter imagem para numpy
    img_np = np.array(image)
    
    # Converter para BGR para OpenCV
    if len(img_np.shape) == 3:
        img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    else:
        img_bgr = cv2.cvtColor(img_np, cv2.COLOR_GRAY2BGR)
    
    # Criar ou processar máscara
    if mask is not None:
        mask_np = np.array(mask)
        if len(mask_np.shape) == 3:
            mask_gray = cv2.cvtColor(mask_np, cv2.COLOR_RGB2GRAY)
        else:
            mask_gray = mask_np
        
        # Binarizar máscara
        _, mask_bin = cv2.threshold(mask_gray, 10, 255, cv2.THRESH_BINARY)
    else:
        # Sem máscara - criar uma pequena no centro para demonstração
        h, w = img_bgr.shape[:2]
        mask_bin = np.zeros((h, w), dtype=np.uint8)
        cv2.circle(mask_bin, (w//2, h//2), 20, 255, -1)
    
    # Escolher método
    if method == "telea":
        method_code = cv2.INPAINT_TELEA
    else:
        method_code = cv2.INPAINT_NS
    
    # Aplicar inpainting
    result_bgr = cv2.inpaint(img_bgr, mask_bin, radius, method_code)
    
    # Converter de volta para RGB
    result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)
    
    return Image.fromarray(result_rgb)

# Interface Gradio
with gr.Blocks(title="Content-Aware Fill") as demo:
    gr.Markdown("# 🎨 Content-Aware Fill Simples")
    gr.Markdown("Remova objetos de imagens usando algoritmos de inpainting")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="Imagem de Entrada", type="pil")
            method = gr.Radio(
                choices=["telea", "navier-stokes"],
                value="telea",
                label="Método"
            )
            radius = gr.Slider(1, 20, 3, label="Raio")
            btn = gr.Button("Processar", variant="primary")
        
        with gr.Column():
            output_img = gr.Image(label="Resultado", type="pil")
    
    # Exemplos
    examples = gr.Examples(
        examples=[
            ["https://images.unsplash.com/photo-1501854140801-50d01698950b?w=400"],
            ["https://images.unsplash.com/photo-1441974231531-c6227db76b6e?w=400"],
            ["https://images.unsplash.com/photo-1470071459604-3b5ec3a7fe05?w=400"]
        ],
        inputs=[input_img],
        outputs=[output_img],
        fn=content_aware_fill,
        cache_examples=False
    )
    
    # Conectar botão
    btn.click(
        fn=content_aware_fill,
        inputs=[input_img, None, method, radius],
        outputs=output_img
    )
    
    # Instruções
    gr.Markdown("""
    ### Como usar:
    1. Faça upload de uma imagem
    2. Escolha o método (telea = rápido, navier-stokes = qualidade)
    3. Ajuste o raio se necessário
    4. Clique em Processar
    
    **Nota:** Esta versão remove uma área central padrão.
    Para áreas específicas, crie uma máscara preta com branco na área a remover.
    """)

# Iniciar
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)