ZoeDepth_slim / app.py
chongjie's picture
Update app.py
d33303c
import gradio as gr
import torch
from utils import colorize
from PIL import Image
import tempfile
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model = torch.hub.load('isl-org/ZoeDepth', "ZoeD_N", pretrained=True).to(DEVICE).eval()
def predict_depth(model, image):
depth = model.infer_pil(image)
return depth
def on_submit(image):
depth = predict_depth(model, image)
colored_depth = colorize(depth, cmap='gray_r')
tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
raw_depth = Image.fromarray((depth*256).astype('uint16'))
raw_depth.save(tmp.name)
return [colored_depth, tmp.name]
iface = gr.Interface(
fn=on_submit,
inputs=gr.inputs.Image(type='pil', label="Input Image"),
outputs=[
gr.outputs.Image(type='numpy', label="Depth Map"),
gr.outputs.File(label="16-bit raw depth, multiplier:256")
],
title="# ZoeDepth",
description="""Unofficial demo for **ZoeDepth: Zero-shot Transfer by Combining Relative and Metric Depth**.""",
css="""
#img-display-container {
max-height: 50vh;
}
#img-display-input {
max-height: 40vh;
}
#img-display-output {
max-height: 40vh;
}
"""
)
if __name__ == '__main__':
iface.launch()