update

app.py

@@ -51,7 +51,7 @@ with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo:
             <a href="https://arxiv.org/abs/xxxx.xxxxx">
                 <img alt="📄 arXiv" src="https://img.shields.io/badge/📄-arXiv-b31b1b">
             </a>
-            <a href="https://github.com/horizon-research/EmbodiedGen">
+            <a href="https://github.com/HorizonRobotics/EmbodiedGen">
                 <img alt="💻 GitHub" src="https://img.shields.io/badge/GitHub-000000?logo=github">
             </a>
             <a href="https://www.youtube.com/watch?v=SnHhzHeb_aI">
@@ -68,7 +68,7 @@ with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo:
     )
 
     gr.HTML(image_css)
-    gr.HTML(lighting_css)
+    # gr.HTML(lighting_css)
     with gr.Row():
         with gr.Column(scale=2):
             with gr.Tabs() as input_tabs:

embodied_gen/data/backproject_v2.py

@@ -606,7 +606,17 @@ def parse_args():
         "--delight", action="store_true", help="Use delighting model."
     )
     parser.add_argument(
-        "--smooth_texture", type=bool, default=True, help="Smooth the texture."
+        "--no_smooth_texture",
+        action="store_true",
+        help="Do not smooth the texture.",
+    )
+    parser.add_argument(
+        "--save_glb_path", type=str, default=None, help="Save glb path."
+    )
+    parser.add_argument(
+        "--no_save_delight_img",
+        action="store_true",
+        help="Disable saving delight image",
     )
 
     args, unknown = parser.parse_known_args()
@@ -642,7 +652,8 @@ def entrypoint(
         save_dir = os.path.dirname(args.output_path)
         os.makedirs(save_dir, exist_ok=True)
         color_grid = delight_model(color_grid)
-        color_grid.save(f"{save_dir}/color_grid_delight.png")
+        if not args.no_save_delight_img:
+            color_grid.save(f"{save_dir}/color_grid_delight.png")
 
     multiviews = get_images_from_grid(color_grid, img_size=512)
 
@@ -675,11 +686,15 @@ def entrypoint(
         view_weights=view_weights,
         render_wh=camera_params.resolution_hw,
         texture_wh=args.texture_wh,
-        smooth_texture=args.smooth_texture,
+        smooth_texture=not args.no_smooth_texture,
     )
 
     textured_mesh = texture_backer(multiviews, mesh, args.output_path)
 
+    if args.save_glb_path is not None:
+        os.makedirs(os.path.dirname(args.save_glb_path), exist_ok=True)
+        textured_mesh.export(args.save_glb_path)
+
     return textured_mesh
 
 

embodied_gen/data/utils.py

@@ -19,6 +19,7 @@ import math
 import os
 import random
 import zipfile
+from shutil import rmtree
 from typing import List, Tuple, Union
 
 import cv2
@@ -66,6 +67,7 @@ __all__ = [
     "gamma_shs",
     "resize_pil",
     "trellis_preprocess",
+    "delete_dir",
 ]
 
 
@@ -994,3 +996,14 @@ def zip_files(input_paths: list[str], output_zip: str) -> str:
                 zipf.write(input_path, arcname=arcname)
 
     return output_zip
+
+
+def delete_dir(folder_path: str, keep_subs: list[str] = None) -> None:
+    for item in os.listdir(folder_path):
+        if keep_subs is not None and item in keep_subs:
+            continue
+        item_path = os.path.join(folder_path, item)
+        if os.path.isdir(item_path):
+            rmtree(item_path)
+        else:
+            os.remove(item_path)

embodied_gen/scripts/imageto3d.py

@@ -20,12 +20,13 @@ import logging
 import os
 import sys
 from glob import glob
+from shutil import copy, copytree
 
 import numpy as np
 import trimesh
 from PIL import Image
 from embodied_gen.data.backproject_v2 import entrypoint as backproject_api
-from embodied_gen.data.utils import trellis_preprocess
+from embodied_gen.data.utils import delete_dir, trellis_preprocess
 from embodied_gen.models.delight_model import DelightingModel
 from embodied_gen.models.gs_model import GaussianOperator
 from embodied_gen.models.segment_model import (
@@ -97,9 +98,6 @@ def parse_args():
         required=True,
         help="Root directory for saving outputs.",
     )
-    parser.add_argument(
-        "--no_mesh", action="store_true", help="Do not output mesh files."
-    )
     parser.add_argument(
         "--height_range",
         type=str,
@@ -116,6 +114,7 @@ def parse_args():
     parser.add_argument("--skip_exists", action="store_true")
     parser.add_argument("--strict_seg", action="store_true")
     parser.add_argument("--version", type=str, default=VERSION)
+    parser.add_argument("--remove_intermediate", type=bool, default=True)
     args = parser.parse_args()
 
     assert (
@@ -136,7 +135,7 @@ if __name__ == "__main__":
         try:
             filename = os.path.basename(image_path).split(".")[0]
             output_root = args.output_root
-            if args.image_root is not None:
+            if args.image_root is not None or len(args.image_path) > 1:
                 output_root = os.path.join(output_root, filename)
             os.makedirs(output_root, exist_ok=True)
 
@@ -189,117 +188,121 @@ if __name__ == "__main__":
             video_path = os.path.join(output_root, "gs_mesh.mp4")
             merge_images_video(color_images, normal_images, video_path)
 
-            if not args.no_mesh:
-                # Save the raw Gaussian model
-                gs_path = mesh_out.replace(".obj", "_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",
-                )
-                color_path = os.path.join(output_root, "color.png")
-                render_gs_api(aligned_gs_path, color_path)
-
-                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(output_root, f"{filename}.obj")
-                mesh.export(mesh_obj_path)
-
-                mesh = backproject_api(
-                    delight_model=DELIGHT,
-                    imagesr_model=IMAGESR_MODEL,
-                    color_path=color_path,
-                    mesh_path=mesh_obj_path,
-                    output_path=mesh_obj_path,
-                    skip_fix_mesh=False,
-                    delight=True,
-                    texture_wh=[2048, 2048],
+            # Save the raw Gaussian model
+            gs_path = mesh_out.replace(".obj", "_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",
+            )
+            color_path = os.path.join(output_root, "color.png")
+            render_gs_api(aligned_gs_path, color_path)
+
+            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(output_root, f"{filename}.obj")
+            mesh.export(mesh_obj_path)
+
+            mesh = backproject_api(
+                delight_model=DELIGHT,
+                imagesr_model=IMAGESR_MODEL,
+                color_path=color_path,
+                mesh_path=mesh_obj_path,
+                output_path=mesh_obj_path,
+                skip_fix_mesh=False,
+                delight=True,
+                texture_wh=[2048, 2048],
+            )
+
+            mesh_glb_path = os.path.join(output_root, f"{filename}.glb")
+            mesh.export(mesh_glb_path)
+
+            urdf_convertor = URDFGenerator(GPT_CLIENT, render_view_num=4)
+            asset_attrs = {
+                "version": VERSION,
+                "gs_model": f"{urdf_convertor.output_mesh_dir}/{filename}_gs.ply",
+            }
+            if args.height_range:
+                min_height, max_height = map(
+                    float, args.height_range.split("-")
                 )
-
-                mesh_glb_path = os.path.join(output_root, f"{filename}.glb")
-                mesh.export(mesh_glb_path)
-
-                urdf_convertor = URDFGenerator(GPT_CLIENT, render_view_num=4)
-                asset_attrs = {
-                    "version": VERSION,
-                    "gs_model": f"{urdf_convertor.output_mesh_dir}/{filename}_gs.ply",
-                }
-                if args.height_range:
-                    min_height, max_height = map(
-                        float, args.height_range.split("-")
-                    )
-                    asset_attrs["min_height"] = min_height
-                    asset_attrs["max_height"] = max_height
-                if args.mass_range:
-                    min_mass, max_mass = map(float, args.mass_range.split("-"))
-                    asset_attrs["min_mass"] = min_mass
-                    asset_attrs["max_mass"] = max_mass
-                if args.asset_type:
-                    asset_attrs["category"] = args.asset_type
-                if args.version:
-                    asset_attrs["version"] = args.version
-
-                urdf_path = urdf_convertor(
-                    mesh_path=mesh_obj_path,
-                    output_root=f"{output_root}/URDF_{filename}",
-                    **asset_attrs,
-                )
-
-                # Rescale GS and save to URDF/mesh folder.
-                real_height = urdf_convertor.get_attr_from_urdf(
-                    urdf_path, attr_name="real_height"
-                )
-                out_gs = f"{output_root}/URDF_{filename}/{urdf_convertor.output_mesh_dir}/{filename}_gs.ply"  # noqa
-                GaussianOperator.resave_ply(
-                    in_ply=aligned_gs_path,
-                    out_ply=out_gs,
-                    real_height=real_height,
-                    device="cpu",
-                )
-
-                # Quality check and update .urdf file.
-                mesh_out = f"{output_root}/URDF_{filename}/{urdf_convertor.output_mesh_dir}/{filename}.obj"  # noqa
-                trimesh.load(mesh_out).export(mesh_out.replace(".obj", ".glb"))
-                # image_paths = render_asset3d(
-                #     mesh_path=mesh_out,
-                #     output_root=f"{output_root}/URDF_{filename}",
-                #     output_subdir="qa_renders",
-                #     num_images=8,
-                #     elevation=(30, -30),
-                #     distance=5.5,
-                # )
-
-                image_dir = f"{output_root}/URDF_{filename}/{urdf_convertor.output_render_dir}/image_color"  # noqa
-                image_paths = glob(f"{image_dir}/*.png")
-                images_list = []
-                for checker in CHECKERS:
-                    images = image_paths
-                    if isinstance(checker, ImageSegChecker):
-                        images = [
-                            f"{output_root}/{filename}_raw.png",
-                            f"{output_root}/{filename}_cond.png",
-                        ]
-                    images_list.append(images)
-
-                results = BaseChecker.validate(CHECKERS, images_list)
-                urdf_convertor.add_quality_tag(urdf_path, results)
+                asset_attrs["min_height"] = min_height
+                asset_attrs["max_height"] = max_height
+            if args.mass_range:
+                min_mass, max_mass = map(float, args.mass_range.split("-"))
+                asset_attrs["min_mass"] = min_mass
+                asset_attrs["max_mass"] = max_mass
+            if args.asset_type:
+                asset_attrs["category"] = args.asset_type
+            if args.version:
+                asset_attrs["version"] = args.version
+
+            urdf_root = f"{output_root}/URDF_{filename}"
+            urdf_path = urdf_convertor(
+                mesh_path=mesh_obj_path,
+                output_root=urdf_root,
+                **asset_attrs,
+            )
+
+            # Rescale GS and save to URDF/mesh folder.
+            real_height = urdf_convertor.get_attr_from_urdf(
+                urdf_path, attr_name="real_height"
+            )
+            out_gs = f"{urdf_root}/{urdf_convertor.output_mesh_dir}/{filename}_gs.ply"  # noqa
+            GaussianOperator.resave_ply(
+                in_ply=aligned_gs_path,
+                out_ply=out_gs,
+                real_height=real_height,
+                device="cpu",
+            )
+
+            # Quality check and update .urdf file.
+            mesh_out = f"{urdf_root}/{urdf_convertor.output_mesh_dir}/{filename}.obj"  # noqa
+            trimesh.load(mesh_out).export(mesh_out.replace(".obj", ".glb"))
+
+            image_dir = f"{urdf_root}/{urdf_convertor.output_render_dir}/image_color"  # noqa
+            image_paths = glob(f"{image_dir}/*.png")
+            images_list = []
+            for checker in CHECKERS:
+                images = image_paths
+                if isinstance(checker, ImageSegChecker):
+                    images = [
+                        f"{output_root}/{filename}_raw.png",
+                        f"{output_root}/{filename}_cond.png",
+                    ]
+                images_list.append(images)
+
+            results = BaseChecker.validate(CHECKERS, images_list)
+            urdf_convertor.add_quality_tag(urdf_path, results)
+
+            # Organize the final result files
+            result_dir = f"{output_root}/result"
+            os.makedirs(result_dir, exist_ok=True)
+            copy(urdf_path, f"{result_dir}/{os.path.basename(urdf_path)}")
+            copytree(
+                f"{urdf_root}/{urdf_convertor.output_mesh_dir}",
+                f"{result_dir}/{urdf_convertor.output_mesh_dir}",
+            )
+            copy(video_path, f"{result_dir}/video.mp4")
+            if args.remove_intermediate:
+                delete_dir(output_root, keep_subs=["result"])
 
         except Exception as e:
             logger.error(f"Failed to process {image_path}: {e}, skip.")

embodied_gen/scripts/texture_gen.sh

@@ -56,8 +56,10 @@ python embodied_gen/scripts/render_mv.py \
 backproject-cli --mesh_path ${mesh_path} \
     --color_path ${output_root}/multi_view/color_sample0.png \
     --output_path "${output_root}/texture_mesh/${uuid}.obj" \
+    --save_glb_path "${output_root}/texture_mesh/${uuid}.glb" \
     --skip_fix_mesh \
-    --delight
+    --delight \
+    --no_save_delight_img
 
 # Step 4: final rendering of textured mesh
 drender-cli --mesh_path "${output_root}/texture_mesh/${uuid}.obj" \
@@ -67,3 +69,12 @@ drender-cli --mesh_path "${output_root}/texture_mesh/${uuid}.obj" \
     --with_mtl \
     --gen_color_mp4 \
     --pbr_light_factor 1.2
+
+# Organize folders
+rm -rf ${output_root}/condition
+video_path="${output_root}/texture_mesh/${uuid}/color.mp4"
+if [ -f "${video_path}" ]; then
+    cp "${video_path}" "${output_root}/texture_mesh/color.mp4"
+    echo "Resave video to ${output_root}/texture_mesh/color.mp4"
+fi
+rm -rf ${output_root}/texture_mesh/${uuid}