update

app.py

@@ -120,7 +120,7 @@ with gr.Blocks(
                     )
                     project_delight = gr.Checkbox(
                         label="Backproject delighting",
-                        value=True,
+                        value=False,
                     )
                 gr.Markdown("Geo Structure Generation")
                 with gr.Row():

asset3d_gen/data/backproject_v2.py

@@ -220,6 +220,7 @@ class TextureBacker:
         texture_wh: tuple[int, int] = (2048, 2048),
         bake_angle_thresh: int = 75,
         mask_thresh: float = 0.5,
+        smooth_texture: bool = True,
     ):
 
         self.camera_params = camera_params
@@ -229,6 +230,7 @@ class TextureBacker:
         self.render_wh = render_wh
         self.texture_wh = texture_wh
         self.mask_thresh = mask_thresh
+        self.smooth_texture = smooth_texture
 
         self.bake_angle_thresh = bake_angle_thresh
         self.bake_unreliable_kernel_size = int(
@@ -468,7 +470,9 @@ class TextureBacker:
         texture_np, mask_np = self.compute_texture(colors, mesh)
 
         texture_np = self.uv_inpaint(mesh, texture_np, mask_np)
-        texture_np = post_process_texture(texture_np)
+        if self.smooth_texture:
+            texture_np = post_process_texture(texture_np)
+
         vertices, faces, uv_map = self.get_mesh_np_attrs(
             mesh, self.scale, self.center
         )
@@ -551,7 +555,11 @@ def parse_args():
     parser.add_argument(
         "--delight", action="store_true", help="Use delighting model."
     )
-    args = parser.parse_args()
+    parser.add_argument(
+        "--smooth_texture", type=bool, default=True, help="Smooth the texture."
+    )
+
+    args, unknown = parser.parse_known_args()
 
     return args
 
@@ -619,6 +627,7 @@ def entrypoint(
         view_weights=view_weights,
         render_wh=camera_params.resolution_hw,
         texture_wh=args.texture_wh,
+        smooth_texture=args.smooth_texture,
     )
 
     textured_mesh = texture_backer(multiviews, mesh, args.output_path)

asset3d_gen/data/utils.py

@@ -933,21 +933,11 @@ def get_images_from_grid(
     return images
 
 
-# def post_process_texture(texture: np.ndarray, iter: int = 2) -> np.ndarray:
-#     for _ in range(iter):
-#         texture = cv2.fastNlMeansDenoisingColored(texture, None, 13, 13, 9, 27)
-#         texture = cv2.bilateralFilter(
-#             texture, d=9, sigmaColor=80, sigmaSpace=80
-#         )
-
-#     return texture
-
-
 def post_process_texture(texture: np.ndarray, iter: int = 1) -> np.ndarray:
     for _ in range(iter):
-        texture = cv2.fastNlMeansDenoisingColored(texture, None, 5, 5, 7, 19)
+        texture = cv2.fastNlMeansDenoisingColored(texture, None, 2, 2, 7, 15)
         texture = cv2.bilateralFilter(
-            texture, d=7, sigmaColor=50, sigmaSpace=50
+            texture, d=5, sigmaColor=20, sigmaSpace=20
         )
 
     return texture

asset3d_gen/scripts/imageto3d.py

@@ -0,0 +1,344 @@
+import argparse
+import logging
+import os
+import sys
+from glob import glob
+
+import numpy as np
+import trimesh
+from PIL import Image
+from asset3d_gen.data.backproject_v2 import entrypoint as backproject_api
+from asset3d_gen.models.delight_model import DelightingModel
+from asset3d_gen.models.gs_model import GaussianOperator
+from asset3d_gen.models.segment_model import (
+    BMGG14Remover,
+    RembgRemover,
+    SAMPredictor,
+    trellis_preprocess,
+)
+from asset3d_gen.models.sr_model import ImageRealESRGAN
+from asset3d_gen.scripts.render_gs import entrypoint as render_gs_api
+from asset3d_gen.utils.gpt_clients import GPT_CLIENT
+from asset3d_gen.utils.process_media import (
+    merge_images_video,
+    render_asset3d,
+    render_mesh,
+    render_video,
+)
+from asset3d_gen.utils.tags import VERSION
+from asset3d_gen.validators.quality_checkers import (
+    BaseChecker,
+    ImageAestheticChecker,
+    ImageSegChecker,
+    MeshGeoChecker,
+)
+from asset3d_gen.validators.urdf_convertor import URDFGenerator
+
+current_file_path = os.path.abspath(__file__)
+current_dir = os.path.dirname(current_file_path)
+sys.path.append(os.path.join(current_dir, "../.."))
+from thirdparty.TRELLIS.trellis.pipelines import TrellisImageTo3DPipeline
+from thirdparty.TRELLIS.trellis.renderers.mesh_renderer import MeshRenderer
+from thirdparty.TRELLIS.trellis.representations import (
+    Gaussian,
+    MeshExtractResult,
+)
+from thirdparty.TRELLIS.trellis.representations.gaussian.general_utils import (
+    build_scaling_rotation,
+    inverse_sigmoid,
+    strip_symmetric,
+)
+from thirdparty.TRELLIS.trellis.utils import postprocessing_utils
+from thirdparty.TRELLIS.trellis.utils.render_utils import (
+    render_frames,
+    yaw_pitch_r_fov_to_extrinsics_intrinsics,
+)
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+os.environ["TORCH_EXTENSIONS_DIR"] = os.path.expanduser(
+    "~/.cache/torch_extensions"
+)
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "false"
+os.environ["SPCONV_ALGO"] = "native"
+
+
+DELIGHT = DelightingModel()
+IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+
+RBG_REMOVER = RembgRemover()
+RBG14_REMOVER = BMGG14Remover()
+SAM_PREDICTOR = SAMPredictor(model_type="vit_h", device="cpu")
+PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
+    "JeffreyXiang/TRELLIS-image-large"
+)
+PIPELINE.cuda()
+SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
+GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
+AESTHETIC_CHECKER = ImageAestheticChecker()
+CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
+TMP_DIR = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Image to 3D pipeline args.")
+    parser.add_argument(
+        "--image_path", type=str, nargs="+", help="Path to the input images."
+    )
+    parser.add_argument(
+        "--image_root", type=str, help="Path to the input images folder."
+    )
+    parser.add_argument(
+        "--output_root",
+        type=str,
+        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,
+        default=None,
+        help="The hight in meter to restore the mesh real size.",
+    )
+    parser.add_argument(
+        "--mass_range",
+        type=str,
+        default=None,
+        help="The mass in kg to restore the mesh real weight.",
+    )
+    parser.add_argument("--asset_type", type=str, default=None)
+    parser.add_argument("--skip_exists", action="store_true")
+    parser.add_argument("--strict_seg", action="store_true")
+    parser.add_argument("--version", type=str, default=VERSION)
+    args = parser.parse_args()
+
+    assert (
+        args.image_path or args.image_root
+    ), "Please provide either --image_path or --image_root."
+    if not args.image_path:
+        args.image_path = glob(os.path.join(args.image_root, "*.png"))
+        args.image_path += glob(os.path.join(args.image_root, "*.jpg"))
+        args.image_path += glob(os.path.join(args.image_root, "*.jpeg"))
+
+    return args
+
+
+def get_segmented_image(
+    image,
+    sam_remover,
+    rbg_remover,
+    seg_checker,
+    image_path,
+    seg_path,
+    mode="loose",
+) -> Image.Image:
+    def _is_valid_seg(img: Image.Image) -> bool:
+        return img.mode == "RGBA" and seg_checker([image_path, seg_path])[0]
+
+    seg_image = sam_remover(image, save_path=seg_path)
+    if not _is_valid_seg(seg_image):
+        logger.warning(
+            f"Failed to segment {image_path} by SAM, retry with `rembg`."
+        )  # noqa
+        seg_image = rbg_remover(image, save_path=seg_path)
+
+        if not _is_valid_seg(seg_image):
+            if mode == "strict":
+                raise RuntimeError(
+                    f"Failed to segment {image_path} by SAM and rembg, abort."
+                )
+            logger.warning(
+                f"Failed to segment {image_path} by rembg, use raw image."
+            )  # noqa
+            seg_image = image.convert("RGBA")
+            seg_image.save(seg_path)
+
+    return seg_image
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    for image_path in args.image_path:
+        try:
+            filename = os.path.basename(image_path).split(".")[0]
+            output_root = args.output_root
+            if args.image_root is not None:
+                output_root = os.path.join(output_root, filename)
+            os.makedirs(output_root, exist_ok=True)
+
+            mesh_out = f"{output_root}/{filename}.obj"
+            if args.skip_exists and os.path.exists(mesh_out):
+                logger.info(
+                    f"Skip {image_path}, already processed in {mesh_out}"
+                )
+                continue
+
+            image = Image.open(image_path)
+            image.save(f"{output_root}/{filename}_raw.png")
+
+            # Segmentation: Get segmented image using SAM or Rembg.
+            seg_path = f"{output_root}/{filename}_cond.png"
+            if image.mode != "RGBA":
+                seg_image = RBG_REMOVER(image, save_path=seg_path)
+                seg_image = trellis_preprocess(seg_image)
+            else:
+                seg_image = image
+                seg_image.save(seg_path)
+
+            # Run the pipeline
+            try:
+                outputs = PIPELINE.run(
+                    seg_image,
+                    preprocess_image=False,
+                    # Optional parameters
+                    # seed=1,
+                    # sparse_structure_sampler_params={
+                    #     "steps": 12,
+                    #     "cfg_strength": 7.5,
+                    # },
+                    # slat_sampler_params={
+                    #     "steps": 12,
+                    #     "cfg_strength": 3,
+                    # },
+                )
+            except Exception as e:
+                logger.error(
+                    f"[Pipeline Failed] process {image_path}: {e}, skip."
+                )
+                continue
+
+            # Render and save color and mesh videos
+            gs_model = outputs["gaussian"][0]
+            mesh_model = outputs["mesh"][0]
+            color_images = render_video(gs_model)["color"]
+            normal_images = render_video(mesh_model)["normal"]
+            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],
+                )
+
+                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)
+
+        except Exception as e:
+            logger.error(f"Failed to process {image_path}: {e}, skip.")
+            continue
+
+    logger.info(f"Processing complete. Outputs saved to {args.output_root}")

asset3d_gen/scripts/render_gs.py

@@ -75,7 +75,7 @@ def parse_args():
         help="Output image size for single view in color grid (default: 512)",
     )
 
-    args = parser.parse_args()
+    args, unknown = parser.parse_known_args()
 
     return args
 

asset3d_gen/utils/process_media.py

@@ -3,6 +3,7 @@ import logging
 import math
 import os
 import subprocess
+import sys
 from glob import glob
 from io import BytesIO
 from typing import Union
@@ -12,7 +13,23 @@ import imageio
 import numpy as np
 import PIL.Image as Image
 import spaces
+import torch
 from moviepy.editor import VideoFileClip, clips_array
+from tqdm import tqdm
+
+current_file_path = os.path.abspath(__file__)
+current_dir = os.path.dirname(current_file_path)
+sys.path.append(os.path.join(current_dir, "../.."))
+from thirdparty.TRELLIS.trellis.pipelines import TrellisImageTo3DPipeline
+from thirdparty.TRELLIS.trellis.renderers.mesh_renderer import MeshRenderer
+from thirdparty.TRELLIS.trellis.representations import (
+    Gaussian,
+    MeshExtractResult,
+)
+from thirdparty.TRELLIS.trellis.utils.render_utils import (
+    render_frames,
+    yaw_pitch_r_fov_to_extrinsics_intrinsics,
+)
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -24,6 +41,8 @@ __all__ = [
     "filter_small_connected_components",
     "filter_image_small_connected_components",
     "combine_images_to_base64",
+    "render_mesh",
+    "render_video",
 ]
 
 
@@ -176,6 +195,59 @@ def combine_images_to_base64(
     return base64.b64encode(buffer.getvalue()).decode("utf-8")
 
 
+@spaces.GPU
+def render_mesh(sample, extrinsics, intrinsics, options={}, **kwargs):
+    renderer = MeshRenderer()
+    renderer.rendering_options.resolution = options.get("resolution", 512)
+    renderer.rendering_options.near = options.get("near", 1)
+    renderer.rendering_options.far = options.get("far", 100)
+    renderer.rendering_options.ssaa = options.get("ssaa", 4)
+    rets = {}
+    for extr, intr in tqdm(zip(extrinsics, intrinsics), desc="Rendering"):
+        res = renderer.render(sample, extr, intr)
+        if "normal" not in rets:
+            rets["normal"] = []
+        normal = torch.lerp(
+            torch.zeros_like(res["normal"]), res["normal"], res["mask"]
+        )
+        normal = np.clip(
+            normal.detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255
+        ).astype(np.uint8)
+        rets["normal"].append(normal)
+
+    return rets
+
+
+@spaces.GPU
+def render_video(
+    sample,
+    resolution=512,
+    bg_color=(0, 0, 0),
+    num_frames=300,
+    r=2,
+    fov=40,
+    **kwargs,
+):
+    yaws = torch.linspace(0, 2 * 3.1415, num_frames)
+    yaws = yaws.tolist()
+    pitch = [0.5] * num_frames
+    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(
+        yaws, pitch, r, fov
+    )
+    render_fn = (
+        render_mesh if isinstance(sample, MeshExtractResult) else render_frames
+    )
+    result = render_fn(
+        sample,
+        extrinsics,
+        intrinsics,
+        {"resolution": resolution, "bg_color": bg_color},
+        **kwargs,
+    )
+
+    return result
+
+
 if __name__ == "__main__":
     # Example usage:
     merge_video_video(

common.py

@@ -26,7 +26,7 @@ from asset3d_gen.models.segment_model import (
     SAMPredictor,
     trellis_preprocess,
 )
-from asset3d_gen.models.sr_model import ImageRealESRGAN
+from asset3d_gen.models.sr_model import ImageRealESRGAN, ImageStableSR
 from asset3d_gen.scripts.render_gs import entrypoint as render_gs_api
 from asset3d_gen.scripts.render_mv import build_texture_gen_pipe, infer_pipe
 from asset3d_gen.scripts.text2image import (
@@ -39,6 +39,8 @@ from asset3d_gen.utils.process_media import (
     filter_image_small_connected_components,
     merge_images_video,
     render_asset3d,
+    render_mesh,
+    render_video,
 )
 from asset3d_gen.utils.tags import VERSION
 from asset3d_gen.validators.quality_checkers import (
@@ -84,6 +86,7 @@ os.environ["SPCONV_ALGO"] = "native"
 MAX_SEED = 100000
 DELIGHT = DelightingModel()
 IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+# IMAGESR_MODEL = ImageStableSR()
 
 
 def patched_setup_functions(self):
@@ -234,59 +237,6 @@ def end_session(req: gr.Request) -> None:
         shutil.rmtree(user_dir)
 
 
-@spaces.GPU
-def render_mesh(sample, extrinsics, intrinsics, options={}, **kwargs):
-    renderer = MeshRenderer()
-    renderer.rendering_options.resolution = options.get("resolution", 512)
-    renderer.rendering_options.near = options.get("near", 1)
-    renderer.rendering_options.far = options.get("far", 100)
-    renderer.rendering_options.ssaa = options.get("ssaa", 4)
-    rets = {}
-    for extr, intr in tqdm(zip(extrinsics, intrinsics), desc="Rendering"):
-        res = renderer.render(sample, extr, intr)
-        if "normal" not in rets:
-            rets["normal"] = []
-        normal = torch.lerp(
-            torch.zeros_like(res["normal"]), res["normal"], res["mask"]
-        )
-        normal = np.clip(
-            normal.detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255
-        ).astype(np.uint8)
-        rets["normal"].append(normal)
-
-    return rets
-
-
-@spaces.GPU
-def render_video(
-    sample,
-    resolution=512,
-    bg_color=(0, 0, 0),
-    num_frames=300,
-    r=2,
-    fov=40,
-    **kwargs,
-):
-    yaws = torch.linspace(0, 2 * 3.1415, num_frames)
-    yaws = yaws.tolist()
-    pitch = [0.5] * num_frames
-    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(
-        yaws, pitch, r, fov
-    )
-    render_fn = (
-        render_mesh if isinstance(sample, MeshExtractResult) else render_frames
-    )
-    result = render_fn(
-        sample,
-        extrinsics,
-        intrinsics,
-        {"resolution": resolution, "bg_color": bg_color},
-        **kwargs,
-    )
-
-    return result
-
-
 @spaces.GPU
 def preprocess_image_fn(
     image: str | np.ndarray | Image.Image, rmbg_tag: str = "rembg"
@@ -495,11 +445,11 @@ def image_to_3d(
 
 @spaces.GPU
 def extract_3d_representations(
-    state: dict, enable_delight: bool, req: gr.Request
+    state: dict, enable_delight: bool, texture_size: int, req: gr.Request
 ):
     output_root = TMP_DIR
     output_root = os.path.join(output_root, str(req.session_hash))
-    gs_model, mesh_model = unpack_state(state)
+    gs_model, mesh_model = unpack_state(state, device="cuda")
 
     mesh = postprocessing_utils.to_glb(
         gs_model,