Spaces:
Sleeping
Sleeping
| # 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() | |