Update app.py

app.py

@@ -1,47 +1,274 @@
-import os
+import json
+import math
+import shutil
 import zipfile
 import tempfile
 from pathlib import Path
 
+import cv2
 import gradio as gr
+import numpy as np
+import torch
 from PIL import Image
 from transformers import pipeline
 
-# Depth Anything V2 é uma boa escolha para depth monocular.
-# A API do Transformers mostra o uso com pipeline("depth-estimation") e pipe(image)["depth"].
-pipe = pipeline(
-    task="depth-estimation",
-    model="depth-anything/Depth-Anything-V2-Small-hf"
-)
+# ----------------------------
+# Config
+# ----------------------------
+DEPTH_MODEL_ID = "depth-anything/Depth-Anything-V2-base-hf"
+DEFAULT_LAYER_COUNT = 6
+EXPORT_ROOT = Path(tempfile.gettempdir()) / "alvore_studio_exports"
+EXPORT_ROOT.mkdir(parents=True, exist_ok=True)
 
-def generate_cinematic_depth(files):
-    if not files:
-        return None
+DEVICE = 0 if torch.cuda.is_available() else -1
+DEPTH_PIPE = pipeline("depth-estimation", model=DEPTH_MODEL_ID, device=DEVICE)
 
-    out_dir = Path(tempfile.mkdtemp(prefix="parallax_depth_"))
-    zip_path = out_dir / "Leva_Parallax_Pro.zip"
 
-    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zipf:
-        for file_obj in files:
-            img = Image.open(file_obj.name).convert("RGB")
-            depth_img = pipe(img)["depth"]
+# ----------------------------
+# Utils
+# ----------------------------
+def pil_rgb(path: str) -> Image.Image:
+    return Image.open(path).convert("RGB")
 
-            base_name = Path(file_obj.name).stem
-            depth_name = f"depth_{base_name}.png"
-            depth_path = out_dir / depth_name
 
-            depth_img.save(depth_path)
-            zipf.write(depth_path, arcname=depth_name)
+def to_np_rgb(img: Image.Image) -> np.ndarray:
+    return np.array(img, dtype=np.uint8)
 
-    return str(zip_path)
 
-with gr.Blocks() as demo:
-    gr.Markdown("# 🎥 Eternity Engine v3.0 — Cinematic Depth")
-    gr.Markdown("Envie uma leva de imagens e baixe um ZIP com os mapas de profundidade.")
-    files = gr.File(file_count="multiple", label="1. Envie a leva de imagens")
-    out = gr.File(label="2. Baixe as máscaras de profundidade (ZIP)")
-    btn = gr.Button("Gerar Depth Maps")
+def save_png(path: Path, arr: np.ndarray) -> None:
+    if arr.ndim == 2:
+        Image.fromarray(arr.astype(np.uint8), mode="L").save(path)
+    elif arr.shape[2] == 3:
+        Image.fromarray(arr.astype(np.uint8), mode="RGB").save(path)
+    elif arr.shape[2] == 4:
+        Image.fromarray(arr.astype(np.uint8), mode="RGBA").save(path)
+    else:
+        raise ValueError(f"Formato inesperado: {arr.shape}")
 
-    btn.click(generate_cinematic_depth, inputs=files, outputs=out)
 
-demo.launch()
+def resize_to(arr: np.ndarray, size: tuple[int, int]) -> np.ndarray:
+    w, h = size
+    return cv2.resize(arr, (w, h), interpolation=cv2.INTER_CUBIC)
+
+
+def normalize_depth(depth: np.ndarray) -> np.ndarray:
+    depth = depth.astype(np.float32)
+    mn, mx = float(depth.min()), float(depth.max())
+    return (depth - mn) / (mx - mn + 1e-8)
+
+
+def infer_depth(img_pil: Image.Image) -> np.ndarray:
+    with torch.no_grad():
+        out = DEPTH_PIPE(img_pil)
+
+    pred = out["predicted_depth"]
+    if hasattr(pred, "detach"):
+        pred = pred.detach().cpu().numpy()
+    pred = np.squeeze(pred).astype(np.float32)
+
+    depth = normalize_depth(pred)
+    return depth
+
+
+def auto_inpaint_from_depth(orig_rgb: np.ndarray, depth: np.ndarray) -> np.ndarray:
+    """
+    Se o inpaint não vier pronto, cria um fallback razoável baseado na região mais próxima.
+    """
+    mask = (depth > 0.67).astype(np.uint8) * 255
+    kernel = np.ones((9, 9), np.uint8)
+    mask = cv2.dilate(mask, kernel, iterations=2)
+    mask = cv2.GaussianBlur(mask, (0, 0), 3.0)
+    return cv2.inpaint(orig_rgb, mask, 3, cv2.INPAINT_TELEA)
+
+
+def make_layer_stack(orig_rgb: np.ndarray, inpaint_rgb: np.ndarray, depth: np.ndarray, layer_count: int) -> list[np.ndarray]:
+    """
+    Cria camadas RGBA.
+    1 = perto, 0 = longe.
+    """
+    centers = np.linspace(0.0, 1.0, layer_count, dtype=np.float32)
+    sigma = max(0.08, 0.42 / max(layer_count, 2))
+
+    weights = []
+    for c in centers:
+        w = np.exp(-0.5 * ((depth - c) / sigma) ** 2)
+        weights.append(w)
+    weights = np.stack(weights, axis=-1)
+    weights /= (weights.sum(axis=-1, keepdims=True) + 1e-8)
+
+    layers = []
+    for i, c in enumerate(centers):
+        alpha = weights[..., i]
+        alpha = cv2.GaussianBlur(alpha, (0, 0), 1.1)
+        alpha = np.clip(alpha ** 0.85, 0.0, 1.0)
+
+        # fundo -> inpaint, frente -> original
+        src = (inpaint_rgb.astype(np.float32) * (1.0 - c) + orig_rgb.astype(np.float32) * c)
+        rgb = np.clip(src, 0, 255).astype(np.uint8)
+        a8 = (alpha * 255.0).astype(np.uint8)
+        rgba = np.dstack([rgb, a8])
+        layers.append(rgba)
+
+    return layers
+
+
+def create_scene_assets(orig_path: str, inpaint_path: str | None, depth_path: str | None, scene_idx: int, layer_count: int, use_ai_depth: bool):
+    scene_name = f"scene_{scene_idx:03d}"
+    scene_dir = EXPORT_ROOT / scene_name
+    scene_dir.mkdir(parents=True, exist_ok=True)
+
+    orig_img = pil_rgb(orig_path)
+    orig_rgb = to_np_rgb(orig_img)
+    size = orig_img.size  # (w, h)
+
+    # Inpaint opcional
+    if inpaint_path:
+        inpaint_img = pil_rgb(inpaint_path).resize(size, Image.LANCZOS)
+        inpaint_rgb = to_np_rgb(inpaint_img)
+    else:
+        inpaint_rgb = None
+
+    # Depth
+    if depth_path:
+        depth_img = Image.open(depth_path).convert("L").resize(size, Image.LANCZOS)
+        depth = np.array(depth_img, dtype=np.float32) / 255.0
+    elif use_ai_depth:
+        depth = infer_depth(orig_img)
+        if depth.shape != (size[1], size[0]):
+            depth = resize_to(depth, size)
+    else:
+        depth = np.full((size[1], size[0]), 0.5, dtype=np.float32)
+
+    depth = np.clip(depth, 0.0, 1.0)
+
+    # Inpaint fallback
+    if inpaint_rgb is None:
+        inpaint_rgb = auto_inpaint_from_depth(orig_rgb, depth)
+
+    # Assets principais
+    save_png(scene_dir / f"{scene_name}_original.png", orig_rgb)
+    save_png(scene_dir / f"{scene_name}_inpaint.png", inpaint_rgb)
+    save_png(scene_dir / f"{scene_name}_depth.png", (depth * 255.0).astype(np.uint8))
+
+    # Camadas RGBA
+    layer_files = []
+    layers = make_layer_stack(orig_rgb, inpaint_rgb, depth, layer_count=layer_count)
+    for li, layer_rgba in enumerate(layers):
+        fname = f"{scene_name}_layer_{li:02d}.png"
+        save_png(scene_dir / fname, layer_rgba)
+        layer_files.append(fname)
+
+    # Manifest
+    manifest = {
+        "scene": scene_name,
+        "original": f"{scene_name}_original.png",
+        "inpaint": f"{scene_name}_inpaint.png",
+        "depth": f"{scene_name}_depth.png",
+        "layers": layer_files,
+        "layer_count": layer_count,
+        "size": {"w": size[0], "h": size[1]},
+    }
+    (scene_dir / f"{scene_name}_manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
+
+    # Thumb
+    thumb = orig_rgb.copy()
+    thumb = cv2.resize(thumb, (960, int(960 * size[1] / max(size[0], 1))), interpolation=cv2.INTER_AREA)
+    save_png(scene_dir / f"{scene_name}_thumb.png", thumb)
+
+    return scene_dir, manifest
+
+
+def zip_folder(folder: Path) -> Path:
+    zip_path = folder.with_suffix(".zip")
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for p in folder.rglob("*"):
+            if p.is_file():
+                zf.write(p, arcname=p.relative_to(folder))
+    return zip_path
+
+
+def process_batch(originals, inpaints, depths, layer_count, use_ai_depth):
+    originals = list(originals or [])
+    inpaints = list(inpaints or [])
+    depths = list(depths or [])
+
+    if not originals:
+        raise gr.Error("Envie pelo menos uma imagem original.")
+
+    session_dir = EXPORT_ROOT / f"session_{tempfile.mkstemp(prefix='alvore_')[1].split('/')[-1]}"
+    session_dir.mkdir(parents=True, exist_ok=True)
+
+    scene_rows = []
+    preview_paths = []
+
+    for i, orig in enumerate(originals, start=1):
+        inp = inpaints[i - 1] if i - 1 < len(inpaints) else None
+        dep = depths[i - 1] if i - 1 < len(depths) else None
+
+        scene_dir, manifest = create_scene_assets(
+            orig_path=orig,
+            inpaint_path=inp,
+            depth_path=dep,
+            scene_idx=i,
+            layer_count=layer_count,
+            use_ai_depth=use_ai_depth,
+        )
+
+        preview_paths.append(str(scene_dir / f"scene_{i:03d}_thumb.png"))
+        scene_rows.append(
+            f"Cena {i:03d} | original={Path(orig).name} | "
+            f"inpaint={Path(inp).name if inp else 'auto'} | "
+            f"depth={Path(dep).name if dep else ('IA' if use_ai_depth else 'flat')} | "
+            f"layers={len(manifest['layers'])}"
+        )
+
+    zip_path = zip_folder(session_dir)
+
+    return (
+        preview_paths,
+        str(zip_path),
+        "\n".join(scene_rows),
+        str(session_dir),
+    )
+
+
+# ----------------------------
+# Gradio UI
+# ----------------------------
+with gr.Blocks(title="ALVORE STUDIO LAYER BUILDER") as demo:
+    gr.Markdown(
+        """
+# ALVORE STUDIO LAYER BUILDER
+
+Use assim:
+- **Imagem original**: cena completa, com personagem.
+- **Inpaint**: mesma cena sem personagem, se você já tiver.
+- **Depth**: mapa branco/perto, preto/longe. Se não vier, a IA gera.
+- **Layer count**: 6 é um bom ponto de partida de estúdio.
+        """
+    )
+
+    with gr.Row():
+        originals = gr.File(label="1. IMAGENS ORIGINAIS", file_count="multiple", file_types=["image"], type="filepath")
+        inpaints = gr.File(label="2. INPAINTS", file_count="multiple", file_types=["image"], type="filepath")
+        depths = gr.File(label="3. DEPTH MAPS", file_count="multiple", file_types=["image"], type="filepath")
+
+    with gr.Row():
+        layer_count = gr.Slider(4, 8, value=6, step=1, label="Número de camadas")
+        use_ai_depth = gr.Checkbox(value=True, label="Gerar depth por IA quando faltar")
+
+    run_btn = gr.Button("GERAR CAMADAS EM LOTE")
+
+    gallery = gr.Gallery(label="Prévia", columns=3, height=240)
+    zip_out = gr.File(label="ZIP do lote")
+    log_out = gr.Textbox(label="Log", lines=8)
+    folder_out = gr.Textbox(label="Pasta da sessão", lines=1)
+
+    run_btn.click(
+        fn=process_batch,
+        inputs=[originals, inpaints, depths, layer_count, use_ai_depth],
+        outputs=[gallery, zip_out, log_out, folder_out],
+    )
+
+if __name__ == "__main__":
+    demo.launch()