initial commit

.gitignore

@@ -0,0 +1,2 @@
+.idea
+.DS_Store

README.md

@@ -1,13 +1,24 @@
 ---
-title: Marigold
-emoji: 🦀
-colorFrom: gray
-colorTo: yellow
+title: Marigold Depth Estimation
+emoji: 🏵️
+colorFrom: blue
+colorTo: red
 sdk: gradio
-sdk_version: 4.8.0
+sdk_version: 3.44.4
 app_file: app.py
-pinned: false
+pinned: true
 license: cc-by-sa-4.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+This is a demo of the monocular depth estimation pipeline, described in the paper titled ["Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation"](https://arxiv.org/abs/2312.02145)
+
+```
+@misc{ke2023repurposing,
+      title={Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation}, 
+      author={Bingxin Ke and Anton Obukhov and Shengyu Huang and Nando Metzger and Rodrigo Caye Daudt and Konrad Schindler},
+      year={2023},
+      eprint={2312.02145},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

app.py

@@ -0,0 +1,139 @@
+import os
+import shutil
+
+import gradio as gr
+
+
+desc = """
+    <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 pane on the left side, or expore examples listed in the bottom.  
+    </p>
+"""
+
+
+def init_persistence(purge=False):
+    if not os.path.exists('/data'):
+        return
+    os.environ['ckpt_dir'] = "/data/Marigold_ckpt"
+    os.environ['TRANSFORMERS_CACHE'] = "/data/hfcache"
+    os.environ['HF_DATASETS_CACHE'] = "/data/hfcache"
+    os.environ['HF_HOME'] = "/data/hfcache"
+    if purge:
+        os.system("rm -rf /data/Marigold_ckpt/*")
+
+
+def download_code_weights():
+    os.system('git clone https://github.com/prs-eth/Marigold.git')
+    os.system('cd Marigold && bash script/download_weights.sh')
+    os.system('echo /data && ls -la /data')
+    os.system('echo /data/Marigold_ckpt && ls -la /data/Marigold_ckpt')
+    os.system('echo /data/Marigold_ckpt/Marigold_v1_merged && ls -la /data/Marigold_ckpt/Marigold_v1_merged')
+
+
+def find_first_png(directory):
+    for file in os.listdir(directory):
+        if file.lower().endswith(".png"):
+            return os.path.join(directory, file)
+    return None
+
+
+def marigold_process(path_input, path_out_png=None, path_out_obj=None, path_out_2_png=None):
+    if path_out_png is not None and path_out_obj is not None and path_out_2_png is not None:
+        return path_out_png, path_out_obj, path_out_2_png
+
+    path_input_dir = path_input + ".input"
+    path_output_dir = path_input + ".output"
+    os.makedirs(path_input_dir, exist_ok=True)
+    os.makedirs(path_output_dir, exist_ok=True)
+    shutil.copy(path_input, path_input_dir)
+
+    persistence_args = ""
+    if os.path.exists('/data'):
+        persistence_args = "--checkpoint /data/Marigold_ckpt/Marigold_v1_merged"
+
+    os.system(
+        f"cd Marigold && python3 run.py "
+        f"{persistence_args} "
+        f"--input_rgb_dir \"{path_input_dir}\" "
+        f"--output_dir \"{path_output_dir}\" "
+        f"--n_infer 5 "
+        f"--denoise_steps 10 "
+    )
+
+    # depth_colored, depth_bw, depth_npy
+    path_out_colored = find_first_png(path_output_dir + "/depth_colored")
+    assert path_out_colored is not None, "Processing failed"
+    path_out_bw = find_first_png(path_output_dir + "/depth_bw")
+    assert path_out_bw is not None, "Processing failed"
+
+    return path_out_colored, path_out_bw
+
+
+iface = gr.Interface(
+    title="Marigold Depth Estimation",
+    description=desc,
+    thumbnail="marigold_logo_square.jpg",
+    fn=marigold_process,
+    inputs=[
+        gr.Image(
+            label="Input Image",
+            type="filepath",
+        ),
+        gr.File(
+            label="Predicted depth (red-near, blue-far)",
+            visible=False,
+        ),
+        gr.File(
+            label="Predicted depth (16-bit PNG)",
+            visible=False,
+        ),
+    ],
+    outputs=[
+        gr.Image(
+            label="Predicted depth (red-near, blue-far)",
+            type="pil",
+        ),
+        gr.Image(
+            label="Predicted depth (16-bit PNG)",
+            type="pil",
+            elem_classes="imgdownload",
+        ),
+    ],
+    allow_flagging="never",
+    # examples=[
+    #     [
+    #         os.path.join(os.path.dirname(__file__), "files/test.png"),
+    #         os.path.join(os.path.dirname(__file__), "files/test.png.out.png"),
+    #         os.path.join(os.path.dirname(__file__), "files/test.png.out.2.png"),
+    #     ],
+    # ],
+    css="""
+    .viewport {
+        aspect-ratio: 4/3;
+    }
+    .imgdownload {
+        height: 32px;
+    }
+    """,
+    cache_examples=True,
+)
+
+
+if __name__ == "__main__":
+    init_persistence()
+    download_code_weights()
+    iface.queue().launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -0,0 +1,24 @@
+gradio==3.44.4
+gradio_client==0.5.1
+trimesh==3.23.5
+
+accelerate
+diffusers==0.20.1
+h5py
+matplotlib
+numpy==1.26.1
+omegaconf
+opencv-python
+pandas
+scipy==1.11.3
+tabulate
+tensorboard
+torch==2.0.1
+torchaudio
+torchvision
+torchshow
+tqdm
+transformers
+triton
+wandb==0.14.0
+xformers