dpt-large-demo / app.py
Celeste-cj
turn off share
9537292
raw
history blame
1.6 kB
import os
from PIL import Image
import gradio as gr
import numpy as np
import torch
from transformers import AutoModelForDepthEstimation, DPTImageProcessor
processor = DPTImageProcessor.from_pretrained(
"Intel/dpt-large")
model = AutoModelForDepthEstimation.from_pretrained("Intel/dpt-large")
def main(image, input_size=384):
# prepare image for the model
inputs = processor(images=image, return_tensors="pt", do_resize=True, size=(
input_size, input_size), keep_aspect_ratio=True)
print(type(inputs), inputs.data["pixel_values"].shape)
# do inference
with torch.no_grad():
outputs = model(**inputs)
predicted_depth = outputs.predicted_depth
# interpolate to original size
prediction = torch.nn.functional.interpolate(predicted_depth.unsqueeze(
1), size=image.shape[:-1], mode="bicubic").squeeze()
output = prediction.cpu().numpy().copy()
formatted = (output * 255 / output.max()).astype("uint8")
depth = Image.fromarray(formatted)
return depth
title = "Demo: monocular depth estimation with DPT"
description = "This demo uses <a href='https://huggingface.co/Intel/dpt-large' target='_blank'>DPT</a> to estimate depth from monocular image."
examples = [[f"examples/{file}"]
for file in os.listdir("examples") if file[0] != "."]
demo = gr.Interface(fn=main, inputs=[gr.Image(label="Input Image"), gr.Slider(128, 512, value=384, label="Input Size")], outputs="image",
title=title, description=description, examples=examples, cache_examples=True)
demo.launch(debug=True, share=False)