Spaces:
Running
Running
File size: 2,301 Bytes
1fbc150 5f8c34f 1fbc150 5f8c34f f647e40 5f8c34f 1fbc150 |
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 |
import depth_pro
import gradio as gr
import matplotlib.cm as cm
import numpy as np
from depth_pro.depth_pro import DepthProConfig
from PIL import Image
MARKDOWN = """
<div align="center">
<h2><a href="https://arxiv.org/abs/2410.02073">Depth Pro: Sharp Monocular Metric Depth in Less Than a Second</a></h2>
</div>
"""
def run(input_image_path):
config = DepthProConfig(
patch_encoder_preset="dinov2l16_384",
image_encoder_preset="dinov2l16_384",
checkpoint_uri="./depth_pro.pt",
decoder_features=256,
use_fov_head=True,
fov_encoder_preset="dinov2l16_384",
)
# Load model and preprocessing transform
model, transform = depth_pro.create_model_and_transforms(config=config)
model.eval()
# Load and preprocess an image
image, _, f_px = depth_pro.load_rgb(input_image_path)
image = transform(image)
# Run inference
prediction = model.infer(image, f_px=f_px)
depth_map = prediction["depth"].squeeze().cpu().numpy()
focallength_px = prediction["focallength_px"]
depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
colormap = cm.get_cmap("viridis")
depth_map = colormap(depth_map)
depth_map = (depth_map[:, :, :3] * 255).astype(np.uint8)
depth_map = Image.fromarray(depth_map)
return depth_map, focallength_px.item()
with gr.Blocks() as demo:
gr.Markdown(MARKDOWN)
with gr.Row():
with gr.Column():
input_image_path = gr.Image(
label="Input Image", type="filepath", sources=["upload"]
)
with gr.Column():
with gr.Column():
output_depth_map = gr.Image(label="Depth Map")
output_focal_length = gr.Number(label="Focal Length")
with gr.Row():
btn = gr.Button("Run")
btn.click(
run, inputs=[input_image_path], outputs=[output_depth_map, output_focal_length]
)
examples = gr.Examples(
examples=[
"assets/input_one.webp",
"assets/input_two.webp",
"assets/input_three.webp",
],
fn=run,
inputs=[input_image_path],
outputs=[output_depth_map, output_focal_length],
cache_examples=True,
)
demo.launch()
|