Upload app.py

app.py

@@ -0,0 +1,210 @@
+# Copyright (c) 2023-2024, Zexin He
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+from PIL import Image
+import numpy as np
+import gradio as gr
+
+
+def assert_input_image(input_image):
+    if input_image is None:
+        raise gr.Error("No image selected or uploaded!")
+
+def prepare_working_dir():
+    import tempfile
+    working_dir = tempfile.TemporaryDirectory()
+    return working_dir
+
+def init_preprocessor():
+    from openlrm.utils.preprocess import Preprocessor
+    global preprocessor
+    preprocessor = Preprocessor()
+
+def preprocess_fn(image_in: np.ndarray, remove_bg: bool, recenter: bool, working_dir):
+    image_raw = os.path.join(working_dir.name, "raw.png")
+    with Image.fromarray(image_in) as img:
+        img.save(image_raw)
+    image_out = os.path.join(working_dir.name, "rembg.png")
+    success = preprocessor.preprocess(image_path=image_raw, save_path=image_out, rmbg=remove_bg, recenter=recenter)
+    assert success, f"Failed under preprocess_fn!"
+    return image_out
+
+
+def demo_openlrm(infer_impl):
+
+    def core_fn(image: str, source_cam_dist: float, working_dir):
+        dump_video_path = os.path.join(working_dir.name, "output.mp4")
+        dump_mesh_path = os.path.join(working_dir.name, "output.ply")
+        infer_impl(
+            image_path=image,
+            source_cam_dist=source_cam_dist,
+            export_video=True,
+            export_mesh=False,
+            dump_video_path=dump_video_path,
+            dump_mesh_path=dump_mesh_path,
+        )
+        return dump_video_path
+
+    def example_fn(image: np.ndarray):
+        from gradio.utils import get_cache_folder
+        working_dir = get_cache_folder()
+        image = preprocess_fn(
+            image_in=image,
+            remove_bg=True,
+            recenter=True,
+            working_dir=working_dir,
+        )
+        video = core_fn(
+            image=image,
+            source_cam_dist=2.0,
+            working_dir=working_dir,
+        )
+        return image, video
+
+
+    _TITLE = '''OpenLRM: Open-Source Large Reconstruction Models'''
+
+    _DESCRIPTION = '''
+        <div>
+            <a style="display:inline-block" href='https://github.com/3DTopia/OpenLRM'><img src='https://img.shields.io/github/stars/3DTopia/OpenLRM?style=social'/></a>
+            <a style="display:inline-block; margin-left: .5em" href="https://huggingface.co/zxhezexin"><img src='https://img.shields.io/badge/Model-Weights-blue'/></a>
+        </div>
+        OpenLRM is an open-source implementation of Large Reconstruction Models.
+
+        <strong>Image-to-3D in 10 seconds with A100!</strong>
+
+        <strong>Disclaimer:</strong> This demo uses `openlrm-mix-base-1.1` model with 288x288 rendering resolution here for a quick demonstration.
+    '''
+
+    with gr.Blocks(analytics_enabled=False) as demo:
+
+        # HEADERS
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown('# ' + _TITLE)
+        with gr.Row():
+            gr.Markdown(_DESCRIPTION)
+
+        # DISPLAY
+        with gr.Row():
+
+            with gr.Column(variant='panel', scale=1):
+                with gr.Tabs(elem_id="openlrm_input_image"):
+                    with gr.TabItem('Input Image'):
+                        with gr.Row():
+                            input_image = gr.Image(label="Input Image", image_mode="RGBA", width="auto", sources="upload", type="numpy", elem_id="content_image")
+
+            with gr.Column(variant='panel', scale=1):
+                with gr.Tabs(elem_id="openlrm_processed_image"):
+                    with gr.TabItem('Processed Image'):
+                        with gr.Row():
+                            processed_image = gr.Image(label="Processed Image", image_mode="RGBA", type="filepath", elem_id="processed_image", width="auto", interactive=False)
+
+            with gr.Column(variant='panel', scale=1):
+                with gr.Tabs(elem_id="openlrm_render_video"):
+                    with gr.TabItem('Rendered Video'):
+                        with gr.Row():
+                            output_video = gr.Video(label="Rendered Video", format="mp4", width="auto", autoplay=True)
+
+        # SETTING
+        with gr.Row():
+            with gr.Column(variant='panel', scale=1):
+                with gr.Tabs(elem_id="openlrm_attrs"):
+                    with gr.TabItem('Settings'):
+                        with gr.Column(variant='panel'):
+                            gr.Markdown(
+                                """
+                                <strong>Best Practice</strong>:
+                                    Centered objects in reasonable sizes. Try adjusting source camera distances.
+                                """
+                            )
+                            checkbox_rembg = gr.Checkbox(True, label='Remove background')
+                            checkbox_recenter = gr.Checkbox(True, label='Recenter the object')
+                            slider_cam_dist = gr.Slider(1.0, 3.5, value=2.0, step=0.1, label="Source Camera Distance")
+                            submit = gr.Button('Generate', elem_id="openlrm_generate", variant='primary')
+
+        # EXAMPLES
+        with gr.Row():
+            examples = [
+                ['assets/sample_input/owl.png'],
+                ['assets/sample_input/building.png'],
+                ['assets/sample_input/mailbox.png'],
+                ['assets/sample_input/fire.png'],
+                ['assets/sample_input/girl.png'],
+                ['assets/sample_input/lamp.png'],
+                ['assets/sample_input/hydrant.png'],
+                ['assets/sample_input/hotdogs.png'],
+                ['assets/sample_input/traffic.png'],
+                ['assets/sample_input/ceramic.png'],
+            ]
+            gr.Examples(
+                examples=examples,
+                inputs=[input_image], 
+                outputs=[processed_image, output_video],
+                fn=example_fn,
+                cache_examples=bool(os.getenv('SPACE_ID')),
+                examples_per_page=20,
+            )
+
+        working_dir = gr.State()
+        submit.click(
+            fn=assert_input_image,
+            inputs=[input_image],
+            queue=False,
+        ).success(
+            fn=prepare_working_dir,
+            outputs=[working_dir],
+            queue=False,
+        ).success(
+            fn=preprocess_fn,
+            inputs=[input_image, checkbox_rembg, checkbox_recenter, working_dir],
+            outputs=[processed_image],
+        ).success(
+            fn=core_fn,
+            inputs=[processed_image, slider_cam_dist, working_dir],
+            outputs=[output_video],
+        )
+
+        demo.queue()
+        demo.launch()
+
+
+def launch_gradio_app():
+
+    os.environ.update({
+        "APP_ENABLED": "1",
+        "APP_MODEL_NAME": "zxhezexin/openlrm-mix-base-1.1",
+        "APP_INFER": "./configs/infer-gradio.yaml",
+        "APP_TYPE": "infer.lrm",
+        "NUMBA_THREADING_LAYER": 'omp',
+    })
+
+    from openlrm.runners import REGISTRY_RUNNERS
+    from openlrm.runners.infer.base_inferrer import Inferrer
+    InferrerClass : Inferrer = REGISTRY_RUNNERS[os.getenv("APP_TYPE")]
+    with InferrerClass() as inferrer:
+        init_preprocessor()
+        if not bool(os.getenv('SPACE_ID')):
+            from openlrm.utils.proxy import no_proxy
+            demo = no_proxy(demo_openlrm)
+        else:
+            demo = demo_openlrm
+        demo(infer_impl=inferrer.infer_single)
+
+
+if __name__ == '__main__':
+
+    launch_gradio_app()