feat(texture): Refine texture back-projection. (#58)

Author: HochCC

README.md

@@ -331,4 +331,4 @@ EmbodiedGen builds upon the following amazing projects and models:
 
 ## ⚖️ License
 
-This project is licensed under the [Apache License 2.0](LICENSE). See the `LICENSE` file for details.
+This project is licensed under the [Apache License 2.0](docs/LICENSE). See the `LICENSE` file for details.

apps/app_style.py

@@ -4,7 +4,7 @@
 lighting_css = """
 <style>
 #lighter_mesh canvas {
-    filter: brightness(1.9) !important;
+    filter: brightness(2.0) !important;
 }
 </style>
 """

apps/common.py

@@ -32,8 +32,9 @@
 from easydict import EasyDict as edict
 from PIL import Image
 from embodied_gen.data.backproject_v2 import entrypoint as backproject_api
+from embodied_gen.data.backproject_v3 import entrypoint as backproject_api_v3
 from embodied_gen.data.differentiable_render import entrypoint as render_api
-from embodied_gen.data.utils import trellis_preprocess, zip_files
+from embodied_gen.data.utils import resize_pil, trellis_preprocess, zip_files
 from embodied_gen.models.delight_model import DelightingModel
 from embodied_gen.models.gs_model import GaussianOperator
 from embodied_gen.models.segment_model import (
@@ -131,8 +132,8 @@ def build_covariance_from_scaling_rotation(
 Gaussian.setup_functions = patched_setup_functions
 
 
-DELIGHT = DelightingModel()
-IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+# DELIGHT = DelightingModel()
+# IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
 # IMAGESR_MODEL = ImageStableSR()
 if os.getenv("GRADIO_APP") == "imageto3d":
     RBG_REMOVER = RembgRemover()
@@ -169,6 +170,8 @@ def build_covariance_from_scaling_rotation(
     )
     os.makedirs(TMP_DIR, exist_ok=True)
 elif os.getenv("GRADIO_APP") == "texture_edit":
+    DELIGHT = DelightingModel()
+    IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
     PIPELINE_IP = build_texture_gen_pipe(
         base_ckpt_dir="./weights",
         ip_adapt_scale=0.7,
@@ -205,7 +208,7 @@ def preprocess_image_fn(
     elif isinstance(image, np.ndarray):
         image = Image.fromarray(image)
 
-    image_cache = image.copy().resize((512, 512))
+    image_cache = resize_pil(image.copy(), 1024)
 
     bg_remover = RBG_REMOVER if rmbg_tag == "rembg" else RBG14_REMOVER
     image = bg_remover(image)
@@ -221,7 +224,7 @@ def preprocess_sam_image_fn(
         image = Image.fromarray(image)
 
     sam_image = SAM_PREDICTOR.preprocess_image(image)
-    image_cache = Image.fromarray(sam_image).resize((512, 512))
+    image_cache = sam_image.copy()
     SAM_PREDICTOR.predictor.set_image(sam_image)
 
     return sam_image, image_cache
@@ -512,6 +515,60 @@ def extract_3d_representations_v2(
     return mesh_glb_path, gs_path, mesh_obj_path, aligned_gs_path
 
 
+def extract_3d_representations_v3(
+    state: dict,
+    enable_delight: bool,
+    texture_size: int,
+    req: gr.Request,
+):
+    output_root = TMP_DIR
+    user_dir = os.path.join(output_root, str(req.session_hash))
+    gs_model, mesh_model = unpack_state(state, device="cpu")
+
+    filename = "sample"
+    gs_path = os.path.join(user_dir, f"{filename}_gs.ply")
+    gs_model.save_ply(gs_path)
+
+    # Rotate mesh and GS by 90 degrees around Z-axis.
+    rot_matrix = [[0, 0, -1], [0, 1, 0], [1, 0, 0]]
+    gs_add_rot = [[1, 0, 0], [0, -1, 0], [0, 0, -1]]
+    mesh_add_rot = [[1, 0, 0], [0, 0, -1], [0, 1, 0]]
+
+    # Addtional rotation for GS to align mesh.
+    gs_rot = np.array(gs_add_rot) @ np.array(rot_matrix)
+    pose = GaussianOperator.trans_to_quatpose(gs_rot)
+    aligned_gs_path = gs_path.replace(".ply", "_aligned.ply")
+    GaussianOperator.resave_ply(
+        in_ply=gs_path,
+        out_ply=aligned_gs_path,
+        instance_pose=pose,
+        device="cpu",
+    )
+
+    mesh = trimesh.Trimesh(
+        vertices=mesh_model.vertices.cpu().numpy(),
+        faces=mesh_model.faces.cpu().numpy(),
+    )
+    mesh.vertices = mesh.vertices @ np.array(mesh_add_rot)
+    mesh.vertices = mesh.vertices @ np.array(rot_matrix)
+
+    mesh_obj_path = os.path.join(user_dir, f"{filename}.obj")
+    mesh.export(mesh_obj_path)
+
+    mesh = backproject_api_v3(
+        gs_path=aligned_gs_path,
+        mesh_path=mesh_obj_path,
+        output_path=mesh_obj_path,
+        skip_fix_mesh=False,
+        texture_size=texture_size,
+    )
+
+    mesh_glb_path = os.path.join(user_dir, f"{filename}.glb")
+    mesh.export(mesh_glb_path)
+
+    return mesh_glb_path, gs_path, mesh_obj_path, aligned_gs_path
+
+
 def extract_urdf(
     gs_path: str,
     mesh_obj_path: str,

apps/image_to_3d.py

@@ -27,7 +27,7 @@
     VERSION,
     active_btn_by_content,
     end_session,
-    extract_3d_representations_v2,
+    extract_3d_representations_v3,
     extract_urdf,
     get_seed,
     image_to_3d,
@@ -179,17 +179,17 @@
                     )
 
             generate_btn = gr.Button(
-                "🚀 1. Generate(~0.5 mins)",
+                "🚀 1. Generate(~2 mins)",
                 variant="primary",
                 interactive=False,
             )
             model_output_obj = gr.Textbox(label="raw mesh .obj", visible=False)
-            with gr.Row():
-                extract_rep3d_btn = gr.Button(
-                    "🔍 2. Extract 3D Representation(~2 mins)",
-                    variant="primary",
-                    interactive=False,
-                )
+            # with gr.Row():
+            #     extract_rep3d_btn = gr.Button(
+            #         "🔍 2. Extract 3D Representation(~2 mins)",
+            #         variant="primary",
+            #         interactive=False,
+            #     )
             with gr.Accordion(
                 label="Enter Asset Attributes(optional)", open=False
             ):
@@ -207,7 +207,7 @@
                 )
             with gr.Row():
                 extract_urdf_btn = gr.Button(
-                    "🧩 3. Extract URDF with physics(~1 mins)",
+                    "🧩 2. Extract URDF with physics(~1 mins)",
                     variant="primary",
                     interactive=False,
                 )
@@ -230,7 +230,7 @@
                 )
             with gr.Row():
                 download_urdf = gr.DownloadButton(
-                    label="⬇️ 4. Download URDF",
+                    label="⬇️ 3. Download URDF",
                     variant="primary",
                     interactive=False,
                 )
@@ -326,7 +326,7 @@
     image_prompt.change(
         lambda: tuple(
             [
-                gr.Button(interactive=False),
+                # gr.Button(interactive=False),
                 gr.Button(interactive=False),
                 gr.Button(interactive=False),
                 None,
@@ -344,7 +344,7 @@
             ]
         ),
         outputs=[
-            extract_rep3d_btn,
+            # extract_rep3d_btn,
             extract_urdf_btn,
             download_urdf,
             model_output_gs,
@@ -375,7 +375,7 @@
     image_prompt_sam.change(
         lambda: tuple(
             [
-                gr.Button(interactive=False),
+                # gr.Button(interactive=False),
                 gr.Button(interactive=False),
                 gr.Button(interactive=False),
                 None,
@@ -394,7 +394,7 @@
             ]
         ),
         outputs=[
-            extract_rep3d_btn,
+            # extract_rep3d_btn,
             extract_urdf_btn,
             download_urdf,
             model_output_gs,
@@ -447,12 +447,7 @@
         ],
         outputs=[output_buf, video_output],
     ).success(
-        lambda: gr.Button(interactive=True),
-        outputs=[extract_rep3d_btn],
-    )
-
-    extract_rep3d_btn.click(
-        extract_3d_representations_v2,
+        extract_3d_representations_v3,
         inputs=[
             output_buf,
             project_delight,
@@ -495,4 +490,4 @@
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_port=8081)

apps/text_to_3d.py

@@ -27,7 +27,7 @@
     VERSION,
     active_btn_by_text_content,
     end_session,
-    extract_3d_representations_v2,
+    extract_3d_representations_v3,
     extract_urdf,
     get_cached_image,
     get_seed,
@@ -178,17 +178,17 @@
                     )
 
             generate_btn = gr.Button(
-                "🚀 2. Generate 3D(~0.5 mins)",
+                "🚀 2. Generate 3D(~2 mins)",
                 variant="primary",
                 interactive=False,
             )
             model_output_obj = gr.Textbox(label="raw mesh .obj", visible=False)
-            with gr.Row():
-                extract_rep3d_btn = gr.Button(
-                    "🔍 3. Extract 3D Representation(~1 mins)",
-                    variant="primary",
-                    interactive=False,
-                )
+            # with gr.Row():
+            #     extract_rep3d_btn = gr.Button(
+            #         "🔍 3. Extract 3D Representation(~1 mins)",
+            #         variant="primary",
+            #         interactive=False,
+            #     )
             with gr.Accordion(
                 label="Enter Asset Attributes(optional)", open=False
             ):
@@ -206,13 +206,13 @@
                 )
             with gr.Row():
                 extract_urdf_btn = gr.Button(
-                    "🧩 4. Extract URDF with physics(~1 mins)",
+                    "🧩 3. Extract URDF with physics(~1 mins)",
                     variant="primary",
                     interactive=False,
                 )
             with gr.Row():
                 download_urdf = gr.DownloadButton(
-                    label="⬇️ 5. Download URDF",
+                    label="⬇️ 4. Download URDF",
                     variant="primary",
                     interactive=False,
                 )
@@ -336,7 +336,7 @@
     generate_img_btn.click(
         lambda: tuple(
             [
-                gr.Button(interactive=False),
+                # gr.Button(interactive=False),
                 gr.Button(interactive=False),
                 gr.Button(interactive=False),
                 gr.Button(interactive=False),
@@ -358,7 +358,7 @@
             ]
         ),
         outputs=[
-            extract_rep3d_btn,
+            # extract_rep3d_btn,
             extract_urdf_btn,
             download_urdf,
             generate_btn,
@@ -428,12 +428,7 @@
         ],
         outputs=[output_buf, video_output],
     ).success(
-        lambda: gr.Button(interactive=True),
-        outputs=[extract_rep3d_btn],
-    )
-
-    extract_rep3d_btn.click(
-        extract_3d_representations_v2,
+        extract_3d_representations_v3,
         inputs=[
             output_buf,
             project_delight,
@@ -476,4 +471,4 @@
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_port=8082)

apps/texture_edit.py

@@ -381,4 +381,4 @@ def active_btn_by_content(mesh_content: gr.Model3D, text_content: gr.Textbox):
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_port=8083)

apps/visualize_asset.py

@@ -727,7 +727,6 @@ def update_assets(p, s, c):
 
 if __name__ == "__main__":
     demo.launch(
-        server_name="10.34.8.77",
         server_port=8088,
         allowed_paths=[
             "/horizon-bucket/robot_lab/datasets/embodiedgen/assets"

LICENSE docs/LICENSELICENSE renamed to docs/LICENSE

@@ -14,6 +14,8 @@ conda activate embodiedgen
 bash install.sh basic
 ```
 
+Please `huggingface-cli login` to ensure that the ckpts can be downloaded automatically afterwards.
+
 ## ✅ Starting from Docker
 
 We provide a pre-built Docker image on [Docker Hub](https://hub.docker.com/repository/docker/wangxinjie/embodiedgen) with a configured environment for your convenience. For more details, please refer to [Docker documentation](https://github.com/HorizonRobotics/EmbodiedGen/tree/master/docker).

docs/install.md

@@ -589,6 +589,8 @@ def convert(self, urdf_path: str, output_file: str):
         stage = Usd.Stage.Open(usd_path)
         layer = stage.GetRootLayer()
         with Usd.EditContext(stage, layer):
+            base_prim = stage.GetPseudoRoot().GetChildren()[0]
+            base_prim.SetMetadata("kind", "component")
             for prim in stage.Traverse():
                 # Change texture path to relative path.
                 if prim.GetName() == "material_0":