File size: 2,907 Bytes
fe45438
6e49e13
 
 
 
 
85c906a
08ba32c
 
fe45438
08ba32c
fe45438
08ba32c
 
 
 
 
 
 
fe45438
6e49e13
08ba32c
 
 
85c906a
cf63ce6
08ba32c
 
 
 
 
 
 
 
 
 
 
 
cf63ce6
08ba32c
 
fe45438
08ba32c
 
fe45438
08ba32c
fe45438
08ba32c
 
 
 
 
fe45438
6e49e13
08ba32c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a1fafd3
08ba32c
 
 
ca0222d
08ba32c
 
 
cf63ce6
08ba32c
cf63ce6
08ba32c
 
 
a1fafd3
08ba32c
 
cf63ce6
08ba32c
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import gradio as gr
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 (CLEAN REBUILD)")

# -------------------------
# 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=MODEL_ID,
    device=device
)

# -------------------------
# 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))

    # leve blur para suavizar banding
    img = img.filter(ImageFilter.GaussianBlur(1.0))

    # melhora contraste sem destruir detalhe
    img = ImageOps.autocontrast(img, cutoff=0.3)

    return img

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, 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"))

    # 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 (CLEAN REBUILD)")

    inp = gr.File(file_count="multiple", type="filepath")
    out_zip = gr.File(label="ZIP")
    preview = gr.Image(label="Preview", type="filepath")

    btn = gr.Button("GERAR DEPTH")
    btn.click(fn=process, inputs=inp, outputs=[out_zip, preview])

demo.launch()