Update app.py

app.py

@@ -2,89 +2,116 @@ import gradio as gr
 import torch
 import numpy as np
 import cv2
-from PIL import Image, ImageFilter, ImageOps
+from PIL import Image
 from transformers import pipeline
 from pathlib import Path
 import zipfile
 import shutil
-import rembg
-import os
 
-print("⚡ ALVORE TURBO-INPAINT ENGINE STARTING...")
+print("AAA INPAINT + DEPTH ENGINE")
 
-device = "cpu" # HF Grátis = CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 1. Depth Anything V2 (Mantido o modelo que você quer para máxima qualidade)
-print("Carregando Depth-Anything-V2-Base...")
-depth_pipe = pipeline(
+pipe = pipeline(
     task="depth-estimation",
     model="depth-anything/Depth-Anything-V2-Base-hf",
-    device=-1 
+    device=device
 )
 
-# 2. Sessão Rembg (Alta Qualidade)
-session = rembg.new_session(model_name="u2net")
-
+# =========================
+# NORMALIZAÇÃO DEPTH
+# =========================
 def normalize_depth(depth):
     d = np.array(depth).astype(np.float32)
     d = (d - d.min()) / (d.max() - d.min() + 1e-6)
     return d
 
-def process_batch(files):
-    if not files: return None
-    out_dir = Path("alvore_turbo_out")
-    if out_dir.exists(): shutil.rmtree(out_dir)
-    out_dir.mkdir()
-    zip_path = "LOTE_ALVORE_TURBO.zip"
+# =========================
+# INPAINT PROFISSIONAL
+# =========================
+def inpaint_pro(image_pil, mask_pil):
+    img = np.array(image_pil)
+    mask = np.array(mask_pil.convert("L"))
+
+    # binariza máscara
+    mask = (mask > 128).astype(np.uint8) * 255
+
+    # 🔥 suaviza borda do buraco
+    kernel = np.ones((5,5), np.uint8)
+    mask = cv2.dilate(mask, kernel, iterations=1)
+
+    img_cv = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+    # 🔥 TELEA (detalhe)
+    telea = cv2.inpaint(img_cv, mask, 3, cv2.INPAINT_TELEA)
+
+    # 🔥 NAVIER (consistência)
+    ns = cv2.inpaint(img_cv, mask, 3, cv2.INPAINT_NS)
+
+    # 🔥 mistura inteligente
+    blended = cv2.addWeighted(telea, 0.6, ns, 0.4, 0)
+
+    # 🔥 refino leve (anti artefato)
+    blended = cv2.bilateralFilter(blended, 5, 50, 50)
+
+    return Image.fromarray(cv2.cvtColor(blended, cv2.COLOR_BGR2RGB))
+
+# =========================
+# PROCESS
+# =========================
+def process(images, masks):
+    if not images or not masks:
+        return None
+
+    out = Path("AAA_output")
+    if out.exists():
+        shutil.rmtree(out)
+    out.mkdir()
+
+    zip_path = "AAA_RESULT.zip"
 
     with zipfile.ZipFile(zip_path, 'w') as zipf:
-        for i, file in enumerate(files):
-            name = Path(file.name).stem
-            print(f"📦 [{i+1}/{len(files)}] {name}")
-
-            # Carregar Imagem
-            img_orig = Image.open(file.name).convert("RGB")
-            img_cv = cv2.cvtColor(np.array(img_orig), cv2.COLOR_RGB2BGR)
-
-            # 1. FG - REMBG (Qualidade Máxima)
-            fg_rgba = rembg.remove(img_orig, session=session)
-            fg_path = out_dir / f"{name}_FG.png"
-            fg_rgba.save(fg_path)
-            zipf.write(fg_path, fg_path.name)
-
-            # 2. BG - INPAINT MATEMÁTICO (Instantâneo)
-            # Extrair máscara do FG e dilatar
-            mask_alpha = np.array(fg_rgba.getchannel('A'))
-            kernel = np.ones((15,15), np.uint8)
-            mask_cv = cv2.dilate(mask_alpha, kernel, iterations=1)
-            
-            # O "Mágica": Inpaint via OpenCV (Método Telea)
-            # Leva milisegundos e limpa o personagem para o fundo
-            bg_cv = cv2.inpaint(img_cv, mask_cv, 3, cv2.INPAINT_TELEA)
-            bg_res = Image.fromarray(cv2.cvtColor(bg_cv, cv2.COLOR_BGR2RGB))
-            
-            bg_path = out_dir / f"{name}_BG.png"
-            bg_res.save(bg_path)
-            zipf.write(bg_path, bg_path.name)
-
-            # 3. DEPTH (Qualidade Máxima no Fundo Limpo)
-            depth_raw = depth_pipe(bg_res)["depth"]
-            d_np = normalize_depth(depth_raw)
-            depth_map = Image.fromarray((d_np * 255).astype(np.uint8))
-            depth_map = depth_map.filter(ImageFilter.GaussianBlur(0.6))
-            
-            depth_path = out_dir / f"{name}_DEPTH.png"
-            depth_map.save(depth_path)
+        for img_file, mask_file in zip(images, masks):
+
+            name = Path(img_file.name).stem
+
+            img = Image.open(img_file.name).convert("RGB")
+            mask = Image.open(mask_file.name)
+
+            # 💥 INPAINT REAL (SEM IA)
+            clean = inpaint_pro(img, mask)
+
+            # 💥 DEPTH
+            depth_raw = pipe(clean)["depth"]
+            depth_np = normalize_depth(depth_raw)
+
+            depth_img = Image.fromarray((depth_np * 255).astype(np.uint8))
+
+            # salvar
+            img_path = out / f"clean_{name}.png"
+            depth_path = out / f"depth_{name}.png"
+
+            clean.save(img_path)
+            depth_img.save(depth_path)
+
+            zipf.write(img_path, img_path.name)
             zipf.write(depth_path, depth_path.name)
 
+    print("✅ PIPELINE COMPLETO FINALIZADO")
     return zip_path
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# ⚡ ALVORE TURBO-BATCH (No-Wait Inpaint)")
-    gr.Markdown("Rembg e Depth em Alta Qualidade + Inpaint Matemático Instantâneo.")
-    inp = gr.File(file_count="multiple")
+# =========================
+# UI
+# =========================
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 AAA INPAINT + DEPTH (NO FAKE AI)")
+
+    inp_img = gr.File(file_count="multiple", label="IMAGENS (BG com buraco)")
+    inp_mask = gr.File(file_count="multiple", label="MÁSCARAS (buraco branco)")
+
     out = gr.File()
-    btn = gr.Button("🚀 GERAR LOTE TURBO", variant="primary")
-    btn.click(fn=process_batch, inputs=inp, outputs=out)
 
-demo.launch()
+    btn = gr.Button("PROCESSAR AAA")
+    btn.click(fn=process, inputs=[inp_img, inp_mask], outputs=out)
+
+demo.launch()