Spaces:
Sleeping
Sleeping
File size: 4,123 Bytes
83f1324 3415f45 83f1324 3415f45 83f1324 3415f45 83f1324 3415f45 83f1324 3415f45 83f1324 3415f45 83f1324 3415f45 83f1324 3415f45 83f1324 3415f45 83f1324 3415f45 | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | import os
# Must be set BEFORE importing cv2 to enable the EXR codec.
os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1")
import tempfile
import spaces
import gradio as gr
import numpy as np
import torch
import cv2
import tifffile
import matplotlib
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
MODEL_ID = "LiheYoung/depth-anything-large-hf"
processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForDepthEstimation.from_pretrained(MODEL_ID)
model.eval()
cmap = matplotlib.colormaps.get_cmap("Spectral_r")
def _write_exr(path, depth_f32):
"""Write single-value float32 depth to EXR (RGB channels all = depth).
Try OpenCV first, then the OpenEXR package. Return path or None."""
try:
ok = cv2.imwrite(path, depth_f32) # cv2 writes float32 -> EXR
if ok and os.path.exists(path) and os.path.getsize(path) > 0:
return path
except Exception as e:
print("cv2 EXR write failed:", e)
try:
import OpenEXR, Imath
h, w = depth_f32.shape
header = OpenEXR.Header(w, h)
ftype = Imath.Channel(Imath.PixelType(Imath.PixelType.FLOAT))
header["channels"] = {c: ftype for c in ("R", "G", "B")}
out = OpenEXR.OutputFile(path, header)
buf = np.ascontiguousarray(depth_f32).tobytes()
out.writePixels({"R": buf, "G": buf, "B": buf})
out.close()
return path
except Exception as e:
print("OpenEXR write failed:", e)
return None
@spaces.GPU
@torch.no_grad()
def predict_depth(image: Image.Image):
if image is None:
return None, []
image = image.convert("RGB")
model.to("cuda")
inputs = processor(images=image, return_tensors="pt").to("cuda")
outputs = model(**inputs)
predicted_depth = outputs.predicted_depth
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=image.size[::-1], # (H, W)
mode="bicubic",
align_corners=False,
).squeeze()
# Raw float32 depth (full precision, NOT normalized) for TIFF/EXR.
depth = prediction.detach().cpu().numpy().astype(np.float32)
dmin, dmax = float(depth.min()), float(depth.max())
norm = (depth - dmin) / (dmax - dmin + 1e-8) # 0..1 for preview + 16-bit PNG
# 8-bit colored preview (display only — precision doesn't matter here).
colored = Image.fromarray((cmap(norm)[:, :, :3] * 255).astype(np.uint8))
tmp = tempfile.mkdtemp(prefix="depth_")
png16_path = os.path.join(tmp, "depth_16bit.png")
tiff_path = os.path.join(tmp, "depth_float32.tiff")
exr_path = os.path.join(tmp, "depth_float32.exr")
# 16-bit PNG: normalized depth mapped to the full 0..65535 range.
cv2.imwrite(png16_path, (norm * 65535.0).astype(np.uint16))
# 32-bit float TIFF: raw depth, fully lossless.
tifffile.imwrite(tiff_path, depth)
# 32-bit float EXR: raw depth (best-effort).
exr_ok = _write_exr(exr_path, depth)
files = [png16_path, tiff_path]
if exr_ok:
files.append(exr_ok)
return colored, files
title = "# Depth Anything (ZeroGPU) — high-bit-depth export"
description = (
"Monocular depth with Depth-Anything-Large on ZeroGPU.\n\n"
"**Downloads preserve full precision** — the colored image is an 8-bit preview only:\n"
"- `depth_16bit.png` — 16-bit grayscale, depth normalized to 0–65535\n"
"- `depth_float32.tiff` — 32-bit float, raw depth values (lossless)\n"
"- `depth_float32.exr` — 32-bit float OpenEXR, raw depth values"
)
with gr.Blocks(title="Depth Anything") as demo:
gr.Markdown(title)
gr.Markdown(description)
with gr.Row():
input_image = gr.Image(label="Input", type="pil")
depth_color = gr.Image(label="Depth (colored preview, 8-bit)", type="pil")
downloads = gr.Files(label="High-bit-depth downloads (PNG16 / TIFF32 / EXR32)")
run = gr.Button("Compute depth", variant="primary")
run.click(fn=predict_depth, inputs=input_image, outputs=[depth_color, downloads])
if __name__ == "__main__":
demo.queue().launch()
|