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, ): 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 run_demo_server(pipe): process_pipe = functools.partial(process, pipe) os.environ["GRADIO_ALLOW_FLAGGING"] = "never" with gr.Blocks( analytics_enabled=False, title="Geowizard Depth and Normal Estimation", css=""" #download { height: 118px; } .slider .inner { width: 5px; background: #FFF; } .viewport { aspect-ratio: 4/3; } """, ) as demo: gr.Markdown( """ttt """ ) 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( """

Depth Maps

TBD result (see Pro Tips below).

""", 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( # """ #
Pro Tips
#
    #
  1. Re-render with new parameters: Click "Clear 3D" and then "Create 3D".
    2. #
  2. Adjust 3D scale and cut-off focus: 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.
    3. #
  3. Increase details: Decrease size of the smoothing filter (also increases noise).
    4. #
# """ # ) # 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, ) 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()