batching_and_3d_printing (#4)

- add fine control (7e6bc52ce03b857bb9117339e936f295b840be37)

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+files/einstein_depth_fp32.npy filter=lfs diff=lfs merge=lfs -text
+files/einstein_depth_16bit.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,2 +1,5 @@
 .idea
-.DS_Store
+.DS_Store
+__pycache__
+gradio_cached_examples
+Marigold

CONTRIBUTING.md

@@ -0,0 +1,15 @@
+## Contributing instructions
+
+We appreciate your interest in contributing. Please follow these guidelines:
+
+1. **Discuss Changes:** Start a GitHub issue to talk about your proposed change before proceeding.
+
+2. **Pull Requests:** Avoid unsolicited PRs. Discussion helps align with project goals.
+
+3. **License Agreement:** By submitting a PR, you accept our LICENSE terms.
+
+4. **Legal Compatibility:** Ensure your change complies with our project's objectives and licensing.
+
+5. **Attribution:** Credit third-party code in your PR if used.
+
+Please, feel free to reach out for questions or assistance. Your contributions are valued, and we're excited to work together to enhance this project!

LICENSE.txt

@@ -0,0 +1,177 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS

README.md

@@ -4,7 +4,7 @@ emoji: 🏵️
 colorFrom: blue
 colorTo: red
 sdk: gradio
-sdk_version: 3.44.4
+sdk_version: 4.9.1
 app_file: app.py
 pinned: true
 license: cc-by-sa-4.0

app.py

@@ -1,127 +1,564 @@
+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
 
-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 download_code():
-    os.system('git clone https://github.com/prs-eth/Marigold.git')
-
-
-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_vis=None, path_out_pred=None):
-    if path_out_vis is not None and path_out_pred is not None:
-        return path_out_vis, path_out_pred
-
-    path_input_dir = path_input + ".input"
-    path_output_dir = path_input + ".output"
-    os.makedirs(path_input_dir, exist_ok=True)
+from extrude import extrude_depth_3d
+
+
+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,
+        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)
-    shutil.copy(path_input, path_input_dir)
 
-    os.system(
-        f"cd Marigold && python3 run.py "
-        f"--input_rgb_dir \"{path_input_dir}\" "
-        f"--output_dir \"{path_output_dir}\" "
+    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],
     )
 
-    path_out_vis = find_first_png(path_output_dir + "/depth_colored")
-    assert path_out_vis is not None, "Processing failed"
-    path_out_pred = find_first_png(path_output_dir + "/depth_bw")
-    assert path_out_pred is not None, "Processing failed"
-
-    return path_out_vis, path_out_pred
-
-
-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.Image(
-            label="Predicted depth (red-near, blue-far)",
-            type="filepath",
-            visible=False,
-        ),
-        gr.Image(
-            label="Predicted depth",
-            type="filepath",
-            visible=False,
-        ),
-    ],
-    outputs=[
-        gr.Image(
-            label="Predicted depth (red-near, blue-far)",
-            type="pil",
-        ),
-        gr.Image(
-            label="Predicted depth",
-            type="pil",
-            elem_classes="imgdownload",
-        ),
-    ],
-    allow_flagging="never",
-    examples=[
-        [
-            os.path.join(os.path.dirname(__file__), "files/bee.jpg"),
-            os.path.join(os.path.dirname(__file__), "files/bee_vis.jpg"),
-            os.path.join(os.path.dirname(__file__), "files/bee_pred.jpg"),
-        ],
-        [
-            os.path.join(os.path.dirname(__file__), "files/cat.jpg"),
-            os.path.join(os.path.dirname(__file__), "files/cat_vis.jpg"),
-            os.path.join(os.path.dirname(__file__), "files/cat_pred.jpg"),
-        ],
-        [
-            os.path.join(os.path.dirname(__file__), "files/swings.jpg"),
-            os.path.join(os.path.dirname(__file__), "files/swings_vis.jpg"),
-            os.path.join(os.path.dirname(__file__), "files/swings_pred.jpg"),
-        ],
-    ],
-    css="""
-    .viewport {
-        aspect-ratio: 4/3;
-    }
-    .imgdownload {
-        height: 64px;
-    }
-    """,
-    cache_examples=True,
-)
+
+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
+                ],
+                [
+                    "files/cat.jpg",
+                    10,  # ensemble_size
+                    10,  # denoise_steps
+                    768,  # processing_res
+                    "files/cat_depth_16bit.png",
+                    "files/cat_depth_fp32.npy",
+                    "files/cat_depth_colored.png",
+                    0.0,  # plane_near
+                    0.3,  # plane_far
+                    20,  # embossing
+                    3,  # filter_size
+                    0,  # frame_near
+                ],
+                [
+                    "files/swings.jpg",
+                    10,  # ensemble_size
+                    10,  # denoise_steps
+                    768,  # processing_res
+                    "files/swings_depth_16bit.png",
+                    "files/swings_depth_fp32.npy",
+                    "files/swings_depth_colored.png",
+                    0.05,  # plane_near
+                    0.25,  # plane_far
+                    10,  # embossing
+                    1,  # filter_size
+                    0,  # frame_near
+                ],
+                [
+                    "files/einstein.jpg",
+                    10,  # ensemble_size
+                    10,  # denoise_steps
+                    768,  # processing_res
+                    "files/einstein_depth_16bit.png",
+                    "files/einstein_depth_fp32.npy",
+                    "files/einstein_depth_colored.png",
+                    0.0,  # plane_near
+                    0.5,  # plane_far
+                    50,  # embossing
+                    3,  # filter_size
+                    -15,  # 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().launch(server_name="0.0.0.0", server_port=7860)
+
+
+def prefetch_hf_cache(pipe):
+    process(pipe, "files/bee.jpg", 1, 1, 64)
+    shutil.rmtree("files/bee_output")
+
+
+def main():
+    REPO_URL = "https://github.com/prs-eth/Marigold.git"
+    REPO_HASH = "22437a9d"
+    REPO_DIR = "Marigold"
+    CHECKPOINT = "Bingxin/Marigold"
+
+    if os.path.isdir(REPO_DIR):
+        shutil.rmtree(REPO_DIR)
+    repo = git.Repo.clone_from(REPO_URL, REPO_DIR)
+    repo.git.checkout(REPO_HASH)
+
+    sys.path.append(os.path.join(os.getcwd(), REPO_DIR))
+
+    from marigold import MarigoldPipeline
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    pipe = MarigoldPipeline.from_pretrained(CHECKPOINT)
+    try:
+        import xformers
+        pipe.enable_xformers_memory_efficient_attention()
+    except:
+        pass  # run without xformers
+
+    pipe = pipe.to(device)
+    prefetch_hf_cache(pipe)
+    run_demo_server(pipe)
 
 
 if __name__ == "__main__":
-    download_code()
-    iface.queue().launch(server_name="0.0.0.0", server_port=7860)
+    main()

extrude.py

@@ -0,0 +1,322 @@
+import math
+import os
+
+import numpy as np
+import pygltflib
+import trimesh
+from PIL import Image, ImageFilter
+
+
+def quaternion_multiply(q1, q2):
+    x1, y1, z1, w1 = q1
+    x2, y2, z2, w2 = q2
+    return [
+        w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
+        w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
+        w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
+        w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
+    ]
+
+
+def glb_add_lights(path_input, path_output):
+    """
+    Adds directional lights in the horizontal plane to the glb file.
+    :param path_input: path to input glb
+    :param path_output: path to output glb
+    :return: None
+    """
+    glb = pygltflib.GLTF2().load(path_input)
+
+    N = 3  # default max num lights in Babylon.js is 4
+    angle_step = 2 * math.pi / N
+    elevation_angle = math.radians(75)
+
+    light_colors = [
+        [1.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0],
+        [0.0, 0.0, 1.0],
+    ]
+
+    lights_extension = {
+        "lights": [
+            {"type": "directional", "color": light_colors[i], "intensity": 2.0}
+            for i in range(N)
+        ]
+    }
+
+    if "KHR_lights_punctual" not in glb.extensionsUsed:
+        glb.extensionsUsed.append("KHR_lights_punctual")
+    glb.extensions["KHR_lights_punctual"] = lights_extension
+
+    light_nodes = []
+    for i in range(N):
+        angle = i * angle_step
+
+        pos_rot = [0.0, 0.0, math.sin(angle / 2), math.cos(angle / 2)]
+        elev_rot = [math.sin(elevation_angle / 2), 0.0, 0.0, math.cos(elevation_angle / 2)]
+        rotation = quaternion_multiply(pos_rot, elev_rot)
+
+        node = {
+            "rotation": rotation,
+            "extensions": {"KHR_lights_punctual": {"light": i}},
+        }
+        light_nodes.append(node)
+
+    light_node_indices = list(range(len(glb.nodes), len(glb.nodes) + N))
+    glb.nodes.extend(light_nodes)
+
+    root_node_index = glb.scenes[glb.scene].nodes[0]
+    root_node = glb.nodes[root_node_index]
+    if hasattr(root_node, "children"):
+        root_node.children.extend(light_node_indices)
+    else:
+        root_node.children = light_node_indices
+
+    glb.save(path_output)
+
+
+def extrude_depth_3d(
+    path_rgb,
+    path_depth,
+    output_model_scale=100,
+    filter_size=3,
+    coef_near=0.0,
+    coef_far=1.0,
+    emboss=0.3,
+    f_thic=0.05,
+    f_near=-0.15,
+    f_back=0.01,
+    vertex_colors=True,
+    scene_lights=True,
+):
+    f_far_inner = -emboss
+    f_far_outer = f_far_inner - f_back
+
+    f_near = max(f_near, f_far_inner)
+
+    depth_image = Image.open(path_depth)
+    assert depth_image.mode == "I", depth_image.mode
+    depth_image = depth_image.filter(ImageFilter.MedianFilter(size=filter_size))
+
+    w, h = depth_image.size
+    d_max = max(w, h)
+    depth_image = np.array(depth_image).astype(np.double)
+    z_min, z_max = np.min(depth_image), np.max(depth_image)
+    depth_image = (depth_image.astype(np.double) - z_min) / (z_max - z_min)
+    depth_image[depth_image < coef_near] = coef_near
+    depth_image[depth_image > coef_far] = coef_far
+    depth_image = emboss * (depth_image - coef_near) / (coef_far - coef_near)
+    rgb_image = np.array(
+        Image.open(path_rgb).convert("RGB").resize((w, h), Image.Resampling.LANCZOS)
+    )
+
+    w_norm = w / float(d_max - 1)
+    h_norm = h / float(d_max - 1)
+    w_half = w_norm / 2
+    h_half = h_norm / 2
+
+    x, y = np.meshgrid(np.arange(w), np.arange(h))
+    x = x / float(d_max - 1) - w_half  # [-w_half, w_half]
+    y = -y / float(d_max - 1) + h_half  # [-h_half, h_half]
+    z = -depth_image  # -depth_emboss (far) - 0 (near)
+    vertices_2d = np.stack((x, y, z), axis=-1)
+    vertices = vertices_2d.reshape(-1, 3)
+    colors = rgb_image[:, :, :3].reshape(-1, 3) / 255.0
+
+    faces = []
+    for y in range(h - 1):
+        for x in range(w - 1):
+            idx = y * w + x
+            faces.append([idx, idx + w, idx + 1])
+            faces.append([idx + 1, idx + w, idx + 1 + w])
+
+    # OUTER frame
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [-w_half - f_thic, -h_half - f_thic, f_near],  # 00
+            [-w_half - f_thic, -h_half - f_thic, f_far_outer],  # 01
+            [w_half + f_thic, -h_half - f_thic, f_near],  # 02
+            [w_half + f_thic, -h_half - f_thic, f_far_outer],  # 03
+            [w_half + f_thic, h_half + f_thic, f_near],  # 04
+            [w_half + f_thic, h_half + f_thic, f_far_outer],  # 05
+            [-w_half - f_thic, h_half + f_thic, f_near],  # 06
+            [-w_half - f_thic, h_half + f_thic, f_far_outer],  # 07
+        ],
+        axis=0,
+    )
+    faces.extend(
+        [
+            [nv + 0, nv + 1, nv + 2],
+            [nv + 2, nv + 1, nv + 3],
+            [nv + 2, nv + 3, nv + 4],
+            [nv + 4, nv + 3, nv + 5],
+            [nv + 4, nv + 5, nv + 6],
+            [nv + 6, nv + 5, nv + 7],
+            [nv + 6, nv + 7, nv + 0],
+            [nv + 0, nv + 7, nv + 1],
+        ]
+    )
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * 8, axis=0)
+
+    # INNER frame
+
+    nv = len(vertices)
+    vertices_left_data = vertices_2d[:, 0]  # H x 3
+    vertices_left_frame = vertices_2d[:, 0].copy()  # H x 3
+    vertices_left_frame[:, 2] = f_near
+    vertices = np.append(vertices, vertices_left_data, axis=0)
+    vertices = np.append(vertices, vertices_left_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * h), axis=0)
+    for i in range(h - 1):
+        nvi_d = nv + i
+        nvi_f = nvi_d + h
+        faces.append([nvi_d, nvi_f, nvi_d + 1])
+        faces.append([nvi_d + 1, nvi_f, nvi_f + 1])
+
+    nv = len(vertices)
+    vertices_right_data = vertices_2d[:, -1]  # H x 3
+    vertices_right_frame = vertices_2d[:, -1].copy()  # H x 3
+    vertices_right_frame[:, 2] = f_near
+    vertices = np.append(vertices, vertices_right_data, axis=0)
+    vertices = np.append(vertices, vertices_right_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * h), axis=0)
+    for i in range(h - 1):
+        nvi_d = nv + i
+        nvi_f = nvi_d + h
+        faces.append([nvi_d, nvi_d + 1, nvi_f])
+        faces.append([nvi_d + 1, nvi_f + 1, nvi_f])
+
+    nv = len(vertices)
+    vertices_top_data = vertices_2d[0, :]  # H x 3
+    vertices_top_frame = vertices_2d[0, :].copy()  # H x 3
+    vertices_top_frame[:, 2] = f_near
+    vertices = np.append(vertices, vertices_top_data, axis=0)
+    vertices = np.append(vertices, vertices_top_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * w), axis=0)
+    for i in range(w - 1):
+        nvi_d = nv + i
+        nvi_f = nvi_d + w
+        faces.append([nvi_d, nvi_d + 1, nvi_f])
+        faces.append([nvi_d + 1, nvi_f + 1, nvi_f])
+
+    nv = len(vertices)
+    vertices_bottom_data = vertices_2d[-1, :]  # H x 3
+    vertices_bottom_frame = vertices_2d[-1, :].copy()  # H x 3
+    vertices_bottom_frame[:, 2] = f_near
+    vertices = np.append(vertices, vertices_bottom_data, axis=0)
+    vertices = np.append(vertices, vertices_bottom_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * w), axis=0)
+    for i in range(w - 1):
+        nvi_d = nv + i
+        nvi_f = nvi_d + w
+        faces.append([nvi_d, nvi_f, nvi_d + 1])
+        faces.append([nvi_d + 1, nvi_f, nvi_f + 1])
+
+    # FRONT frame
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [-w_half - f_thic, -h_half - f_thic, f_near],
+            [-w_half - f_thic, h_half + f_thic, f_near],
+        ],
+        axis=0,
+    )
+    vertices = np.append(vertices, vertices_left_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + h), axis=0)
+    for i in range(h - 1):
+        faces.append([nv, nv + 2 + i + 1, nv + 2 + i])
+    faces.append([nv, nv + 2, nv + 1])
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [w_half + f_thic, h_half + f_thic, f_near],
+            [w_half + f_thic, -h_half - f_thic, f_near],
+        ],
+        axis=0,
+    )
+    vertices = np.append(vertices, vertices_right_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + h), axis=0)
+    for i in range(h - 1):
+        faces.append([nv, nv + 2 + i, nv + 2 + i + 1])
+    faces.append([nv, nv + h + 1, nv + 1])
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [w_half + f_thic, h_half + f_thic, f_near],
+            [-w_half - f_thic, h_half + f_thic, f_near],
+        ],
+        axis=0,
+    )
+    vertices = np.append(vertices, vertices_top_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + w), axis=0)
+    for i in range(w - 1):
+        faces.append([nv, nv + 2 + i, nv + 2 + i + 1])
+    faces.append([nv, nv + 1, nv + 2])
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [-w_half - f_thic, -h_half - f_thic, f_near],
+            [w_half + f_thic, -h_half - f_thic, f_near],
+        ],
+        axis=0,
+    )
+    vertices = np.append(vertices, vertices_bottom_frame, axis=0)
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + w), axis=0)
+    for i in range(w - 1):
+        faces.append([nv, nv + 2 + i + 1, nv + 2 + i])
+    faces.append([nv, nv + 1, nv + w + 1])
+
+    # BACK frame
+
+    nv = len(vertices)
+    vertices = np.append(
+        vertices,
+        [
+            [-w_half - f_thic, -h_half - f_thic, f_far_outer],  # 00
+            [w_half + f_thic, -h_half - f_thic, f_far_outer],  # 01
+            [w_half + f_thic, h_half + f_thic, f_far_outer],  # 02
+            [-w_half - f_thic, h_half + f_thic, f_far_outer],  # 03
+        ],
+        axis=0,
+    )
+    faces.extend(
+        [
+            [nv + 0, nv + 2, nv + 1],
+            [nv + 2, nv + 0, nv + 3],
+        ]
+    )
+    colors = np.append(colors, [[0.5, 0.5, 0.5]] * 4, axis=0)
+
+    trimesh_kwargs = {}
+    if vertex_colors:
+        trimesh_kwargs["vertex_colors"] = colors
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, **trimesh_kwargs)
+
+    mesh.merge_vertices()
+
+    current_max_dimension = max(mesh.extents)
+    scaling_factor = output_model_scale / current_max_dimension
+    mesh.apply_scale(scaling_factor)
+
+    path_out_base = os.path.splitext(path_depth)[0].replace("_16bit", "")
+    path_out_glb = path_out_base + ".glb"
+    path_out_stl = path_out_base + ".stl"
+
+    mesh.export(path_out_glb, file_type="glb")
+    if scene_lights:
+        glb_add_lights(path_out_glb, path_out_glb)
+
+    mesh.export(path_out_stl, file_type="stl")
+
+    return path_out_glb, path_out_stl

files/bee_depth_fp32.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82bfba1b08942eea990bcda9f120bc75274f432de56d221fadbf428f85c6f22c
+size 1398912

files/cat_depth_fp32.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b44fba3ec47bc4591d47ee2ddf8594934a89728fb8fc8d50bc9596dfe2cee831
+size 1550292

files/einstein_depth_fp32.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6306511f81c545ef9997176e067ed44d0ee6c4277788f19ffc109586be88a2f7
+size 4194432

files/swings_depth_fp32.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5728d846cd554d4a5e1d0e5f71d622135bca36164026b8e49668acdfa20e070
+size 1398912

requirements.txt

@@ -1,24 +1,13 @@
-gradio==3.44.4
-gradio_client==0.5.1
-trimesh==3.23.5
+gradio==4.9.1
+gradio-imageslider==0.0.16
+GitPython==3.1.40
+pygltflib==1.16.1
+trimesh==4.0.5
 
-accelerate
-diffusers==0.20.1
-h5py
-matplotlib
-numpy==1.26.1
-omegaconf
-opencv-python
-pandas
-scipy==1.11.3
-tabulate
-tensorboard
+accelerate>=0.22.0
+diffusers>=0.20.1
+matplotlib==3.8.2
+scipy==1.11.4
 torch==2.0.1
-torchaudio
-torchvision
-torchshow
-tqdm
-transformers
-triton
-wandb==0.14.0
-xformers
+transformers>=4.32.1
+xformers==0.0.21