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_front_image, input_back_image): | |
| if input_front_image is None: | |
| raise gr.Error("No front image selected or uploaded!") | |
| if input_back_image is None: | |
| raise gr.Error("No back 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_front: np.ndarray, image_in_back: np.ndarray, remove_bg: bool, recenter: bool, working_dir): | |
| # save front image first | |
| image_raw_front = os.path.join(working_dir.name, "raw_front.png") | |
| with Image.fromarray(image_in_front) as img: | |
| img.save(image_raw_front) | |
| image_out_front = os.path.join(working_dir.name, "front/rembg_front.png") | |
| # save back image first | |
| image_raw_back = os.path.join(working_dir.name, "raw_back.png") | |
| with Image.fromarray(image_in_back) as img: | |
| img.save(image_raw_back) | |
| image_out_back = os.path.join(working_dir.name, "back/rembg_back.png") | |
| # process the front and back image. | |
| success_front = preprocessor.preprocess(image_path=image_raw_front, save_path=image_out_front, rmbg=remove_bg, recenter=recenter) | |
| success_back = preprocessor.preprocess(image_path=image_raw_back, save_path=image_out_back, rmbg=remove_bg, recenter=recenter) | |
| assert success_front and success_back, f"Failed under preprocess_fn!" | |
| return image_out_front, image_out_back | |
| def demo_openlrm(infer_impl): | |
| def core_fn(image_front: str, image_back: 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_front, | |
| source_cam_dist=source_cam_dist, | |
| export_video=True, | |
| export_mesh=False, | |
| dump_video_path=dump_video_path, | |
| dump_mesh_path=dump_mesh_path, | |
| image_path_back=image_back, | |
| ) | |
| return dump_video_path | |
| def example_fn(input_front_image: np.ndarray, input_back_image: np.ndarray): | |
| from gradio.utils import get_cache_folder | |
| working_dir = get_cache_folder() | |
| processed_front_image, processed_back_image = preprocess_fn( | |
| image_in_front=input_front_image, | |
| image_in_back=input_back_image, | |
| remove_bg=True, | |
| recenter=True, | |
| working_dir=working_dir, | |
| ) | |
| video = core_fn( | |
| image_front=processed_front_image, | |
| image_back=processed_back_image, | |
| source_cam_dist=2.0, | |
| working_dir=working_dir, | |
| ) | |
| return processed_front_image, processed_back_image, video | |
| _TITLE = '''🔥 🔥 🔥 Tailor3D: Customized 3D Assets Editing and Generation with Dual-Side Images''' | |
| _DESCRIPTION = ''' | |
| <div> | |
| <a style="display:inline-block" href='https://github.com/Qi-Zhangyang/Tailor3D'><img src='https://img.shields.io/github/stars/Qi-Zhangyang/Tailor3D?style=social'/></a> | |
| <a style="display:inline-block; margin-left: .5em" href="https://huggingface.co./alexzyqi"><img src='https://img.shields.io/badge/Model-Weights-blue'/></a> | |
| </div> | |
| We propose Tailor3D, a novel pipeline creating customized 3D assets from editable dual-side images and feed-forward reconstruction methods. | |
| Here we show the final step of Tailor3D. That is given the edited front and beck view of the object. We can produce the 3D object with several seconds. | |
| <strong>Disclaimer:</strong> This demo uses `Tailor3D-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(): | |
| gr.Markdown( | |
| """ | |
| ## 🖼️ Input: This is the input front and back images. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(variant='panel', scale=0.2): | |
| with gr.Tabs(elem_id="tailor3d_input_front_image"): | |
| with gr.TabItem('Input Front-view Image'): | |
| with gr.Row(): | |
| input_front_image = gr.Image(label="Input Front Image", image_mode="RGBA", width="auto", sources="upload", type="numpy", elem_id="content_image") | |
| with gr.Column(variant='panel', scale=0.2): | |
| with gr.Tabs(elem_id="tailor3d_input_back_image"): | |
| with gr.TabItem('Input Back-view Image'): | |
| with gr.Row(): | |
| input_back_image = gr.Image(label="Input Back Image", image_mode="RGBA", width="auto", sources="upload", type="numpy", elem_id="content_image") | |
| with gr.Row(): | |
| gr.Markdown( | |
| """ | |
| ## 🛠️ Preprocess: Remove the background and center the object. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(variant='panel', scale=0.2): | |
| with gr.Tabs(elem_id="tailor3d_processed_image"): | |
| with gr.TabItem('Processed Front-view Image'): | |
| with gr.Row(): | |
| processed_front_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=0.2): | |
| with gr.Tabs(elem_id="tailor3d_processed_image"): | |
| with gr.TabItem('Processed Back-view Image'): | |
| with gr.Row(): | |
| processed_back_image = gr.Image(label="Processed Image", image_mode="RGBA", type="filepath", elem_id="processed_image", width="auto", interactive=False) | |
| with gr.Row(): | |
| gr.Markdown( | |
| """ | |
| ## 🚀 Output: The rendering video of the 3D object. | |
| Note that the output is the 3D mesh, for convience, we showcase it through a video that circles around. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(variant='panel', scale=0.2): | |
| with gr.Tabs(elem_id="tailor3d_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(): | |
| gr.Markdown( | |
| """ | |
| ## Example in the paper. | |
| """ | |
| ) | |
| with gr.Row(): | |
| examples = [ | |
| ['assets/sample_input/demo/front/boy_astronaut.png', 'assets/sample_input/demo/back/boy_astronaut.png'], | |
| ['assets/sample_input/demo/front/boy_chinese_style.png', 'assets/sample_input/demo/back/boy_chinese_style.png'], | |
| ['assets/sample_input/demo/front/kunkun_soccer.png', 'assets/sample_input/demo/back/kunkun_soccer.png'], | |
| ['assets/sample_input/demo/front/marvel_superman.png', 'assets/sample_input/demo/back/marvel_superman.png'], | |
| ] | |
| for example in examples: | |
| with gr.Column(scale=1): | |
| gr.Examples( | |
| examples=[example], | |
| inputs=[input_front_image, input_back_image], | |
| outputs=[processed_front_image, processed_back_image, output_video], | |
| fn=example_fn, | |
| cache_examples=bool(os.getenv('SPACE_ID')), | |
| examples_per_page=3, | |
| ) | |
| # EXAMPLES | |
| with gr.Row(): | |
| examples = [ | |
| ['assets/sample_input/demo/front/lego_astronaut.png', 'assets/sample_input/demo/back/lego_astronaut.png'], | |
| ['assets/sample_input/demo/front/lego_wizard.png', 'assets/sample_input/demo/back/lego_wizard.png'], | |
| ['assets/sample_input/demo/front/bird.png', 'assets/sample_input/demo/back/bird.png'], | |
| ['assets/sample_input/demo/front/bird_whitered.png', 'assets/sample_input/demo/back/bird_whitered.png'], | |
| ] | |
| for example in examples: | |
| with gr.Column(scale=0.3): | |
| gr.Examples( | |
| examples=[example], | |
| inputs=[input_front_image, input_back_image], | |
| outputs=[processed_front_image, processed_back_image, output_video], | |
| fn=example_fn, | |
| cache_examples=bool(os.getenv('SPACE_ID')), | |
| examples_per_page=3, | |
| ) | |
| working_dir = gr.State() | |
| submit.click( | |
| fn=assert_input_image, | |
| inputs=[input_front_image, input_back_image], | |
| queue=False, | |
| ).success( | |
| fn=prepare_working_dir, | |
| outputs=[working_dir], | |
| queue=False, | |
| ).success( | |
| fn=preprocess_fn, | |
| inputs=[input_front_image, input_back_image, checkbox_rembg, checkbox_recenter, working_dir], | |
| outputs=[processed_front_image, processed_back_image], | |
| ).success( | |
| fn=core_fn, | |
| inputs=[processed_front_image, processed_back_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": "alexzyqi/Tailor3D-Base-1.0", | |
| "APP_PRETRAIN_MODEL_NAME": "zxhezexin/openlrm-mix-base-1.1", | |
| "APP_INFER": "./configs/infer-gradio-base.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() | |