Spaces:
Build error
Build error
import functools | |
import os | |
import shutil | |
import sys | |
import git | |
import gradio as gr | |
import numpy as np | |
import torch as torch | |
from PIL import Image | |
from gradio_imageslider import ImageSlider | |
def process( | |
pipe, | |
path_input, | |
ensemble_size, | |
denoise_steps, | |
processing_res, | |
path_out_16bit=None, | |
path_out_fp32=None, | |
path_out_vis=None, | |
_input_3d_plane_near=None, | |
_input_3d_plane_far=None, | |
_input_3d_embossing=None, | |
_input_3d_filter_size=None, | |
_input_3d_frame_near=None, | |
): | |
if path_out_vis is not None: | |
return ( | |
[path_out_16bit, path_out_vis], | |
[path_out_16bit, path_out_fp32, path_out_vis], | |
) | |
input_image = Image.open(path_input) | |
pipe_out = pipe( | |
input_image, | |
ensemble_size=ensemble_size, | |
denoising_steps=denoise_steps, | |
processing_res=processing_res, | |
batch_size=1 if processing_res == 0 else 0, | |
show_progress_bar=True, | |
) | |
depth_pred = pipe_out.depth_np | |
depth_colored = pipe_out.depth_colored | |
depth_16bit = (depth_pred * 65535.0).astype(np.uint16) | |
path_output_dir = os.path.splitext(path_input)[0] + "_output" | |
os.makedirs(path_output_dir, exist_ok=True) | |
name_base = os.path.splitext(os.path.basename(path_input))[0] | |
path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_depth_fp32.npy") | |
path_out_16bit = os.path.join(path_output_dir, f"{name_base}_depth_16bit.png") | |
path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.png") | |
np.save(path_out_fp32, depth_pred) | |
Image.fromarray(depth_16bit).save(path_out_16bit, mode="I;16") | |
depth_colored.save(path_out_vis) | |
return ( | |
[path_out_16bit, path_out_vis], | |
[path_out_16bit, path_out_fp32, path_out_vis], | |
) | |
def process_3d( | |
input_image, | |
files, | |
size_longest_px, | |
size_longest_cm, | |
filter_size, | |
plane_near, | |
plane_far, | |
embossing, | |
frame_thickness, | |
frame_near, | |
frame_far, | |
): | |
if input_image is None or len(files) < 1: | |
raise gr.Error("Please upload an image (or use examples) and compute depth first") | |
if plane_near >= plane_far: | |
raise gr.Error("NEAR plane must have a value smaller than the FAR plane") | |
def _process_3d(size_longest_px, filter_size, vertex_colors, scene_lights, output_model_scale=None): | |
image_rgb = input_image | |
image_depth = files[0] | |
image_rgb_basename, image_rgb_ext = os.path.splitext(image_rgb) | |
image_depth_basename, image_depth_ext = os.path.splitext(image_depth) | |
image_rgb_content = Image.open(image_rgb) | |
image_rgb_w, image_rgb_h = image_rgb_content.width, image_rgb_content.height | |
image_rgb_d = max(image_rgb_w, image_rgb_h) | |
image_new_w = size_longest_px * image_rgb_w // image_rgb_d | |
image_new_h = size_longest_px * image_rgb_h // image_rgb_d | |
image_rgb_new = image_rgb_basename + f"_{size_longest_px}" + image_rgb_ext | |
image_depth_new = image_depth_basename + f"_{size_longest_px}" + image_depth_ext | |
image_rgb_content.resize((image_new_w, image_new_h), Image.LANCZOS).save( | |
image_rgb_new | |
) | |
Image.open(image_depth).resize((image_new_w, image_new_h), Image.LANCZOS).save( | |
image_depth_new | |
) | |
path_glb, path_stl = extrude_depth_3d( | |
image_rgb_new, | |
image_depth_new, | |
output_model_scale=size_longest_cm * 10 if output_model_scale is None else output_model_scale, | |
filter_size=filter_size, | |
coef_near=plane_near, | |
coef_far=plane_far, | |
emboss=embossing / 100, | |
f_thic=frame_thickness / 100, | |
f_near=frame_near / 100, | |
f_back=frame_far / 100, | |
vertex_colors=vertex_colors, | |
scene_lights=scene_lights, | |
) | |
return path_glb, path_stl | |
path_viewer_glb, _ = _process_3d(256, filter_size, vertex_colors=False, scene_lights=True, output_model_scale=1) | |
path_files_glb, path_files_stl = _process_3d(size_longest_px, filter_size, vertex_colors=True, scene_lights=False) | |
# sanitize 3d viewer glb path to keep babylon.js happy | |
path_viewer_glb_sanitized = os.path.join(os.path.dirname(path_viewer_glb), "preview.glb") | |
if path_viewer_glb_sanitized != path_viewer_glb: | |
os.rename(path_viewer_glb, path_viewer_glb_sanitized) | |
path_viewer_glb = path_viewer_glb_sanitized | |
return path_viewer_glb, [path_files_glb, path_files_stl] | |
def run_demo_server(pipe): | |
process_pipe = functools.partial(process, pipe) | |
os.environ["GRADIO_ALLOW_FLAGGING"] = "never" | |
with gr.Blocks( | |
analytics_enabled=False, | |
title="Marigold Depth Estimation", | |
css=""" | |
#download { | |
height: 118px; | |
} | |
.slider .inner { | |
width: 5px; | |
background: #FFF; | |
} | |
.viewport { | |
aspect-ratio: 4/3; | |
} | |
""", | |
) as demo: | |
gr.Markdown( | |
""" | |
<h1 align="center">Marigold Depth Estimation</h1> | |
<p align="center"> | |
<a title="Website" href="https://marigoldmonodepth.github.io/" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
<img src="https://www.obukhov.ai/img/badges/badge-website.svg"> | |
</a> | |
<a title="arXiv" href="https://arxiv.org/abs/2312.02145" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
<img src="https://www.obukhov.ai/img/badges/badge-pdf.svg"> | |
</a> | |
<a title="Github" href="https://github.com/prs-eth/marigold" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
<img src="https://img.shields.io/github/stars/prs-eth/marigold?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars"> | |
</a> | |
<a title="Social" href="https://twitter.com/antonobukhov1" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
<img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social"> | |
</a> | |
</p> | |
<p align="justify"> | |
Marigold is the new state-of-the-art depth estimator for images in the wild. | |
Upload your image into the <b>left</b> side, or click any of the <b>examples</b> below. | |
The result will be computed and appear on the <b>right</b> in the output comparison window. | |
<b style="color: red;">NEW</b>: Scroll down to the new 3D printing part of the demo! | |
</p> | |
""" | |
) | |
with gr.Row(): | |
with gr.Column(): | |
input_image = gr.Image( | |
label="Input Image", | |
type="filepath", | |
) | |
with gr.Accordion("Advanced options", open=False): | |
ensemble_size = gr.Slider( | |
label="Ensemble size", | |
minimum=1, | |
maximum=20, | |
step=1, | |
value=10, | |
) | |
denoise_steps = gr.Slider( | |
label="Number of denoising steps", | |
minimum=1, | |
maximum=20, | |
step=1, | |
value=10, | |
) | |
processing_res = gr.Radio( | |
[ | |
("Native", 0), | |
("Recommended", 768), | |
], | |
label="Processing resolution", | |
value=768, | |
) | |
input_output_16bit = gr.File( | |
label="Predicted depth (16-bit)", | |
visible=False, | |
) | |
input_output_fp32 = gr.File( | |
label="Predicted depth (32-bit)", | |
visible=False, | |
) | |
input_output_vis = gr.File( | |
label="Predicted depth (red-near, blue-far)", | |
visible=False, | |
) | |
with gr.Row(): | |
submit_btn = gr.Button(value="Compute Depth", variant="primary") | |
clear_btn = gr.Button(value="Clear") | |
with gr.Column(): | |
output_slider = ImageSlider( | |
label="Predicted depth (red-near, blue-far)", | |
type="filepath", | |
show_download_button=True, | |
show_share_button=True, | |
interactive=False, | |
elem_classes="slider", | |
position=0.25, | |
) | |
files = gr.Files( | |
label="Depth outputs", | |
elem_id="download", | |
interactive=False, | |
) | |
demo_3d_header = gr.Markdown( | |
""" | |
<h3 align="center">3D Printing Depth Maps</h3> | |
<p align="justify"> | |
This part of the demo uses Marigold depth maps estimated in the previous step to create a | |
3D-printable model. The models are watertight, with correct normals, and exported in the STL format. | |
We recommended creating the first model with the default parameters and iterating on it until the best | |
result (see Pro Tips below). | |
</p> | |
""", | |
render=False, | |
) | |
demo_3d = gr.Row(render=False) | |
with demo_3d: | |
with gr.Column(): | |
with gr.Accordion("3D printing demo: Main options", open=True): | |
plane_near = gr.Slider( | |
label="Relative position of the near plane (between 0 and 1)", | |
minimum=0.0, | |
maximum=1.0, | |
step=0.001, | |
value=0.0, | |
) | |
plane_far = gr.Slider( | |
label="Relative position of the far plane (between near and 1)", | |
minimum=0.0, | |
maximum=1.0, | |
step=0.001, | |
value=1.0, | |
) | |
embossing = gr.Slider( | |
label="Embossing level", | |
minimum=0, | |
maximum=100, | |
step=1, | |
value=20, | |
) | |
with gr.Accordion("3D printing demo: Advanced options", open=False): | |
size_longest_px = gr.Slider( | |
label="Size (px) of the longest side", | |
minimum=256, | |
maximum=1024, | |
step=256, | |
value=512, | |
) | |
size_longest_cm = gr.Slider( | |
label="Size (cm) of the longest side", | |
minimum=1, | |
maximum=100, | |
step=1, | |
value=10, | |
) | |
filter_size = gr.Slider( | |
label="Size (px) of the smoothing filter", | |
minimum=1, | |
maximum=5, | |
step=2, | |
value=3, | |
) | |
frame_thickness = gr.Slider( | |
label="Frame thickness", | |
minimum=0, | |
maximum=100, | |
step=1, | |
value=5, | |
) | |
frame_near = gr.Slider( | |
label="Frame's near plane offset", | |
minimum=-100, | |
maximum=100, | |
step=1, | |
value=1, | |
) | |
frame_far = gr.Slider( | |
label="Frame's far plane offset", | |
minimum=1, | |
maximum=10, | |
step=1, | |
value=1, | |
) | |
with gr.Row(): | |
submit_3d = gr.Button(value="Create 3D", variant="primary") | |
clear_3d = gr.Button(value="Clear 3D") | |
gr.Markdown( | |
""" | |
<h5 align="center">Pro Tips</h5> | |
<ol> | |
<li><b>Re-render with new parameters</b>: Click "Clear 3D" and then "Create 3D".</li> | |
<li><b>Adjust 3D scale and cut-off focus</b>: Set the frame's near plane offset to the | |
minimum and use 3D preview to evaluate depth scaling. Repeat until the scale is correct and | |
everything important is in the focus. Set the optimal value for frame's near | |
plane offset as a last step.</li> | |
<li><b>Increase details</b>: Decrease size of the smoothing filter (also increases noise).</li> | |
</ol> | |
""" | |
) | |
with gr.Column(): | |
viewer_3d = gr.Model3D( | |
camera_position=(75.0, 90.0, 1.25), | |
elem_classes="viewport", | |
label="3D preview (low-res, relief highlight)", | |
interactive=False, | |
) | |
files_3d = gr.Files( | |
label="3D model outputs (high-res)", | |
elem_id="download", | |
interactive=False, | |
) | |
blocks_settings_depth = [ensemble_size, denoise_steps, processing_res] | |
blocks_settings_3d = [plane_near, plane_far, embossing, size_longest_px, size_longest_cm, filter_size, | |
frame_thickness, frame_near, frame_far] | |
blocks_settings = blocks_settings_depth + blocks_settings_3d | |
map_id_to_default = {b._id: b.value for b in blocks_settings} | |
inputs = [ | |
input_image, | |
ensemble_size, | |
denoise_steps, | |
processing_res, | |
input_output_16bit, | |
input_output_fp32, | |
input_output_vis, | |
plane_near, | |
plane_far, | |
embossing, | |
filter_size, | |
frame_near, | |
] | |
outputs = [ | |
submit_btn, | |
input_image, | |
output_slider, | |
files, | |
] | |
def submit_depth_fn(*args): | |
out = list(process_pipe(*args)) | |
out = [gr.Button(interactive=False), gr.Image(interactive=False)] + out | |
return out | |
submit_btn.click( | |
fn=submit_depth_fn, | |
inputs=inputs, | |
outputs=outputs, | |
concurrency_limit=1, | |
) | |
gr.Examples( | |
fn=submit_depth_fn, | |
examples=[ | |
[ | |
"files/bee.jpg", | |
10, # ensemble_size | |
10, # denoise_steps | |
768, # processing_res | |
"files/bee_depth_16bit.png", | |
"files/bee_depth_fp32.npy", | |
"files/bee_depth_colored.png", | |
0.0, # plane_near | |
0.5, # plane_far | |
20, # embossing | |
3, # filter_size | |
0, # frame_near | |
], | |
], | |
inputs=inputs, | |
outputs=outputs, | |
cache_examples=True, | |
) | |
demo_3d_header.render() | |
demo_3d.render() | |
def clear_fn(): | |
out = [] | |
for b in blocks_settings: | |
out.append(map_id_to_default[b._id]) | |
out += [ | |
gr.Button(interactive=True), | |
gr.Button(interactive=True), | |
gr.Image(value=None, interactive=True), | |
None, None, None, None, None, None, None, | |
] | |
return out | |
clear_btn.click( | |
fn=clear_fn, | |
inputs=[], | |
outputs=blocks_settings + [ | |
submit_btn, | |
submit_3d, | |
input_image, | |
input_output_16bit, | |
input_output_fp32, | |
input_output_vis, | |
output_slider, | |
files, | |
viewer_3d, | |
files_3d, | |
], | |
) | |
def submit_3d_fn(*args): | |
out = list(process_3d(*args)) | |
out = [gr.Button(interactive=False)] + out | |
return out | |
submit_3d.click( | |
fn=submit_3d_fn, | |
inputs=[ | |
input_image, | |
files, | |
size_longest_px, | |
size_longest_cm, | |
filter_size, | |
plane_near, | |
plane_far, | |
embossing, | |
frame_thickness, | |
frame_near, | |
frame_far, | |
], | |
outputs=[submit_3d, viewer_3d, files_3d], | |
concurrency_limit=1, | |
) | |
def clear_3d_fn(): | |
return [gr.Button(interactive=True), None, None] | |
clear_3d.click( | |
fn=clear_3d_fn, | |
inputs=[], | |
outputs=[submit_3d, viewer_3d, files_3d], | |
) | |
demo.queue( | |
api_open=False, | |
).launch( | |
server_name="0.0.0.0", | |
server_port=7860, | |
) | |
def main(): | |
REPO_URL = "https://github.com/lemonaddie/geowizard.git" | |
CHECKPOINT = "lemonaddie/Geowizard" | |
REPO_DIR = "geowizard" | |
if os.path.isdir(REPO_DIR): | |
shutil.rmtree(REPO_DIR) | |
repo = git.Repo.clone_from(REPO_URL, REPO_DIR) | |
sys.path.append(os.path.join(os.getcwd(), REPO_DIR)) | |
from pipeline.depth_normal_pipeline_clip_cfg import DepthNormalEstimationPipeline | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
pipe = DepthNormalEstimationPipeline.from_pretrained(CHECKPOINT) | |
try: | |
import xformers | |
pipe.enable_xformers_memory_efficient_attention() | |
except: | |
pass # run without xformers | |
pipe = pipe.to(device) | |
run_demo_server(pipe) | |
if __name__ == "__main__": | |
main() |