Spaces:
Sleeping
Sleeping
File size: 4,104 Bytes
9eca32a 7c27268 9d05699 7c27268 9d05699 7c27268 9d05699 7c27268 9d05699 02d2686 9d05699 02d2686 9d05699 3a13793 48d0430 9d05699 7c27268 9d05699 7c27268 6ceac94 9d05699 7c27268 187c6c8 7c27268 9d05699 7c27268 6ceac94 7c27268 9d05699 7c27268 9d05699 7c27268 187c6c8 7c27268 187c6c8 7c27268 9d05699 7c27268 |
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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
from huggingface_hub import whoami
import gradio as gr
import cv2
import matplotlib
import numpy as np
import os
from PIL import Image
import spaces
import torch
import tempfile
from gradio_imageslider import ImageSlider
from huggingface_hub import hf_hub_download
from depth_anything_v2.dpt import DepthAnythingV2
css = """
#img-display-container {
max-height: 100vh;
}
#img-display-input {
max-height: 80vh;
}
#img-display-output {
max-height: 80vh;
}
#download {
height: 62px;
}
"""
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model_configs = {
'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]},
}
encoder2name = {
'vits': 'Small',
'vitb': 'Base',
'vitl': 'Large',
'vitg': 'Giant',
}
encoder = 'vits'
model_name = encoder2name[encoder]
model = DepthAnythingV2(**model_configs[encoder])
filepath = hf_hub_download(
repo_id=f"depth-anything/Depth-Anything-V2-{model_name}",
filename=f"depth_anything_v2_{encoder}.pth",
repo_type="model"
)
state_dict = torch.load(filepath, map_location="cpu")
model.load_state_dict(state_dict)
model = model.to(DEVICE).eval()
title = "# Depth Anything V2"
description = """Official demo for **Depth Anything V2**.
Please refer to our [paper](https://arxiv.org/abs/2406.09414),
[project page](https://depth-anything-v2.github.io),
and [github](https://github.com/DepthAnything/Depth-Anything-V2) for more details."""
@spaces.GPU
def predict_depth(image):
return model.infer_image(image)
# -------------------------------------
# OLD GRADIO COMPATIBILITY PATCH
# -------------------------------------
if not hasattr(gr.Blocks, "get_api_info"):
gr.Blocks.get_api_info = lambda self: {}
# -------------------------------------
with gr.Blocks(css=css) as demo:
gr.Markdown(title)
gr.Markdown(description)
gr.Markdown("### Depth Prediction demo")
with gr.Row():
input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input')
depth_image_slider = ImageSlider(
label="Depth Map with Slider View",
elem_id='img-display-output',
position=0.5
)
submit = gr.Button(value="Compute Depth")
gray_depth_file = gr.File(label="Grayscale depth map", elem_id="download")
raw_file = gr.File(label="16-bit raw output (can be considered as disparity)", elem_id="download")
cmap = matplotlib.colormaps.get_cmap('Spectral_r')
def on_submit(image):
original_image = image.copy()
depth = predict_depth(image[:, :, ::-1])
raw_depth = Image.fromarray(depth.astype('uint16'))
tmp_raw_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
raw_depth.save(tmp_raw_depth.name)
depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
depth = depth.astype(np.uint8)
colored_depth = (cmap(depth)[:, :, :3] * 255).astype(np.uint8)
gray_depth = Image.fromarray(depth)
tmp_gray_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
gray_depth.save(tmp_gray_depth.name)
return [(original_image, colored_depth), tmp_gray_depth.name, tmp_raw_depth.name]
submit.click(
on_submit,
inputs=[input_image],
outputs=[depth_image_slider, gray_depth_file, raw_file]
)
if os.path.exists('assets/examples'):
example_files = sorted(os.listdir('assets/examples'))
example_files = [os.path.join('assets/examples', f) for f in example_files]
gr.Examples(
cache_examples=False,
examples=example_files,
inputs=[input_image],
outputs=[depth_image_slider, gray_depth_file, raw_file],
fn=on_submit
)
if __name__ == '__main__':
demo.queue().launch(share=True)
|