Update app.py

app.py

@@ -1,44 +1,104 @@
 import gradio as gr
-import torch
 import numpy as np
 from PIL import Image, ImageFilter, ImageOps
 from transformers import pipeline
 from pathlib import Path
 import zipfile
 import shutil
+import os
+import cv2
 
-print("DEPTH AAA ENGINE (PURE MODE)")
+print("DEPTH AAA ENGINE (CLEAN REBUILD)")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# -------------------------
+# CONFIG
+# -------------------------
+MODEL_ID = "depth-anything/Depth-Anything-V2-Base-hf"
+
+# Use cuda if available, otherwise cpu
+device = 0 if os.environ.get("CUDA_VISIBLE_DEVICES", "") != "" else -1
 
 pipe = pipeline(
-task="depth-estimation",
-model="depth-anything/Depth-Anything-V2-Base-hf",
-device=device
+    task="depth-estimation",
+    model=MODEL_ID,
+    device=device
 )
 
-#========================= NORMALIZAÇÃO PROFISSIONAL ========================= 
+# -------------------------
+# HELPERS
+# -------------------------
+def normalize_depth(depth_img: Image.Image) -> np.ndarray:
+    d = np.array(depth_img).astype(np.float32)
+    d_min = d.min()
+    d_max = d.max()
+    d = (d - d_min) / (d_max - d_min + 1e-6)
+    return d
+
+def refine_depth(depth_np: np.ndarray) -> Image.Image:
+    img = Image.fromarray((depth_np * 255).astype(np.uint8))
 
-def normalize_depth(depth):
-d = np.array(depth).astype(np.float32)
+    # leve blur para suavizar banding
+    img = img.filter(ImageFilter.GaussianBlur(1.0))
 
-# normalização estável d = (d - d.min()) / (d.max() - d.min() + 1e-6) return d ========================= REFINO (SEM QUEBRAR O MAPA) ========================= 
+    # melhora contraste sem destruir detalhe
+    img = ImageOps.autocontrast(img, cutoff=0.3)
 
-def refine_depth(depth_np):
-img = Image.fromarray((depth_np * 255).astype(np.uint8))
+    return img
 
-# 🔥 suavização leve (ANTI-BANDING) img = img.filter(ImageFilter.GaussianBlur(1.0)) # 🔥 contraste leve (mantém detalhe) img = ImageOps.autocontrast(img, cutoff=0.3) return img ========================= PIPELINE ========================= 
+def colorize_depth(depth_gray: Image.Image) -> Image.Image:
+    arr = np.array(depth_gray)
+    color = cv2.applyColorMap(arr, cv2.COLORMAP_INFERNO)
+    color = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(color)
 
 def process(files):
-if not files:
-return None
+    if not files:
+        return None, None
+
+    out_dir = Path("depth_output")
+    if out_dir.exists():
+        shutil.rmtree(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    zip_path = "DEPTH_RESULT.zip"
+
+    preview_items = []
+
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for item in files:
+            path = item if isinstance(item, str) else getattr(item, "name", item)
+            stem = Path(path).stem
+
+            img = Image.open(path).convert("RGB")
+
+            # depth-estimation pipeline returns a dict with depth image
+            result = pipe(img)
+            depth_raw = result["depth"]
+
+            depth_np = normalize_depth(depth_raw)
+            depth_final = refine_depth(depth_np)
+
+            out_path = out_dir / f"{stem}_depth.png"
+            depth_final.save(out_path)
+            zipf.write(out_path, out_path.name)
+
+            preview_items.append((str(out_path), f"{stem}_depth"))
 
-out = Path("depth_clean_output") if out.exists(): shutil.rmtree(out) out.mkdir() zip_path = "DEPTH_AAA_ONLY.zip" with zipfile.ZipFile(zip_path, 'w') as zipf: for file in files: name = Path(file.name).stem img = Image.open(file.name).convert("RGB") # 💥 DEPTH DIRETO DO BG (VOCÊ JÁ VAI FAZER ISSO CERTO) depth_raw = pipe(img)["depth"] # normalização real depth_np = normalize_depth(depth_raw) # refinamento seguro depth_final = refine_depth(depth_np) # salvar path = out / f"depth_{name}.png" depth_final.save(path) zipf.write(path, path.name) print("✅ DEPTH PERFEITO GERADO") return zip_path ========================= UI ========================= 
+    # show first result as preview if available
+    preview = preview_items[0][0] if preview_items else None
+    return zip_path, preview
 
+# -------------------------
+# UI
+# -------------------------
 with gr.Blocks() as demo:
-gr.Markdown("# 🧠 DEPTH AAA (PARALLAX READY)")
+    gr.Markdown("# 🧠 DEPTH AAA (CLEAN REBUILD)")
 
-inp = gr.File(file_count="multiple") out = gr.File() btn = gr.Button("GERAR DEPTH") btn.click(fn=process, inputs=inp, outputs=out) 
+    inp = gr.File(file_count="multiple", type="filepath")
+    out_zip = gr.File(label="ZIP")
+    preview = gr.Image(label="Preview", type="filepath")
 
-demo.launch()
+    btn = gr.Button("GERAR DEPTH")
+    btn.click(fn=process, inputs=inp, outputs=[out_zip, preview])
 
+demo.launch()