Spaces:
Running
Running
import gradio as gr | |
from transformers import DPTFeatureExtractor, DPTForDepthEstimation | |
import torch | |
import numpy as np | |
from PIL import Image | |
import open3d as o3d | |
from pathlib import Path | |
import os | |
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large") | |
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large") | |
def process_image(image_path): | |
image_path = Path(image_path) | |
image = Image.open(image_path) | |
# prepare image for the model | |
encoding = feature_extractor(image, return_tensors="pt") | |
# forward pass | |
with torch.no_grad(): | |
outputs = model(**encoding) | |
predicted_depth = outputs.predicted_depth | |
# interpolate to original size | |
prediction = torch.nn.functional.interpolate( | |
predicted_depth.unsqueeze(1), | |
size=image.size[::-1], | |
mode="bicubic", | |
align_corners=False, | |
).squeeze() | |
output = prediction.cpu().numpy() | |
depth_image = (output * 255 / np.max(output)).astype("uint8") | |
img = Image.fromarray(depth_image) | |
return [img] | |
title = "Demo: zero-shot depth estimation with DPT + 3D Point Cloud" | |
description = "This demo is a variation from the original <a href='https://huggingface.co/spaces/nielsr/dpt-depth-estimation' target='_blank'>DPT Demo</a>. It uses the DPT model to predict the depth of an image and then uses 3D Point Cloud to create a 3D object." | |
examples = [["examples/" + img] for img in os.listdir("examples/")] | |
iface = gr.Interface( | |
fn=process_image, | |
inputs=[gr.inputs.Image(type="filepath", label="Input Image")], | |
outputs=[predicted_depth], | |
title=title, | |
description=description, | |
examples=examples, | |
allow_flagging="never", | |
cache_examples=False, | |
) | |
iface.launch(debug=True, show_api=False) |