-
Notifications
You must be signed in to change notification settings - Fork 60
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
72 changed files
with
5,240 additions
and
598 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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!</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=os.getenv('SYSTEM') != 'spaces', | ||
| 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 os.getenv('SYSTEM') != 'spaces': | ||
| 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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,8 @@ | ||
| source_size: 336 | ||
| source_cam_dist: 2.0 | ||
| render_size: 288 | ||
| render_views: 160 | ||
| render_fps: 40 | ||
| frame_size: 4 | ||
| mesh_size: 384 | ||
| mesh_thres: 3.0 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,7 @@ | ||
| source_size: 336 | ||
| render_size: 288 | ||
| render_views: 100 | ||
| render_fps: 25 | ||
| frame_size: 2 | ||
| mesh_size: 384 | ||
| mesh_thres: 3.0 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,8 @@ | ||
| source_size: 448 | ||
| source_cam_dist: 2.0 | ||
| render_size: 384 | ||
| render_views: 160 | ||
| render_fps: 40 | ||
| frame_size: 2 | ||
| mesh_size: 384 | ||
| mesh_thres: 3.0 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,8 @@ | ||
| source_size: 224 | ||
| source_cam_dist: 2.0 | ||
| render_size: 192 | ||
| render_views: 160 | ||
| render_fps: 40 | ||
| frame_size: 4 | ||
| mesh_size: 384 | ||
| mesh_thres: 3.0 |
Oops, something went wrong.