app.py

import gradio as gr
import numpy as np
from PIL import Image, ImageFilter, ImageDraw, ImageEnhance, ImageOps
import cv2
import rembg
import tempfile
import os
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ImageProcessor:
    def __init__(self):
        logger.info("🔧 Inicializando processador...")
        try:
            self.session = rembg.new_session("u2net")
            logger.info("✅ Modelo U2Net carregado com sucesso!")
        except Exception as e:
            logger.error(f"❌ Erro ao carregar modelo: {e}")
            self.session = None
    
    def remove_background(self, img):
        """Remove fundo da imagem"""
        logger.info(f"🔄 REMOVER FUNDO - Entrada: {img.size}, modo: {img.mode}")
        try:
            if self.session:
                logger.info("🔄 Chamando rembg.remove()...")
                result = rembg.remove(img, session=self.session)
                logger.info(f"✅ Fundo removido! Resultado: {result.size}, modo: {result.mode}")
                return result
            else:
                logger.warning("⚠️ Sessão não disponível, retornando imagem original")
                return img.convert('RGBA')
        except Exception as e:
            logger.error(f"❌ Erro na remoção de fundo: {e}")
            return img.convert('RGBA')
    
    def create_smooth_mask(self, alpha_channel):
        """Cria máscara suavizada"""
        logger.info("🔄 CRIAR MÁSCARA SUAVIZADA")
        logger.info(f"📊 Alpha shape: {alpha_channel.shape}, dtype: {alpha_channel.dtype}")
        
        # Converte para binário
        mask = (alpha_channel > 128).astype(np.uint8) * 255
        logger.info(f"📊 Máscara binária - Unique values: {np.unique(mask)}")
        
        # Preenche buracos
        contours, hierarchy = cv2.findContours(mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
        logger.info(f"📊 Número de contornos encontrados: {len(contours)}")
        
        if hierarchy is not None:
            for i, contour in enumerate(contours):
                if hierarchy[0][i][3] == -1:  # Contorno externo
                    cv2.drawContours(mask, [contour], 0, 255, -1)
        
        # Suaviza bordas
        mask = cv2.GaussianBlur(mask, (5, 5), 0)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((3,3), np.uint8))
        
        logger.info(f"✅ Máscara criada! Shape final: {mask.shape}")
        return mask
    
    def add_border(self, img, mask, size, color_hex):
        """Adiciona borda ao objeto usando dilatação da máscara"""
        logger.info(f"🔄 ADICIONAR BORDA - size={size}, color_hex={color_hex}")
        
        if size == 0:
            logger.info("⏭️ Borda desativada (tamanho = 0)")
            return img
        
        try:
            # Garante que temos uma string de cor válida
            if not color_hex or not isinstance(color_hex, str):
                color_hex = "#FF0000"
                logger.warning(f"⚠️ Cor inválida, usando padrão: {color_hex}")
            
            logger.info(f"🎨 Cor original recebida: {color_hex}")
            
            # Converte cor hex para RGB
            hex_color = color_hex.lstrip('#')
            logger.info(f"🎨 Hex após strip: {hex_color}")
            
            if len(hex_color) != 6:
                hex_color = "FF0000"
                logger.warning(f"⚠️ Formato de cor inválido, usando vermelho")
            
            r = int(hex_color[0:2], 16)
            g = int(hex_color[2:4], 16)
            b = int(hex_color[4:6], 16)
            border_color = (r, g, b, 255)
            logger.info(f"🎨 Cor da borda RGB: {border_color}")
            
            # Garante que a máscara está no formato correto
            logger.info(f"📊 Máscara dtype antes: {mask.dtype}")
            if mask.dtype != np.uint8:
                mask = mask.astype(np.uint8)
            
            # Cria kernel para dilatação baseado no tamanho da borda
            kernel_size = max(3, size * 2)
            kernel = np.ones((kernel_size, kernel_size), np.uint8)
            logger.info(f"📏 Kernel size: {kernel_size}x{kernel_size}")
            
            dilated_mask = cv2.dilate(mask, kernel, iterations=1)
            logger.info(f"📊 Máscara dilatada - valores únicos: {np.unique(dilated_mask)}")
            
            # Cria máscara apenas da borda
            border_mask = cv2.subtract(dilated_mask, mask)
            border_mask = cv2.GaussianBlur(border_mask, (3, 3), 0)
            
            border_pixels = np.sum(border_mask > 0)
            logger.info(f"📏 Pixels de borda criados: {border_pixels}")
            
            # Cria imagem base com a borda
            img_array = np.array(img).copy()
            height, width = img_array.shape[:2]
            logger.info(f"📊 Imagem array shape: {img_array.shape}")
            
            # Cria camada só com a borda
            border_only = np.zeros((height, width, 4), dtype=np.uint8)
            border_only[:, :, 0] = r
            border_only[:, :, 1] = g
            border_only[:, :, 2] = b
            border_only[:, :, 3] = border_mask
            
            # Log da camada de borda
            logger.info(f"📊 Border_only - R min/max: {border_only[:,:,0].min()}/{border_only[:,:,0].max()}")
            logger.info(f"📊 Border_only - G min/max: {border_only[:,:,1].min()}/{border_only[:,:,1].max()}")
            logger.info(f"📊 Border_only - B min/max: {border_only[:,:,2].min()}/{border_only[:,:,2].max()}")
            
            # Combina: borda primeiro, depois objeto
            result_array = border_only.copy()
            
            # Sobrepõe a imagem original onde ela tem alpha > 0
            alpha_img = img_array[:, :, 3].astype(float) / 255.0
            alpha_border = border_only[:, :, 3].astype(float) / 255.0
            
            logger.info(f"📊 Alpha imagem - min/max: {alpha_img.min():.2f}/{alpha_img.max():.2f}")
            logger.info(f"📊 Alpha borda - min/max: {alpha_border.min():.2f}/{alpha_border.max():.2f}")
            
            for c in range(3):  # RGB
                result_array[:, :, c] = (
                    img_array[:, :, c] * alpha_img +
                    border_only[:, :, c] * alpha_border * (1 - alpha_img)
                ).astype(np.uint8)
            
            # Canal alpha combinado
            result_array[:, :, 3] = np.maximum(img_array[:, :, 3], border_only[:, :, 3])
            
            # Verifica resultado
            logger.info(f"📊 Result array - R min/max: {result_array[:,:,0].min()}/{result_array[:,:,0].max()}")
            logger.info(f"📊 Result array - cor média: {result_array[:,:,:3].mean(axis=(0,1))}")
            
            result = Image.fromarray(result_array, 'RGBA')
            logger.info(f"✅ Borda adicionada! Result: {result.size}, modo: {result.mode}")
            return result
            
        except Exception as e:
            logger.error(f"❌ Erro ao adicionar borda: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return img
    
    def apply_background(self, foreground, bg_type, bg_color_hex, bg_image):
        """Aplica fundo à imagem"""
        logger.info(f"🔄 APLICAR FUNDO - Tipo: '{bg_type}', Cor: '{bg_color_hex}'")
        logger.info(f"📊 Foreground: {foreground.size}, modo: {foreground.mode}")
        
        try:
            if bg_type == "transparent":
                logger.info("✅ Mantendo fundo transparente")
                return foreground
            
            elif bg_type == "color":
                # Valida e corrige cor se necessário
                logger.info(f"🎨 Processando fundo colorido - Valor recebido: {bg_color_hex}")
                
                if not bg_color_hex or not isinstance(bg_color_hex, str):
                    bg_color_hex = "#FFFFFF"
                    logger.warning(f"⚠️ Cor de fundo inválida, usando branco")
                
                logger.info(f"🎨 Cor após validação: {bg_color_hex}")
                
                # Converte hex para RGB
                hex_color = bg_color_hex.lstrip('#')
                logger.info(f"🎨 Hex após strip: {hex_color}")
                
                if len(hex_color) != 6:
                    hex_color = "FFFFFF"
                    logger.warning("⚠️ Formato de cor inválido, usando branco")
                
                r = int(hex_color[0:2], 16)
                g = int(hex_color[2:4], 16)
                b = int(hex_color[4:6], 16)
                
                logger.info(f"🎨 RGB do fundo: ({r}, {g}, {b})")
                
                # Cria fundo sólido
                bg = Image.new('RGBA', foreground.size, (r, g, b, 255))
                logger.info(f"📊 Background criado: {bg.size}, cor: {bg.getpixel((0,0))}")
                
                # Combina usando alpha composite
                result = Image.alpha_composite(bg, foreground)
                logger.info(f"✅ Fundo colorido aplicado! Result: {result.size}, modo: {result.mode}")
                
                # Verifica cor resultante
                sample_pixel = result.getpixel((10, 10))
                logger.info(f"📊 Pixel de amostra (10,10): {sample_pixel}")
                return result
            
            elif bg_type == "image":
                if bg_image is None:
                    logger.warning("⚠️ Nenhuma imagem de fundo fornecida, mantendo transparente")
                    return foreground
                
                logger.info(f"🖼️ Aplicando imagem de fundo {bg_image.size}, modo: {bg_image.mode}")
                
                # Converte e redimensiona fundo
                bg = bg_image.convert('RGBA')
                logger.info(f"📊 Background convertido: {bg.size}, modo: {bg.mode}")
                
                if bg.size != foreground.size:
                    logger.info(f"📐 Redimensionando fundo de {bg.size} para {foreground.size}")
                    bg = bg.resize(foreground.size, Image.Resampling.LANCZOS)
                
                # Combina
                result = Image.alpha_composite(bg, foreground)
                logger.info("✅ Imagem de fundo aplicada!")
                return result
            
            else:
                logger.warning(f"⚠️ Tipo de fundo desconhecido: '{bg_type}', mantendo transparente")
                return foreground
                
        except Exception as e:
            logger.error(f"❌ Erro ao aplicar fundo: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return foreground
    
    def process(self, img, border_size, border_color, bg_type, bg_img, bg_color, effect, smooth):
        """Processa a imagem completa"""
        logger.info("="*60)
        logger.info("🚀 INICIANDO PROCESSAMENTO COMPLETO")
        logger.info("="*60)
        
        # Log de todos os parâmetros recebidos
        logger.info(f"📋 PARÂMETROS RECEBIDOS:")
        logger.info(f"  - img: {'None' if img is None else f'{img.size}, {img.mode}'}")
        logger.info(f"  - border_size: {border_size}")
        logger.info(f"  - border_color: {border_color}")
        logger.info(f"  - bg_type: {bg_type}")
        logger.info(f"  - bg_img: {'None' if bg_img is None else f'{bg_img.size}, {bg_img.mode}'}")
        logger.info(f"  - bg_color: {bg_color}")
        logger.info(f"  - smooth: {smooth}")
        
        if img is None:
            logger.warning("⚠️ Nenhuma imagem fornecida")
            return None, None, "⚠️ Por favor, faça upload de uma imagem", None, None
        
        try:
            logger.info(f"📸 Imagem recebida: {img.size}, modo: {img.mode}")
            
            # Passo 1: Remove fundo
            logger.info("="*30)
            logger.info("1️⃣ REMOVENDO FUNDO")
            logger.info("="*30)
            no_bg = self.remove_background(img)
            logger.info(f"📊 No BG result: {no_bg.size}, modo: {no_bg.mode}")
            
            # Passo 2: Extrai e processa máscara
            logger.info("="*30)
            logger.info("2️⃣ PROCESSANDO MÁSCARA")
            logger.info("="*30)
            if no_bg.mode == 'RGBA':
                alpha = np.array(no_bg.split()[3])
                logger.info(f"📊 Alpha channel shape: {alpha.shape}, unique: {np.unique(alpha)}")
            else:
                logger.warning("⚠️ Imagem sem canal alpha, criando máscara completa")
                alpha = np.ones((no_bg.height, no_bg.width), dtype=np.uint8) * 255
            
            # CORREÇÃO CRÍTICA: Verifica se smooth é realmente booleano
            logger.info(f"🔄 Smooth value: {smooth}, type: {type(smooth)}")
            
            if smooth:
                logger.info("🔧 Criando máscara suavizada")
                mask = self.create_smooth_mask(alpha)
            else:
                logger.info("🔧 Criando máscara básica (sem suavização)")
                mask = (alpha > 128).astype(np.uint8) * 255
            
            logger.info(f"📊 Máscara final - shape: {mask.shape}, unique: {np.unique(mask)}")
            
            # Passo 3: Aplica máscara refinada à imagem original
            logger.info("="*30)
            logger.info("3️⃣ APLICANDO MÁSCARA REFINADA")
            logger.info("="*30)
            logger.info("🔄 Convertendo imagem para RGBA...")
            result_rgba = img.convert('RGBA')
            logger.info(f"📊 Imagem convertida: {result_rgba.size}, modo: {result_rgba.mode}")
            
            result_array = np.array(result_rgba)
            logger.info(f"📊 Array original shape: {result_array.shape}")
            
            result_array[:, :, 3] = mask
            logger.info(f"📊 Array após aplicar máscara - Alpha unique: {np.unique(result_array[:,:,3])}")
            
            result = Image.fromarray(result_array, 'RGBA')
            logger.info(f"✅ Máscara aplicada! Result: {result.size}, modo: {result.mode}")
            
            # Passo 4: Adiciona borda
            logger.info("="*30)
            logger.info("4️⃣ ADICIONANDO BORDA")
            logger.info("="*30)
            logger.info(f"🎨 Usando cor da borda: {border_color}")
            bordered = self.add_border(result, mask, border_size, border_color)
            logger.info(f"📊 Bordered result: {bordered.size}, modo: {bordered.mode}")
            
            # Passo 5: Aplica fundo
            logger.info("="*30)
            logger.info("5️⃣ APLICANDO FUNDO")
            logger.info("="*30)
            logger.info(f"🎭 Tipo de fundo: {bg_type}, Cor: {bg_color}")
            final = self.apply_background(bordered, bg_type, bg_color, bg_img)
            logger.info(f"📊 Final result: {final.size}, modo: {final.mode}")
            
            # Passo 6: Salva arquivos
            logger.info("="*30)
            logger.info("6️⃣ SALVANDO ARQUIVOS")
            logger.info("="*30)
            logger.info("💾 Salvando arquivos...")
            final_path = os.path.join(tempfile.gettempdir(), f"final_{os.getpid()}.png")
            nobg_path = os.path.join(tempfile.gettempdir(), f"nobg_{os.getpid()}.png")
            
            final.save(final_path, "PNG")
            bordered.save(nobg_path, "PNG")
            logger.info(f"✅ Arquivos salvos: {final_path}, {nobg_path}")
            
            # Mostra informações mais detalhadas
            if bg_type == "color":
                bg_info = f"{bg_type} ({bg_color})"
            elif bg_type == "image" and bg_img is not None:
                bg_info = f"{bg_type} ({bg_img.size[0]}x{bg_img.size[1]})"
            else:
                bg_info = bg_type
            
            status = f"✅ SUCESSO!\n"
            status += f"📏 Tamanho: {final.size[0]}x{final.size[1]}px\n"
            status += f"🖍️ Borda: {border_size}px - Cor: {border_color}\n"
            status += f"🎭 Fundo: {bg_info}\n"
            status += f"✨ Suavização: {'Ativa' if smooth else 'Desativada'}"
            
            logger.info("="*60)
            logger.info("✅ PROCESSAMENTO CONCLUÍDO COM SUCESSO!")
            logger.info("="*60)
            
            return final, bordered, status, final_path, nobg_path
            
        except Exception as e:
            logger.error("="*60)
            logger.error(f"❌ ERRO CRÍTICO NO PROCESSAMENTO: {e}")
            logger.error("="*60)
            import traceback
            logger.error(traceback.format_exc())
            
            error_msg = f"❌ Erro no processamento:\n{str(e)}\n\nVerifique os logs para detalhes."
            return None, None, error_msg, None, None

# Inicializa processador globalmente
logger.info("🌟 Iniciando aplicação...")
processor = ImageProcessor()

# Interface Gradio
with gr.Blocks(theme=gr.themes.Soft(), title="🎨 Magic Background Remover PRO") as demo:
    gr.Markdown("""
    # 🎨 Magic Background Remover PRO
    ### Remoção automática de fundo com IA + Bordas personalizadas
    
    **Como usar:**
    1. 📤 Faça upload de uma imagem
    2. 🖍️ Ajuste a borda (tamanho e cor)
    3. 🎭 Escolha o tipo de fundo
    4. 📥 Baixe o resultado!
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Controles")
            
            img_in = gr.Image(label="📤 Upload da Imagem", type="pil", height=300)
            
            with gr.Group():
                gr.Markdown("**🖍️ Configurações da Borda**")
                border_size = gr.Slider(0, 50, 10, step=1, label="Espessura (px)")
                border_color = gr.ColorPicker("#FF0000", label="Cor da Borda")
                smooth = gr.Checkbox(True, label="✨ Suavizar bordas (recomendado)")
            
            with gr.Group():
                gr.Markdown("**🎭 Configurações do Fundo**")
                bg_type = gr.Radio(
                    choices=["transparent", "color", "image"], 
                    label="Tipo de Fundo",
                    value="transparent",
                    info="Selecione o tipo de fundo"
                )
                bg_color = gr.ColorPicker("#FFFFFF", label="Cor do Fundo", visible=False)
                bg_img = gr.Image(label="Imagem de Fundo", type="pil", visible=False)
            
            def toggle_bg_controls(bg_type):
                """Controla visibilidade dos controles de fundo"""
                logger.info(f"🔄 TOGGLE BG CONTROLS - bg_type: {bg_type}")
                visible_color = (bg_type == "color")
                visible_img = (bg_type == "image")
                return (
                    gr.ColorPicker(visible=visible_color),
                    gr.Image(visible=visible_img)
                )
            
            bg_type.change(
                fn=toggle_bg_controls,
                inputs=[bg_type],
                outputs=[bg_color, bg_img]
            )
            
            process_btn = gr.Button("🚀 PROCESSAR IMAGEM", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            gr.Markdown("### 🖼️ Resultados")
            
            img_out = gr.Image(label="Imagem Final", height=350)
            img_nobg = gr.Image(label="Apenas sem Fundo (com borda)", height=250)
            
            with gr.Row():
                dl_final = gr.File(label="📥 Download - Imagem Final", visible=False)
                dl_nobg = gr.File(label="📥 Download - Sem Fundo", visible=False)
            
            status = gr.Textbox(
                label="📊 Status do Processamento", 
                lines=6,
                value="✅ Sistema pronto!\n\n📝 Aguardando upload de imagem..."
            )
    
    gr.Markdown("""
    ---
    ### 💡 Dicas para melhores resultados:
    - Use imagens com bom contraste entre objeto e fundo
    - Para retratos: bordas de 8-15px funcionam bem
    - Para produtos: bordas de 3-8px são ideais
    - Ative a suavização para bordas mais naturais
    
    **🤖 Processamento com IA U2Net - Qualidade profissional!**
    """)
    
    # Conecta eventos
    inputs = [img_in, border_size, border_color, bg_type, bg_img, bg_color, gr.State(""), smooth]
    outputs = [img_out, img_nobg, status, dl_final, dl_nobg]
    
    # Função wrapper para log extra
    def process_with_logs(*args):
        logger.info("🔄 EVENTO DISPARADO - Iniciando processamento via interface")
        logger.info(f"📋 Args recebidos: {len(args)}")
        for i, arg in enumerate(args):
            if i == 0:  # img
                logger.info(f"  Arg {i} (img): {'None' if arg is None else f'{arg.size}, {arg.mode}'}")
            elif i == 3:  # bg_type
                logger.info(f"  Arg {i} (bg_type): {arg}")
            elif i == 4:  # bg_img
                logger.info(f"  Arg {i} (bg_img): {'None' if arg is None else f'{arg.size}, {arg.mode}'}")
            elif i == 5:  # bg_color
                logger.info(f"  Arg {i} (bg_color): {arg}")
            elif i == 7:  # smooth
                logger.info(f"  Arg {i} (smooth): {arg}, type: {type(arg)}")
        
        result = processor.process(*args)
        logger.info(f"✅ Processamento via interface concluído")
        return result
    
    # Botão principal
    process_btn.click(
        fn=process_with_logs, 
        inputs=inputs, 
        outputs=outputs,
        show_progress=True
    )
    
    # Reprocessa quando imagem mudar
    img_in.change(
        fn=process_with_logs, 
        inputs=inputs, 
        outputs=outputs,
        show_progress=True
    )
    
    # Reprocessa ao mudar tamanho da borda
    border_size.change(
        fn=process_with_logs, 
        inputs=inputs, 
        outputs=outputs,
        show_progress=True
    )
    
    # Reprocessa ao mudar COR da borda
    border_color.change(
        fn=process_with_logs, 
        inputs=inputs, 
        outputs=outputs,
        show_progress=True
    )
    
    # Reprocessa ao mudar suavização
    smooth.change(
        fn=process_with_logs, 
        inputs=inputs, 
        outputs=outputs,
        show_progress=True
    )
    
    # Reprocessa ao mudar tipo de fundo
    bg_type.change(
        fn=process_with_logs, 
        inputs=inputs, 
        outputs=outputs,
        show_progress=True
    )
    
    # Reprocessa ao mudar COR do fundo
    bg_color.change(
        fn=process_with_logs, 
        inputs=inputs, 
        outputs=outputs,
        show_progress=True
    )
    
    # Reprocessa ao mudar IMAGEM de fundo
    bg_img.change(
        fn=process_with_logs, 
        inputs=inputs, 
        outputs=outputs,
        show_progress=True
    )

logger.info("✅ Interface Gradio criada com sucesso!")
demo.launch(debug=True)