app.py

import gradio as gr
import numpy as np
import random
import logging
import os
from glob import glob

import numpy as np
import torch
from PIL import Image
from tqdm.auto import tqdm

from depthmaster import DepthMasterPipeline
from depthmaster.modules.unet_2d_condition import UNet2DConditionModel

def load_example(example_image):
    # 返回选中的图片
    return example_image


device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "zysong212/DepthMaster"  # Replace to the model you would like to use

# if torch.cuda.is_available():
#     torch_dtype = torch.float16
# else:
torch_dtype = torch.float32

# pipe = DepthMasterPipeline.from_pretrained('eval', torch_dtype=torch_dtype)
# unet = UNet2DConditionModel.from_pretrained(os.path.join('eval', f'unet'))
pipe = DepthMasterPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
unet = UNet2DConditionModel.from_pretrained(model_repo_id, subfolder="unet", torch_dtype=torch_dtype)
pipe.unet = unet

try:
    pipe.enable_xformers_memory_efficient_attention()
except ImportError:
    pass  # run without xformers

pipe = pipe.to(device)

# MAX_SEED = np.iinfo(np.int32).max
# MAX_IMAGE_SIZE = 1024


# @spaces.GPU #[uncomment to use ZeroGPU]
def infer(
    input_image,
    progress=gr.Progress(track_tqdm=True),
):
    # if randomize_seed:
    #     seed = random.randint(0, MAX_SEED)

    # generator = torch.Generator().manual_seed(seed)

    # image = pipe(
    #     prompt=prompt,
    #     negative_prompt=negative_prompt,
    #     guidance_scale=guidance_scale,
    #     num_inference_steps=num_inference_steps,
    #     width=width,
    #     height=height,
    #     generator=generator,
    # ).images[0]
    pipe_out = pipe(
            input_image,
            processing_res=768,
            match_input_res=True,
            batch_size=1,
            color_map="Spectral",
            show_progress_bar=True,
            resample_method="bilinear",
        )

    # depth_pred: np.ndarray = pipe_out.depth_np
    depth_colored: Image.Image = pipe_out.depth_colored


    return depth_colored


# 默认图像路径
example_images = [
    "wild_example/000000000776.jpg",
    "wild_example/800x.jpg",
    "wild_example/000000055950.jpg",
    "wild_example/53441037037_c2cbd91ad2_k.jpg",
    "wild_example/53501906161_6109e3da29_b.jpg",
    "wild_example/m_1e31af1c.jpg",
    "wild_example/sg-11134201-7rd5x-lvlh48byidbqca.jpg"
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
#example-gallery {
    height: 80px; /* 设置缩略图高度 */
    width: auto;  /* 保持宽高比 */
    margin: 0 auto;  /* 图片间距 */
    cursor: pointer; /* 鼠标指针变为手型 */
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("# DepthMaster")
    gr.Markdown("Official demo for DepthMaster. Please refer to our [paper](https://arxiv.org/abs/2501.02576), [project page](https://indu1ge.github.io/DepthMaster_page/), and [github](https://github.com/indu1ge/DepthMaster) for more details.")
    gr.Markdown(" ### Depth Estimation with DepthMaster.")
    # with gr.Column(elem_id="col-container"):
    #     gr.Markdown(" # Depth Estimation")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil", elem_id="input-image", interactive=True)
        with gr.Column():
            depth_map = gr.Image(label="Depth Map with Slider View", type="pil", interactive=False, elem_id="depth-map")
        
    # 计算按钮
    compute_button = gr.Button("Compute Depth")

    # # 添加示例图片选择器
    # with gr.Row():
    #     gr.Markdown("### example images")
    # with gr.Row(elem_id="example-gallery"):
    #     example_gallery = gr.Gallery(
    #         label="",
    #         value=example_images,
    #         elem_id="example-gallery",
    #         show_label=False,
    #         interactive=True,
    #         columns=10
    #     )

    # 设置默认图片点击后的操作
    # example_gallery.select(
    #     fn=lambda img_path: img_path,  # 回调函数：返回选择的路径
    #     inputs=[],
    #     outputs=input_image  # 输出设置为 Input Image
    # )
    # example_gallery.click(
    #     fn=load_example,  # 选择图片的回调
    #     inputs=[example_gallery],  # 输入：用户点击的图片
    #     outputs=[input_image]  # 输出：更新 Input Image
    # )


    # 设置计算按钮的回调
    compute_button.click(
        fn=infer,  # 回调函数
        inputs=input_image,  # 输入
        outputs=depth_map  # 输出
    )

# 启动 Gradio 应用
demo.launch()
        # with gr.Column(scale=45):
        #     img_in = gr.Image(type="pil")
        # with gr.Column(scale=45):
        #     img_out = 

        # with gr.Row():
        #     prompt = gr.Text(
        #         label="Prompt",
        #         show_label=False,
        #         max_lines=1,
        #         placeholder="Enter your prompt",
        #         container=False,
        #     )

        #     run_button = gr.Button("Run", scale=0, variant="primary")

        # result = gr.Image(label="Result", show_label=False)

        # with gr.Accordion("Advanced Settings", open=False):
        #     negative_prompt = gr.Text(
        #         label="Negative prompt",
        #         max_lines=1,
        #         placeholder="Enter a negative prompt",
        #         visible=False,
        #     )

        #     seed = gr.Slider(
        #         label="Seed",
        #         minimum=0,
        #         maximum=MAX_SEED,
        #         step=1,
        #         value=0,
        #     )

        #     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            # with gr.Row():
            #     width = gr.Slider(
            #         label="Width",
            #         minimum=256,
            #         maximum=MAX_IMAGE_SIZE,
            #         step=32,
            #         value=1024,  # Replace with defaults that work for your model
            #     )

            #     height = gr.Slider(
            #         label="Height",
            #         minimum=256,
            #         maximum=MAX_IMAGE_SIZE,
            #         step=32,
            #         value=1024,  # Replace with defaults that work for your model
            #     )

            # with gr.Row():
            #     guidance_scale = gr.Slider(
            #         label="Guidance scale",
            #         minimum=0.0,
            #         maximum=10.0,
            #         step=0.1,
            #         value=0.0,  # Replace with defaults that work for your model
            #     )

            #     num_inference_steps = gr.Slider(
            #         label="Number of inference steps",
            #         minimum=1,
            #         maximum=50,
            #         step=1,
            #         value=2,  # Replace with defaults that work for your model
            #     )

    #     gr.Examples(examples=examples, inputs=[prompt])
    # gr.on(
    #     triggers=[run_button.click, prompt.submit],
    #     fn=infer,
    #     inputs=[
    #         prompt,
    #         negative_prompt,
    #         seed,
    #         randomize_seed,
    #         # width,
    #         # height,
    #         # guidance_scale,
    #         # num_inference_steps,
    #     ],
    #     outputs=[result, seed],
    # )

# if __name__ == "__main__":
#     demo.launch()