update

app.py

@@ -1,35 +1,8 @@
 import os
-import shutil
-from functools import partial
-
-import gradio as gr
-from common import (
-    MAX_SEED,
-    VERSION,
-    TrellisImageTo3DPipeline,
-    active_btn_by_content,
-    extract_3d_representations_v2,
-    extract_urdf,
-    get_seed,
-    image_to_3d,
-    preprocess_image_fn,
-    preprocess_sam_image_fn,
-    select_point,
-)
-from gradio.themes import Default
-from gradio.themes.utils.colors import slate
-from gradio_litmodel3d import LitModel3D
-from asset3d_gen.models.delight import DelightingModel
-from asset3d_gen.models.segment import RembgRemover, SAMPredictor
-from asset3d_gen.models.super_resolution import ImageRealESRGAN
-from asset3d_gen.utils.gpt_clients import GPT_CLIENT
-from asset3d_gen.validators.quality_checkers import (
-    ImageAestheticChecker,
-    ImageSegChecker,
-    MeshGeoChecker,
-)
-from asset3d_gen.validators.urdf_convertor import URDFGenerator
+import shutil 
 
+import gradio as gr 
+from gradio_litmodel3d import LitModel3D 
 TMP_DIR = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
 )

asset3d_gen/data/backproject.py

@@ -1,503 +0,0 @@
-import argparse
-import logging
-import math
-import os
-from typing import List, Literal, Tuple, Union
-
-import cv2
-import numpy as np
-import nvdiffrast.torch as dr
-import torch
-import trimesh
-import utils3d
-import xatlas
-from tqdm import tqdm
-from asset3d_gen.data.mesh_operator import MeshFixer
-from asset3d_gen.data.utils import (
-    CameraSetting,
-    get_images_from_grid,
-    init_kal_camera,
-    normalize_vertices_array,
-    post_process_texture,
-    save_mesh_with_mtl,
-)
-from asset3d_gen.models.delight import DelightingModel
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-class TextureBaker(object):
-    """Baking textures onto a mesh from multiple observations.
-
-    This class take 3D mesh data, camera settings and texture baking parameters
-    to generate texture map by projecting images to the mesh from diff views.
-    It supports both a fast texture baking approach and a more optimized method
-    with total variation regularization.
-
-    Attributes:
-        vertices (torch.Tensor): The vertices of the mesh.
-        faces (torch.Tensor): The faces of the mesh, defined by vertex indices.
-        uvs (torch.Tensor): The UV coordinates of the mesh.
-        camera_params (CameraSetting): Camera setting (intrinsics, extrinsics).
-        device (str): The device to run computations on ("cpu" or "cuda").
-        w2cs (torch.Tensor): World-to-camera transformation matrices.
-        projections (torch.Tensor): Camera projection matrices.
-
-    Example:
-        >>> vertices, faces, uvs = TextureBaker.parametrize_mesh(vertices, faces)  # noqa
-        >>> texture_backer = TextureBaker(vertices, faces, uvs, camera_params)
-        >>> images = get_images_from_grid(args.input_image, image_size)
-        >>> texture = texture_backer.bake_texture(
-        ...     images, texture_size=args.texture_size, mode=args.baker_mode
-        ... )
-        >>> texture = post_process_texture(texture)
-    """
-
-    def __init__(
-        self,
-        vertices: np.ndarray,
-        faces: np.ndarray,
-        uvs: np.ndarray,
-        camera_params: CameraSetting,
-        device: str = "cuda",
-    ) -> None:
-        self.vertices = (
-            torch.tensor(vertices, device=device)
-            if isinstance(vertices, np.ndarray)
-            else vertices.to(device)
-        )
-        self.faces = (
-            torch.tensor(faces.astype(np.int32), device=device)
-            if isinstance(faces, np.ndarray)
-            else faces.to(device)
-        )
-        self.uvs = (
-            torch.tensor(uvs, device=device)
-            if isinstance(uvs, np.ndarray)
-            else uvs.to(device)
-        )
-        self.camera_params = camera_params
-        self.device = device
-
-        camera = init_kal_camera(camera_params)
-        matrix_mv = camera.view_matrix()  # (n_cam 4 4) world2cam
-        matrix_mv = kaolin_to_opencv_view(matrix_mv)
-        matrix_p = (
-            camera.intrinsics.projection_matrix()
-        )  # (n_cam 4 4) cam2pixel
-        self.w2cs = matrix_mv.to(self.device)
-        self.projections = matrix_p.to(self.device)
-
-    @staticmethod
-    def parametrize_mesh(
-        vertices: np.array, faces: np.array
-    ) -> Union[np.array, np.array, np.array]:
-        vmapping, indices, uvs = xatlas.parametrize(vertices, faces)
-
-        vertices = vertices[vmapping]
-        faces = indices
-
-        return vertices, faces, uvs
-
-    def _bake_fast(self, observations, w2cs, projections, texture_size, masks):
-        texture = torch.zeros(
-            (texture_size * texture_size, 3), dtype=torch.float32
-        ).cuda()
-        texture_weights = torch.zeros(
-            (texture_size * texture_size), dtype=torch.float32
-        ).cuda()
-        rastctx = utils3d.torch.RastContext(backend="cuda")
-        for observation, w2c, projection in tqdm(
-            zip(observations, w2cs, projections),
-            total=len(observations),
-            desc="Texture baking (fast)",
-        ):
-            with torch.no_grad():
-                rast = utils3d.torch.rasterize_triangle_faces(
-                    rastctx,
-                    self.vertices[None],
-                    self.faces,
-                    observation.shape[1],
-                    observation.shape[0],
-                    uv=self.uvs[None],
-                    view=w2c,
-                    projection=projection,
-                )
-                uv_map = rast["uv"][0].detach().flip(0)
-                mask = rast["mask"][0].detach().bool() & masks[0]
-
-            # nearest neighbor interpolation
-            uv_map = (uv_map * texture_size).floor().long()
-            obs = observation[mask]
-            uv_map = uv_map[mask]
-            idx = (
-                uv_map[:, 0] + (texture_size - uv_map[:, 1] - 1) * texture_size
-            )
-            texture = texture.scatter_add(
-                0, idx.view(-1, 1).expand(-1, 3), obs
-            )
-            texture_weights = texture_weights.scatter_add(
-                0,
-                idx,
-                torch.ones(
-                    (obs.shape[0]), dtype=torch.float32, device=texture.device
-                ),
-            )
-
-        mask = texture_weights > 0
-        texture[mask] /= texture_weights[mask][:, None]
-        texture = np.clip(
-            texture.reshape(texture_size, texture_size, 3).cpu().numpy() * 255,
-            0,
-            255,
-        ).astype(np.uint8)
-
-        # inpaint
-        mask = (
-            (texture_weights == 0)
-            .cpu()
-            .numpy()
-            .astype(np.uint8)
-            .reshape(texture_size, texture_size)
-        )
-        texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA)
-
-        return texture
-
-    def _bake_opt(
-        self,
-        observations,
-        w2cs,
-        projections,
-        texture_size,
-        lambda_tv,
-        masks,
-        total_steps,
-    ):
-        rastctx = utils3d.torch.RastContext(backend="cuda")
-        observations = [observations.flip(0) for observations in observations]
-        masks = [m.flip(0) for m in masks]
-        _uv = []
-        _uv_dr = []
-        for observation, w2c, projection in tqdm(
-            zip(observations, w2cs, projections),
-            total=len(w2cs),
-        ):
-            with torch.no_grad():
-                rast = utils3d.torch.rasterize_triangle_faces(
-                    rastctx,
-                    self.vertices[None],
-                    self.faces,
-                    observation.shape[1],
-                    observation.shape[0],
-                    uv=self.uvs[None],
-                    view=w2c,
-                    projection=projection,
-                )
-                _uv.append(rast["uv"].detach())
-                _uv_dr.append(rast["uv_dr"].detach())
-
-        texture = torch.nn.Parameter(
-            torch.zeros(
-                (1, texture_size, texture_size, 3), dtype=torch.float32
-            ).cuda()
-        )
-        optimizer = torch.optim.Adam([texture], betas=(0.5, 0.9), lr=1e-2)
-
-        def cosine_anealing(step, total_steps, start_lr, end_lr):
-            return end_lr + 0.5 * (start_lr - end_lr) * (
-                1 + np.cos(np.pi * step / total_steps)
-            )
-
-        def tv_loss(texture):
-            return torch.nn.functional.l1_loss(
-                texture[:, :-1, :, :], texture[:, 1:, :, :]
-            ) + torch.nn.functional.l1_loss(
-                texture[:, :, :-1, :], texture[:, :, 1:, :]
-            )
-
-        with tqdm(total=total_steps, desc="Texture baking") as pbar:
-            for step in range(total_steps):
-                optimizer.zero_grad()
-                selected = np.random.randint(0, len(w2cs))
-                uv, uv_dr, observation, mask = (
-                    _uv[selected],
-                    _uv_dr[selected],
-                    observations[selected],
-                    masks[selected],
-                )
-                render = dr.texture(texture, uv, uv_dr)[0]
-                loss = torch.nn.functional.l1_loss(
-                    render[mask], observation[mask]
-                )
-                if lambda_tv > 0:
-                    loss += lambda_tv * tv_loss(texture)
-                loss.backward()
-                optimizer.step()
-
-                optimizer.param_groups[0]["lr"] = cosine_anealing(
-                    step, total_steps, 1e-2, 1e-5
-                )
-                pbar.set_postfix({"loss": loss.item()})
-                pbar.update()
-        texture = np.clip(
-            texture[0].flip(0).detach().cpu().numpy() * 255, 0, 255
-        ).astype(np.uint8)
-        mask = 1 - utils3d.torch.rasterize_triangle_faces(
-            rastctx,
-            (self.uvs * 2 - 1)[None],
-            self.faces,
-            texture_size,
-            texture_size,
-        )["mask"][0].detach().cpu().numpy().astype(np.uint8)
-        texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA)
-
-        return texture
-
-    def bake_texture(
-        self,
-        images: List[np.array],
-        texture_size: int = 1024,
-        mode: Literal["fast", "opt"] = "opt",
-        lambda_tv: float = 1e-2,
-        opt_step: int = 2000,
-    ):
-        masks = [np.any(img > 0, axis=-1) for img in images]
-        masks = [torch.tensor(m > 0).bool().to(self.device) for m in masks]
-        images = [
-            torch.tensor(obs / 255.0).float().to(self.device) for obs in images
-        ]
-
-        if mode == "fast":
-            return self._bake_fast(
-                images, self.w2cs, self.projections, texture_size, masks
-            )
-        elif mode == "opt":
-            return self._bake_opt(
-                images,
-                self.w2cs,
-                self.projections,
-                texture_size,
-                lambda_tv,
-                masks,
-                opt_step,
-            )
-        else:
-            raise ValueError(f"Unknown mode: {mode}")
-
-
-def kaolin_to_opencv_view(raw_matrix):
-    R_orig = raw_matrix[:, :3, :3]
-    t_orig = raw_matrix[:, :3, 3]
-
-    R_target = torch.zeros_like(R_orig)
-    R_target[:, :, 0] = R_orig[:, :, 2]
-    R_target[:, :, 1] = R_orig[:, :, 0]
-    R_target[:, :, 2] = R_orig[:, :, 1]
-
-    t_target = t_orig
-
-    target_matrix = (
-        torch.eye(4, device=raw_matrix.device)
-        .unsqueeze(0)
-        .repeat(raw_matrix.size(0), 1, 1)
-    )
-    target_matrix[:, :3, :3] = R_target
-    target_matrix[:, :3, 3] = t_target
-
-    return target_matrix
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Render settings")
-
-    parser.add_argument(
-        "--mesh_path",
-        type=str,
-        nargs="+",
-        required=True,
-        help="Paths to the mesh files for rendering.",
-    )
-    parser.add_argument(
-        "--input_image",
-        type=str,
-        nargs="+",
-        required=True,
-        help="Paths to the mesh files for rendering.",
-    )
-    parser.add_argument(
-        "--output_root",
-        type=str,
-        default="./outputs",
-        help="Root directory for output",
-    )
-    parser.add_argument(
-        "--uuid",
-        type=str,
-        nargs="+",
-        default=None,
-        help="uuid for rendering saving.",
-    )
-    parser.add_argument(
-        "--num_images", type=int, default=6, help="Number of images to render."
-    )
-    parser.add_argument(
-        "--elevation",
-        type=float,
-        nargs="+",
-        default=[20.0, -10.0],
-        help="Elevation angles for the camera (default: [20.0, -10.0])",
-    )
-    parser.add_argument(
-        "--distance",
-        type=float,
-        default=5,
-        help="Camera distance (default: 5)",
-    )
-    parser.add_argument(
-        "--resolution_hw",
-        type=int,
-        nargs=2,
-        default=(512, 512),
-        help="Resolution of the output images (default: (512, 512))",
-    )
-    parser.add_argument(
-        "--fov",
-        type=float,
-        default=30,
-        help="Field of view in degrees (default: 30)",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        choices=["cpu", "cuda"],
-        default="cuda",
-        help="Device to run on (default: `cuda`)",
-    )
-    parser.add_argument(
-        "--texture_size",
-        type=int,
-        default=1024,
-        help="Texture size for texture baking (default: 1024)",
-    )
-    parser.add_argument(
-        "--baker_mode",
-        type=str,
-        default="opt",
-        help="Texture baking mode, `fast` or `opt` (default: opt)",
-    )
-    parser.add_argument(
-        "--opt_step",
-        type=int,
-        default=2500,
-        help="Optimization steps for texture baking (default: 2500)",
-    )
-    parser.add_argument(
-        "--mesh_sipmlify_ratio",
-        type=float,
-        default=0.9,
-        help="Mesh simplification ratio (default: 0.9)",
-    )
-    parser.add_argument(
-        "--no_coor_trans",
-        action="store_true",
-        help="Do not transform the asset coordinate system.",
-    )
-    parser.add_argument(
-        "--delight", action="store_true", help="Use delighting model."
-    )
-    parser.add_argument(
-        "--skip_fix_mesh", action="store_true", help="Fix mesh geometry."
-    )
-
-    args = parser.parse_args()
-
-    if args.uuid is None:
-        args.uuid = []
-        for path in args.mesh_path:
-            uuid = os.path.basename(path).split(".")[0]
-            args.uuid.append(uuid)
-
-    return args
-
-
-def entrypoint() -> None:
-    args = parse_args()
-    camera_params = CameraSetting(
-        num_images=args.num_images,
-        elevation=args.elevation,
-        distance=args.distance,
-        resolution_hw=args.resolution_hw,
-        fov=math.radians(args.fov),
-        device=args.device,
-    )
-
-    for mesh_path, uuid, img_path in zip(
-        args.mesh_path, args.uuid, args.input_image
-    ):
-        mesh = trimesh.load(mesh_path)
-        if isinstance(mesh, trimesh.Scene):
-            mesh = mesh.dump(concatenate=True)
-        vertices, scale, center = normalize_vertices_array(mesh.vertices)
-
-        if not args.no_coor_trans:
-            x_rot = torch.Tensor([[1, 0, 0], [0, 0, 1], [0, -1, 0]])
-            z_rot = torch.Tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]])
-            vertices = vertices @ x_rot
-            vertices = vertices @ z_rot
-
-        faces = mesh.faces.cpu().numpy().astype(np.int32)
-        vertices = vertices.cpu().numpy().astype(np.float32)
-
-        if not args.skip_fix_mesh:
-            mesh_fixer = MeshFixer(vertices, faces, args.device)
-            vertices, faces = mesh_fixer(
-                filter_ratio=args.mesh_sipmlify_ratio,
-                max_hole_size=0.04,
-                resolution=1024,
-                num_views=1000,
-                norm_mesh_ratio=0.5,
-            )
-
-        vertices, faces, uvs = TextureBaker.parametrize_mesh(vertices, faces)
-        texture_backer = TextureBaker(
-            vertices,
-            faces,
-            uvs,
-            camera_params,
-        )
-        images = get_images_from_grid(
-            img_path, img_size=camera_params.resolution_hw[0]
-        )
-        if args.delight:
-            delight_model = DelightingModel(
-                model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
-            )
-            delight_images = [delight_model(img) for img in images]
-            images = [np.array(img) for img in delight_images]
-
-        texture = texture_backer.bake_texture(
-            images=[img[..., :3] for img in images],
-            texture_size=args.texture_size,
-            mode=args.baker_mode,
-            opt_step=args.opt_step,
-        )
-        texture = post_process_texture(texture)
-
-        if not args.no_coor_trans:
-            vertices = vertices @ np.linalg.inv(z_rot)
-            vertices = vertices @ np.linalg.inv(x_rot)
-        vertices = vertices / scale
-        vertices = vertices + center
-
-        output_path = os.path.join(args.output_root, f"{uuid}.obj")
-        mesh = save_mesh_with_mtl(vertices, faces, uvs, texture, output_path)
-
-    return
-
-
-if __name__ == "__main__":
-    entrypoint()

asset3d_gen/data/backproject_v2.py

@@ -1,613 +0,0 @@
-import argparse
-import logging
-import math
-import os
-
-import cv2
-import numpy as np
-import nvdiffrast.torch as dr
-import torch
-import torch.nn.functional as F
-import trimesh
-import xatlas
-from PIL import Image
-from asset3d_gen.data.mesh_operator import MeshFixer
-from asset3d_gen.data.utils import (
-    CameraSetting,
-    DiffrastRender,
-    get_images_from_grid,
-    init_kal_camera,
-    normalize_vertices_array,
-    post_process_texture,
-    save_mesh_with_mtl,
-)
-from asset3d_gen.models.delight import DelightingModel
-from asset3d_gen.models.super_resolution import ImageRealESRGAN
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "TextureBacker",
-]
-
-
-def transform_vertices(
-    mtx: torch.Tensor, pos: torch.Tensor, keepdim: bool = False
-) -> torch.Tensor:
-    """Transform 3D vertices using a projection matrix."""
-    t_mtx = torch.as_tensor(mtx, device=pos.device, dtype=pos.dtype)
-    if pos.size(-1) == 3:
-        pos = torch.cat([pos, torch.ones_like(pos[..., :1])], dim=-1)
-
-    result = pos @ t_mtx.T
-
-    return result if keepdim else result.unsqueeze(0)
-
-
-def _bilinear_interpolation_scattering(
-    image_h: int, image_w: int, coords: torch.Tensor, values: torch.Tensor
-) -> torch.Tensor:
-    """Bilinear interpolation scattering for grid-based value accumulation."""
-    device = values.device
-    dtype = values.dtype
-    C = values.shape[-1]
-
-    indices = coords * torch.tensor(
-        [image_h - 1, image_w - 1], dtype=dtype, device=device
-    )
-    i, j = indices.unbind(-1)
-
-    i0, j0 = (
-        indices.floor()
-        .long()
-        .clamp(0, image_h - 2)
-        .clamp(0, image_w - 2)
-        .unbind(-1)
-    )
-    i1, j1 = i0 + 1, j0 + 1
-
-    w_i = i - i0.float()
-    w_j = j - j0.float()
-    weights = torch.stack(
-        [(1 - w_i) * (1 - w_j), (1 - w_i) * w_j, w_i * (1 - w_j), w_i * w_j],
-        dim=1,
-    )
-
-    indices_comb = torch.stack(
-        [
-            torch.stack([i0, j0], dim=1),
-            torch.stack([i0, j1], dim=1),
-            torch.stack([i1, j0], dim=1),
-            torch.stack([i1, j1], dim=1),
-        ],
-        dim=1,
-    )
-
-    grid = torch.zeros(image_h, image_w, C, device=device, dtype=dtype)
-    cnt = torch.zeros(image_h, image_w, 1, device=device, dtype=dtype)
-
-    for k in range(4):
-        idx = indices_comb[:, k]
-        w = weights[:, k].unsqueeze(-1)
-
-        stride = torch.tensor([image_w, 1], device=device, dtype=torch.long)
-        flat_idx = (idx * stride).sum(-1)
-
-        grid.view(-1, C).scatter_add_(
-            0, flat_idx.unsqueeze(-1).expand(-1, C), values * w
-        )
-        cnt.view(-1, 1).scatter_add_(0, flat_idx.unsqueeze(-1), w)
-
-    mask = cnt.squeeze(-1) > 0
-    grid[mask] = grid[mask] / cnt[mask].repeat(1, C)
-
-    return grid
-
-
-def _texture_inpaint_smooth(
-    texture: np.ndarray,
-    mask: np.ndarray,
-    vertices: np.ndarray,
-    faces: np.ndarray,
-    uv_map: np.ndarray,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Perform texture inpainting using vertex-based color propagation."""
-    image_h, image_w, C = texture.shape
-    N = vertices.shape[0]
-
-    # Initialize vertex data structures
-    vtx_mask = np.zeros(N, dtype=np.float32)
-    vtx_colors = np.zeros((N, C), dtype=np.float32)
-    unprocessed = []
-    adjacency = [[] for _ in range(N)]
-
-    # Build adjacency graph and initial color assignment
-    for face_idx in range(faces.shape[0]):
-        for k in range(3):
-            uv_idx_k = faces[face_idx, k]
-            v_idx = faces[face_idx, k]
-
-            # Convert UV to pixel coordinates with boundary clamping
-            u = np.clip(
-                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
-            )
-            v = np.clip(
-                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
-                0,
-                image_h - 1,
-            )
-
-            if mask[v, u]:
-                vtx_mask[v_idx] = 1.0
-                vtx_colors[v_idx] = texture[v, u]
-            elif v_idx not in unprocessed:
-                unprocessed.append(v_idx)
-
-            # Build undirected adjacency graph
-            neighbor = faces[face_idx, (k + 1) % 3]
-            if neighbor not in adjacency[v_idx]:
-                adjacency[v_idx].append(neighbor)
-            if v_idx not in adjacency[neighbor]:
-                adjacency[neighbor].append(v_idx)
-
-    # Color propagation with dynamic stopping
-    remaining_iters, prev_count = 2, 0
-    while remaining_iters > 0:
-        current_unprocessed = []
-
-        for v_idx in unprocessed:
-            valid_neighbors = [n for n in adjacency[v_idx] if vtx_mask[n] > 0]
-            if not valid_neighbors:
-                current_unprocessed.append(v_idx)
-                continue
-
-            # Calculate inverse square distance weights
-            neighbors_pos = vertices[valid_neighbors]
-            dist_sq = np.sum((vertices[v_idx] - neighbors_pos) ** 2, axis=1)
-            weights = 1 / np.maximum(dist_sq, 1e-8)
-
-            vtx_colors[v_idx] = np.average(
-                vtx_colors[valid_neighbors], weights=weights, axis=0
-            )
-            vtx_mask[v_idx] = 1.0
-
-        # Update iteration control
-        if len(current_unprocessed) == prev_count:
-            remaining_iters -= 1
-        else:
-            remaining_iters = min(remaining_iters + 1, 2)
-        prev_count = len(current_unprocessed)
-        unprocessed = current_unprocessed
-
-    # Generate output texture
-    inpainted_texture, updated_mask = texture.copy(), mask.copy()
-    for face_idx in range(faces.shape[0]):
-        for k in range(3):
-            v_idx = faces[face_idx, k]
-            if not vtx_mask[v_idx]:
-                continue
-
-            # UV coordinate conversion
-            uv_idx_k = faces[face_idx, k]
-            u = np.clip(
-                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
-            )
-            v = np.clip(
-                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
-                0,
-                image_h - 1,
-            )
-
-            inpainted_texture[v, u] = vtx_colors[v_idx]
-            updated_mask[v, u] = 255
-
-    return inpainted_texture, updated_mask
-
-
-class TextureBacker:
-    """Texture baking pipeline for multi-view projection and fusion."""
-
-    def __init__(
-        self,
-        camera_params: CameraSetting,
-        view_weights: list[float],
-        render_wh: tuple[int, int] = (2048, 2048),
-        texture_wh: tuple[int, int] = (2048, 2048),
-        bake_angle_thresh: int = 75,
-        mask_thresh: float = 0.5,
-    ):
-        camera = init_kal_camera(camera_params)
-        mv = camera.view_matrix()  # (n 4 4) world2cam
-        p = camera.intrinsics.projection_matrix()
-        # NOTE: add a negative sign at P[0, 2] as the y axis is flipped in `nvdiffrast` output.  # noqa
-        p[:, 1, 1] = -p[:, 1, 1]
-        renderer = DiffrastRender(
-            p_matrix=p,
-            mv_matrix=mv,
-            resolution_hw=camera_params.resolution_hw,
-            context=dr.RasterizeCudaContext(),
-            mask_thresh=mask_thresh,
-            grad_db=False,
-            device=camera_params.device,
-            antialias_mask=True,
-        )
-        self.camera = camera
-        self.renderer = renderer
-        self.view_weights = view_weights
-        self.device = camera_params.device
-        self.render_wh = render_wh
-        self.texture_wh = texture_wh
-
-        self.bake_angle_thresh = bake_angle_thresh
-        self.bake_unreliable_kernel_size = int(
-            (2 / 512) * max(self.render_wh[0], self.render_wh[1])
-        )
-
-    def load_mesh(self, mesh: trimesh.Trimesh) -> None:
-        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-        self.scale, self.center = scale, center
-
-        vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
-        uvs[:, 1] = 1 - uvs[:, 1]
-        mesh.vertices = mesh.vertices[vmapping]
-        mesh.faces = indices
-        mesh.visual.uv = uvs
-
-        self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
-        self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
-        self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
-
-    def get_mesh_np_attrs(
-        self,
-        scale: float = None,
-        center: np.ndarray = None,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        vertices = self.vertices.cpu().numpy()
-        faces = self.faces.cpu().numpy()
-        uv_map = self.uv_map.cpu().numpy()
-        uv_map[:, 1] = 1.0 - uv_map[:, 1]
-
-        if scale is not None:
-            vertices = vertices / scale
-        if center is not None:
-            vertices = vertices + center
-
-        return vertices, faces, uv_map
-
-    def _render_depth_edges(self, depth_image: torch.Tensor) -> torch.Tensor:
-        depth_image_np = depth_image.cpu().numpy()
-        depth_image_np = (depth_image_np * 255).astype(np.uint8)
-        depth_edges = cv2.Canny(depth_image_np, 30, 80)
-        sketch_image = (
-            torch.from_numpy(depth_edges).to(depth_image.device).float() / 255
-        )
-        sketch_image = sketch_image.unsqueeze(-1)
-
-        return sketch_image
-
-    def compute_enhanced_viewnormal(
-        self, mv_mtx: torch.Tensor, vertices: torch.Tensor, faces: torch.Tensor
-    ) -> torch.Tensor:
-        rast, _ = self.renderer.compute_dr_raster(vertices, faces)
-        rendered_view_normals = []
-        for idx in range(len(mv_mtx)):
-            pos_cam = transform_vertices(mv_mtx[idx], vertices, keepdim=True)
-            pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
-            v0, v1, v2 = (pos_cam[faces[:, i]] for i in range(3))
-            face_norm = F.normalize(
-                torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1
-            )
-            vertex_norm = (
-                torch.from_numpy(
-                    trimesh.geometry.mean_vertex_normals(
-                        len(pos_cam), faces.cpu(), face_norm.cpu()
-                    )
-                )
-                .to(vertices.device)
-                .contiguous()
-            )
-            im_base_normals, _ = dr.interpolate(
-                vertex_norm[None, ...].float(),
-                rast[idx : idx + 1],
-                faces.to(torch.int32),
-            )
-            rendered_view_normals.append(im_base_normals)
-
-        rendered_view_normals = torch.cat(rendered_view_normals, dim=0)
-
-        return rendered_view_normals
-
-    def back_project(
-        self, image, vis_mask, depth, normal, uv
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        image = np.array(image)
-        image = torch.as_tensor(image, device=self.device, dtype=torch.float32)
-        if image.ndim == 2:
-            image = image.unsqueeze(-1)
-        image = image / 255
-
-        depth_inv = (1.0 - depth) * vis_mask
-        sketch_image = self._render_depth_edges(depth_inv)
-
-        cos = F.cosine_similarity(
-            torch.tensor([[0, 0, 1]], device=self.device),
-            normal.view(-1, 3),
-        ).view_as(normal[..., :1])
-        cos[cos < np.cos(np.radians(self.bake_angle_thresh))] = 0
-
-        k = self.bake_unreliable_kernel_size * 2 + 1
-        kernel = torch.ones((1, 1, k, k), device=self.device)
-
-        vis_mask = vis_mask.permute(2, 0, 1).unsqueeze(0).float()
-        vis_mask = F.conv2d(
-            1.0 - vis_mask,
-            kernel,
-            padding=k // 2,
-        )
-        vis_mask = 1.0 - (vis_mask > 0).float()
-        vis_mask = vis_mask.squeeze(0).permute(1, 2, 0)
-
-        sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)
-        sketch_image = F.conv2d(sketch_image, kernel, padding=k // 2)
-        sketch_image = (sketch_image > 0).float()
-        sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
-        vis_mask = vis_mask * (sketch_image < 0.5)
-
-        cos[vis_mask == 0] = 0
-        valid_pixels = (vis_mask != 0).view(-1)
-
-        return (
-            self._scatter_texture(uv, image, valid_pixels),
-            self._scatter_texture(uv, cos, valid_pixels),
-        )
-
-    def _scatter_texture(self, uv, data, mask):
-        def __filter_data(data, mask):
-            return data.view(-1, data.shape[-1])[mask]
-
-        return _bilinear_interpolation_scattering(
-            self.texture_wh[1],
-            self.texture_wh[0],
-            __filter_data(uv, mask)[..., [1, 0]],
-            __filter_data(data, mask),
-        )
-
-    @torch.no_grad()
-    def fast_bake_texture(
-        self, textures: list[torch.Tensor], confidence_maps: list[torch.Tensor]
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        channel = textures[0].shape[-1]
-        texture_merge = torch.zeros(self.texture_wh + [channel]).to(
-            self.device
-        )
-        trust_map_merge = torch.zeros(self.texture_wh + [1]).to(self.device)
-        for texture, cos_map in zip(textures, confidence_maps):
-            view_sum = (cos_map > 0).sum()
-            painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()
-            if painted_sum / view_sum > 0.99:
-                continue
-            texture_merge += texture * cos_map
-            trust_map_merge += cos_map
-        texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)
-
-        return texture_merge, trust_map_merge > 1e-8
-
-    def uv_inpaint(
-        self, texture: torch.Tensor, mask: torch.Tensor
-    ) -> np.ndarray:
-        texture_np = texture.cpu().numpy()
-        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
-        vertices, faces, uv_map = self.get_mesh_np_attrs()
-
-        texture_np, mask_np = _texture_inpaint_smooth(
-            texture_np, mask_np, vertices, faces, uv_map
-        )
-        texture_np = texture_np.clip(0, 1)
-        texture_np = cv2.inpaint(
-            (texture_np * 255).astype(np.uint8),
-            255 - mask_np,
-            3,
-            cv2.INPAINT_NS,
-        )
-
-        return texture_np
-
-    def __call__(
-        self,
-        colors: list[Image.Image],
-        mesh: trimesh.Trimesh,
-        output_path: str,
-    ) -> trimesh.Trimesh:
-        self.load_mesh(mesh)
-        rendered_depth, masks = self.renderer.render_depth(
-            self.vertices, self.faces
-        )
-        norm_deps = self.renderer.normalize_map_by_mask(rendered_depth, masks)
-        render_uvs, _ = self.renderer.render_uv(
-            self.vertices, self.faces, self.uv_map
-        )
-        view_normals = self.compute_enhanced_viewnormal(
-            self.renderer.mv_mtx, self.vertices, self.faces
-        )
-
-        textures, weighted_cos_maps = [], []
-        for color, mask, dep, normal, uv, weight in zip(
-            colors,
-            masks,
-            norm_deps,
-            view_normals,
-            render_uvs,
-            self.view_weights,
-        ):
-            texture, cos_map = self.back_project(color, mask, dep, normal, uv)
-            textures.append(texture)
-            weighted_cos_maps.append(weight * (cos_map**4))
-
-        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
-        texture_np = self.uv_inpaint(texture, mask)
-        texture_np = post_process_texture(texture_np)
-        vertices, faces, uv_map = self.get_mesh_np_attrs(
-            self.scale, self.center
-        )
-
-        textured_mesh = save_mesh_with_mtl(
-            vertices, faces, uv_map, texture_np, output_path
-        )
-
-        return textured_mesh
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Backproject texture")
-    parser.add_argument(
-        "--color_path",
-        type=str,
-        help="Multiview color image in 6x512x512 file path",
-    )
-    parser.add_argument(
-        "--mesh_path",
-        type=str,
-        help="Mesh path, .obj, .glb or .ply",
-    )
-    parser.add_argument(
-        "--output_path",
-        type=str,
-        help="Output mesh path with suffix",
-    )
-    parser.add_argument(
-        "--num_images", type=int, default=6, help="Number of images to render."
-    )
-    parser.add_argument(
-        "--elevation",
-        nargs=2,
-        type=float,
-        default=[20.0, -10.0],
-        help="Elevation angles for the camera (default: [20.0, -10.0])",
-    )
-    parser.add_argument(
-        "--distance",
-        type=float,
-        default=5,
-        help="Camera distance (default: 5)",
-    )
-    parser.add_argument(
-        "--resolution_hw",
-        type=int,
-        nargs=2,
-        default=(2048, 2048),
-        help="Resolution of the output images (default: (2048, 2048))",
-    )
-    parser.add_argument(
-        "--fov",
-        type=float,
-        default=30,
-        help="Field of view in degrees (default: 30)",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        choices=["cpu", "cuda"],
-        default="cuda",
-        help="Device to run on (default: `cuda`)",
-    )
-    parser.add_argument(
-        "--skip_fix_mesh", action="store_true", help="Fix mesh geometry."
-    )
-    parser.add_argument(
-        "--texture_wh",
-        nargs=2,
-        type=int,
-        default=[2048, 2048],
-        help="Texture resolution width and height",
-    )
-    parser.add_argument(
-        "--mesh_sipmlify_ratio",
-        type=float,
-        default=0.9,
-        help="Mesh simplification ratio (default: 0.9)",
-    )
-    parser.add_argument(
-        "--delight", action="store_true", help="Use delighting model."
-    )
-    args = parser.parse_args()
-
-    return args
-
-
-def entrypoint(
-    delight_model: DelightingModel = None,
-    imagesr_model: ImageRealESRGAN = None,
-    **kwargs,
-) -> trimesh.Trimesh:
-    args = parse_args()
-    for k, v in kwargs.items():
-        if hasattr(args, k) and v is not None:
-            setattr(args, k, v)
-
-    # Setup camera parameters.
-    camera_params = CameraSetting(
-        num_images=args.num_images,
-        elevation=args.elevation,
-        distance=args.distance,
-        resolution_hw=args.resolution_hw,
-        fov=math.radians(args.fov),
-        device=args.device,
-    )
-    view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]
-
-    color_grid = Image.open(args.color_path)
-    if args.delight:
-        if delight_model is None:
-            delight_model = DelightingModel(
-                model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
-            )
-        save_dir = os.path.dirname(args.output_path)
-        os.makedirs(save_dir, exist_ok=True)
-        color_grid.save(f"{save_dir}/color_grid.png")
-        color_grid = delight_model(color_grid)
-        color_grid.save(f"{save_dir}/color_grid_delight.png")
-
-    multiviews = get_images_from_grid(color_grid, img_size=512)
-
-    # Use RealESRGAN_x4plus for x4 (512->2048) image super resolution.
-    if imagesr_model is None:
-        imagesr_model = ImageRealESRGAN(outscale=4)
-    multiviews = [imagesr_model(img) for img in multiviews]
-    multiviews = [img.convert("RGB") for img in multiviews]
-    mesh = trimesh.load(args.mesh_path)
-    if isinstance(mesh, trimesh.Scene):
-        mesh = mesh.dump(concatenate=True)
-
-    if not args.skip_fix_mesh:
-        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-        mesh_fixer = MeshFixer(mesh.vertices, mesh.faces, args.device)
-        mesh.vertices, mesh.faces = mesh_fixer(
-            filter_ratio=args.mesh_sipmlify_ratio,
-            max_hole_size=0.04,
-            resolution=1024,
-            num_views=1000,
-            norm_mesh_ratio=0.5,
-        )
-        # Restore scale.
-        mesh.vertices = mesh.vertices / scale
-        mesh.vertices = mesh.vertices + center
-
-    # Baking texture to mesh.
-    texture_backer = TextureBacker(
-        camera_params=camera_params,
-        view_weights=view_weights,
-        render_wh=camera_params.resolution_hw,
-        texture_wh=args.texture_wh,
-    )
-
-    textured_mesh = texture_backer(multiviews, mesh, args.output_path)
-
-    return textured_mesh
-
-
-if __name__ == "__main__":
-    entrypoint()

asset3d_gen/data/backup/backproject_v2 copy.py

@@ -1,650 +0,0 @@
-import argparse
-import logging
-import math
-import os
-
-import cv2
-import numpy as np
-import nvdiffrast.torch as dr
-import torch
-import torch.nn.functional as F
-from torchvision.transforms import functional as tF
-
-import trimesh
-import xatlas
-from PIL import Image
-from asset3d_gen.data.mesh_operator import MeshFixer
-from asset3d_gen.data.utils import (
-    CameraSetting,
-    DiffrastRender,
-    get_images_from_grid,
-    init_kal_camera,
-    normalize_vertices_array,
-    post_process_texture,
-    save_mesh_with_mtl,
-)
-from asset3d_gen.models.delight import DelightingModel
-from asset3d_gen.models.super_resolution import ImageRealESRGAN
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "TextureBacker",
-]
-
-
-def transform_vertices(
-    mtx: torch.Tensor, pos: torch.Tensor, keepdim: bool = False
-) -> torch.Tensor:
-    """Transform 3D vertices using a projection matrix."""
-    t_mtx = torch.as_tensor(mtx, device=pos.device, dtype=pos.dtype)
-    if pos.size(-1) == 3:
-        pos = torch.cat([pos, torch.ones_like(pos[..., :1])], dim=-1)
-
-    result = pos @ t_mtx.T
-
-    return result if keepdim else result.unsqueeze(0)
-
-
-def _bilinear_interpolation_scattering(
-    image_h: int, image_w: int, coords: torch.Tensor, values: torch.Tensor
-) -> torch.Tensor:
-    """Bilinear interpolation scattering for grid-based value accumulation."""
-    device = values.device
-    dtype = values.dtype
-    C = values.shape[-1]
-
-    indices = coords * torch.tensor(
-        [image_h - 1, image_w - 1], dtype=dtype, device=device
-    )
-    i, j = indices.unbind(-1)
-
-    i0, j0 = (
-        indices.floor()
-        .long()
-        .clamp(0, image_h - 2)
-        .clamp(0, image_w - 2)
-        .unbind(-1)
-    )
-    i1, j1 = i0 + 1, j0 + 1
-
-    w_i = i - i0.float()
-    w_j = j - j0.float()
-    weights = torch.stack(
-        [(1 - w_i) * (1 - w_j), (1 - w_i) * w_j, w_i * (1 - w_j), w_i * w_j],
-        dim=1,
-    )
-
-    indices_comb = torch.stack(
-        [
-            torch.stack([i0, j0], dim=1),
-            torch.stack([i0, j1], dim=1),
-            torch.stack([i1, j0], dim=1),
-            torch.stack([i1, j1], dim=1),
-        ],
-        dim=1,
-    )
-
-    grid = torch.zeros(image_h, image_w, C, device=device, dtype=dtype)
-    cnt = torch.zeros(image_h, image_w, 1, device=device, dtype=dtype)
-
-    for k in range(4):
-        idx = indices_comb[:, k]
-        w = weights[:, k].unsqueeze(-1)
-
-        stride = torch.tensor([image_w, 1], device=device, dtype=torch.long)
-        flat_idx = (idx * stride).sum(-1)
-
-        grid.view(-1, C).scatter_add_(
-            0, flat_idx.unsqueeze(-1).expand(-1, C), values * w
-        )
-        cnt.view(-1, 1).scatter_add_(0, flat_idx.unsqueeze(-1), w)
-
-    mask = cnt.squeeze(-1) > 0
-    grid[mask] = grid[mask] / cnt[mask].repeat(1, C)
-
-    return grid
-
-
-def _texture_inpaint_smooth(
-    texture: np.ndarray,
-    mask: np.ndarray,
-    vertices: np.ndarray,
-    faces: np.ndarray,
-    uv_map: np.ndarray,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Perform texture inpainting using vertex-based color propagation."""
-    image_h, image_w, C = texture.shape
-    N = vertices.shape[0]
-
-    # Initialize vertex data structures
-    vtx_mask = np.zeros(N, dtype=np.float32)
-    vtx_colors = np.zeros((N, C), dtype=np.float32)
-    unprocessed = []
-    adjacency = [[] for _ in range(N)]
-
-    # Build adjacency graph and initial color assignment
-    for face_idx in range(faces.shape[0]):
-        for k in range(3):
-            uv_idx_k = faces[face_idx, k]
-            v_idx = faces[face_idx, k]
-
-            # Convert UV to pixel coordinates with boundary clamping
-            u = np.clip(
-                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
-            )
-            v = np.clip(
-                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
-                0,
-                image_h - 1,
-            )
-
-            if mask[v, u]:
-                vtx_mask[v_idx] = 1.0
-                vtx_colors[v_idx] = texture[v, u]
-            elif v_idx not in unprocessed:
-                unprocessed.append(v_idx)
-
-            # Build undirected adjacency graph
-            neighbor = faces[face_idx, (k + 1) % 3]
-            if neighbor not in adjacency[v_idx]:
-                adjacency[v_idx].append(neighbor)
-            if v_idx not in adjacency[neighbor]:
-                adjacency[neighbor].append(v_idx)
-
-    # Color propagation with dynamic stopping
-    remaining_iters, prev_count = 2, 0
-    while remaining_iters > 0:
-        current_unprocessed = []
-
-        for v_idx in unprocessed:
-            valid_neighbors = [n for n in adjacency[v_idx] if vtx_mask[n] > 0]
-            if not valid_neighbors:
-                current_unprocessed.append(v_idx)
-                continue
-
-            # Calculate inverse square distance weights
-            neighbors_pos = vertices[valid_neighbors]
-            dist_sq = np.sum((vertices[v_idx] - neighbors_pos) ** 2, axis=1)
-            weights = 1 / np.maximum(dist_sq, 1e-8)
-
-            vtx_colors[v_idx] = np.average(
-                vtx_colors[valid_neighbors], weights=weights, axis=0
-            )
-            vtx_mask[v_idx] = 1.0
-
-        # Update iteration control
-        if len(current_unprocessed) == prev_count:
-            remaining_iters -= 1
-        else:
-            remaining_iters = min(remaining_iters + 1, 2)
-        prev_count = len(current_unprocessed)
-        unprocessed = current_unprocessed
-
-    # Generate output texture
-    inpainted_texture, updated_mask = texture.copy(), mask.copy()
-    for face_idx in range(faces.shape[0]):
-        for k in range(3):
-            v_idx = faces[face_idx, k]
-            if not vtx_mask[v_idx]:
-                continue
-
-            # UV coordinate conversion
-            uv_idx_k = faces[face_idx, k]
-            u = np.clip(
-                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
-            )
-            v = np.clip(
-                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
-                0,
-                image_h - 1,
-            )
-
-            inpainted_texture[v, u] = vtx_colors[v_idx]
-            updated_mask[v, u] = 255
-
-    return inpainted_texture, updated_mask
-
-
-def interp_tensers(tensors: list[torch.Tensor], target_wh: tuple[int, int]) -> list[torch.Tensor]:
-    for idx in range(len(tensors)):
-        tensor = tensors[idx].permute(2, 0, 1)
-        tensor = tF.resize(tensor, target_wh[::-1], antialias=True)
-        tensors[idx] = tensor.permute(1, 2, 0)
-
-    return tensors
-
-
-class TextureBacker:
-    """Texture baking pipeline for multi-view projection and fusion."""
-
-    def __init__(
-        self,
-        camera_params: CameraSetting,
-        view_weights: list[float],
-        render_wh: tuple[int, int] = (2048, 2048),
-        texture_wh: tuple[int, int] = (2048, 2048),
-        bake_angle_thresh: int = 75,
-        mask_thresh: float = 0.5,
-    ):
-        camera = init_kal_camera(camera_params)
-        mv = camera.view_matrix()  # (n 4 4) world2cam
-        p = camera.intrinsics.projection_matrix()
-        # NOTE: add a negative sign at P[0, 2] as the y axis is flipped in `nvdiffrast` output.  # noqa
-        p[:, 1, 1] = -p[:, 1, 1]
-        self.renderer = DiffrastRender(
-            p_matrix=p,
-            mv_matrix=mv,
-            resolution_hw=camera_params.resolution_hw,
-            context=dr.RasterizeCudaContext(),
-            mask_thresh=mask_thresh,
-            grad_db=False,
-            device=camera_params.device,
-            antialias_mask=True,
-        )
-        self.camera = camera
-        self.view_weights = view_weights
-        self.device = camera_params.device
-        self.render_wh = render_wh
-        self.texture_wh = texture_wh
-
-        self.bake_angle_thresh = bake_angle_thresh
-        self.bake_unreliable_kernel_size = int(
-            (2 / 512) * max(self.render_wh[0], self.render_wh[1])
-        )
-
-    def load_mesh(self, mesh: trimesh.Trimesh) -> None:
-        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-        self.scale, self.center = scale, center
-
-        vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
-        uvs[:, 1] = 1 - uvs[:, 1]
-        mesh.vertices = mesh.vertices[vmapping]
-        mesh.faces = indices
-        mesh.visual.uv = uvs
-
-        self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
-        self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
-        self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
-
-    def get_mesh_np_attrs(
-        self,
-        scale: float = None,
-        center: np.ndarray = None,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        vertices = self.vertices.cpu().numpy()
-        faces = self.faces.cpu().numpy()
-        uv_map = self.uv_map.cpu().numpy()
-        uv_map[:, 1] = 1.0 - uv_map[:, 1]
-
-        if scale is not None:
-            vertices = vertices / scale
-        if center is not None:
-            vertices = vertices + center
-
-        return vertices, faces, uv_map
-
-    def _render_depth_edges(self, depth_image: torch.Tensor) -> torch.Tensor:
-        depth_image_np = depth_image.cpu().numpy()
-        depth_image_np = (depth_image_np * 255).astype(np.uint8)
-        depth_edges = cv2.Canny(depth_image_np, 30, 80)
-        sketch_image = (
-            torch.from_numpy(depth_edges).to(depth_image.device).float() / 255
-        )
-        sketch_image = sketch_image.unsqueeze(-1)
-
-        return sketch_image
-
-    def compute_enhanced_viewnormal(
-        self, mv_mtx: torch.Tensor, vertices: torch.Tensor, faces: torch.Tensor
-    ) -> torch.Tensor:
-        rast, _ = self.renderer.compute_dr_raster(vertices, faces)
-        rendered_view_normals = []
-        for idx in range(len(mv_mtx)):
-            pos_cam = transform_vertices(mv_mtx[idx], vertices, keepdim=True)
-            pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
-            v0, v1, v2 = (pos_cam[faces[:, i]] for i in range(3))
-            face_norm = F.normalize(
-                torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1
-            )
-            vertex_norm = (
-                torch.from_numpy(
-                    trimesh.geometry.mean_vertex_normals(
-                        len(pos_cam), faces.cpu(), face_norm.cpu()
-                    )
-                )
-                .to(vertices.device)
-                .contiguous()
-            )
-            im_base_normals, _ = dr.interpolate(
-                vertex_norm[None, ...].float(),
-                rast[idx : idx + 1],
-                faces.to(torch.int32),
-            )
-            rendered_view_normals.append(im_base_normals)
-
-        rendered_view_normals = torch.cat(rendered_view_normals, dim=0)
-
-        return rendered_view_normals
-
-    def back_project(
-        self, image, vis_mask, depth, normal, uv
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        image = np.array(image)
-        image = torch.as_tensor(image, device=self.device, dtype=torch.float32)
-        if image.ndim == 2:
-            image = image.unsqueeze(-1)
-        image = image / 255
-
-        depth_inv = (1.0 - depth) * vis_mask
-        sketch_image = self._render_depth_edges(depth_inv)
-
-        cos = F.cosine_similarity(
-            torch.tensor([[0, 0, 1]], device=self.device),
-            normal.view(-1, 3),
-        ).view_as(normal[..., :1])
-        cos[cos < np.cos(np.radians(self.bake_angle_thresh))] = 0
-
-        k = self.bake_unreliable_kernel_size * 2 + 1
-        kernel = torch.ones((1, 1, k, k), device=self.device)
-
-        vis_mask = vis_mask.permute(2, 0, 1).unsqueeze(0).float()
-        vis_mask = F.conv2d(
-            1.0 - vis_mask,
-            kernel,
-            padding=k // 2,
-        )
-        vis_mask = 1.0 - (vis_mask > 0).float()
-        vis_mask = vis_mask.squeeze(0).permute(1, 2, 0)
-
-        sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)
-        sketch_image = F.conv2d(sketch_image, kernel, padding=k // 2)
-        sketch_image = (sketch_image > 0).float()
-        sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
-        vis_mask = vis_mask * (sketch_image < 0.5)
-
-        cos[vis_mask == 0] = 0
-        valid_pixels = (vis_mask != 0).view(-1)
-
-        return (
-            self._scatter_texture(uv, image, valid_pixels),
-            self._scatter_texture(uv, cos, valid_pixels),
-        )
-
-    def _scatter_texture(self, uv, data, mask):
-        def __filter_data(data, mask):
-            return data.view(-1, data.shape[-1])[mask]
-
-        return _bilinear_interpolation_scattering(
-            self.texture_wh[1],
-            self.texture_wh[0],
-            __filter_data(uv, mask)[..., [1, 0]],
-            __filter_data(data, mask),
-        )
-
-    @torch.no_grad()
-    def fast_bake_texture(
-        self, textures: list[torch.Tensor], confidence_maps: list[torch.Tensor]
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        channel = textures[0].shape[-1]
-        texture_merge = torch.zeros(self.texture_wh + [channel]).to(
-            self.device
-        )
-        trust_map_merge = torch.zeros(self.texture_wh + [1]).to(self.device)
-        for texture, cos_map in zip(textures, confidence_maps):
-            view_sum = (cos_map > 0).sum()
-            painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()
-            if painted_sum / view_sum > 0.99:
-                continue
-            texture_merge += texture * cos_map
-            trust_map_merge += cos_map
-        texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)
-
-        return texture_merge, trust_map_merge > 1e-8
-
-    def uv_inpaint(
-        self, texture: torch.Tensor, mask: torch.Tensor
-    ) -> np.ndarray:
-        texture_np = texture.cpu().numpy()
-        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
-        vertices, faces, uv_map = self.get_mesh_np_attrs()
-
-        texture_np, mask_np = _texture_inpaint_smooth(
-            texture_np, mask_np, vertices, faces, uv_map
-        )
-        texture_np = texture_np.clip(0, 1)
-        texture_np = cv2.inpaint(
-            (texture_np * 255).astype(np.uint8),
-            255 - mask_np,
-            3,
-            cv2.INPAINT_NS,
-        )
-
-        return texture_np
-
-    def __call__(
-        self,
-        colors: list[Image.Image],
-        mesh: trimesh.Trimesh,
-        output_path: str,
-    ) -> trimesh.Trimesh:
-        import time
-        start = time.time()
-        self.load_mesh(mesh)
-        print("load_mesh", time.time() - start)
-        
-        start = time.time()
-        rendered_depth, masks = self.renderer.render_depth(
-            self.vertices, self.faces
-        )
-        norm_deps = self.renderer.normalize_map_by_mask(rendered_depth, masks)
-        render_uvs, _ = self.renderer.render_uv(
-            self.vertices, self.faces, self.uv_map
-        )
-        view_normals = self.compute_enhanced_viewnormal(
-            self.renderer.mv_mtx, self.vertices, self.faces
-        )
-        print("0", time.time() - start)
-
-        textures, weighted_cos_maps = [], []
-        
-        start = time.time()
-        for color, mask, dep, normal, uv, weight in zip(
-            colors,
-            masks,
-            norm_deps,
-            view_normals,
-            render_uvs,
-            self.view_weights,
-        ):
-            mask, dep, normal, uv = interp_tensers([mask, dep, normal, uv], self.render_wh)
-            texture, cos_map = self.back_project(color, mask, dep, normal, uv)
-            textures.append(texture)
-            weighted_cos_maps.append(weight * (cos_map**4))
-        print("1", time.time() - start)
-        start = time.time()
-        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
-        print("2", time.time() - start)
-        start = time.time()
-        texture_np = self.uv_inpaint(texture, mask)
-        print("3", time.time() - start)
-        start = time.time()
-        texture_np = post_process_texture(texture_np)
-        vertices, faces, uv_map = self.get_mesh_np_attrs(
-            self.scale, self.center
-        )
-
-        textured_mesh = save_mesh_with_mtl(
-            vertices, faces, uv_map, texture_np, output_path
-        )
-        print("4", time.time() - start)
-
-        return textured_mesh
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Backproject texture")
-    parser.add_argument(
-        "--color_path",
-        type=str,
-        help="Multiview color image in 6x512x512 file path",
-    )
-    parser.add_argument(
-        "--mesh_path",
-        type=str,
-        help="Mesh path, .obj, .glb or .ply",
-    )
-    parser.add_argument(
-        "--output_path",
-        type=str,
-        help="Output mesh path with suffix",
-    )
-    parser.add_argument(
-        "--num_images", type=int, default=6, help="Number of images to render."
-    )
-    parser.add_argument(
-        "--elevation",
-        nargs=2,
-        type=float,
-        default=[20.0, -10.0],
-        help="Elevation angles for the camera (default: [20.0, -10.0])",
-    )
-    parser.add_argument(
-        "--distance",
-        type=float,
-        default=5,
-        help="Camera distance (default: 5)",
-    )
-    parser.add_argument(
-        "--resolution_hw",
-        type=int,
-        nargs=2,
-        default=(2048, 2048),
-        help="Resolution of the mesh rendering",
-    )
-    parser.add_argument(
-        "--target_hw",
-        type=int,
-        nargs=2,
-        default=(2048, 2048),
-        help="Target rendering images resolution",
-    )
-    parser.add_argument(
-        "--fov",
-        type=float,
-        default=30,
-        help="Field of view in degrees (default: 30)",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        choices=["cpu", "cuda"],
-        default="cuda",
-        help="Device to run on (default: `cuda`)",
-    )
-    parser.add_argument(
-        "--skip_fix_mesh", action="store_true", help="Fix mesh geometry."
-    )
-    parser.add_argument(
-        "--texture_wh",
-        nargs=2,
-        type=int,
-        default=[2048, 2048],
-        help="Texture resolution width and height",
-    )
-    parser.add_argument(
-        "--mesh_sipmlify_ratio",
-        type=float,
-        default=0.9,
-        help="Mesh simplification ratio (default: 0.9)",
-    )
-    parser.add_argument(
-        "--delight", action="store_true", help="Use delighting model."
-    )
-    args = parser.parse_args()
-
-    return args
-
-
-def entrypoint(
-    delight_model: DelightingModel = None,
-    imagesr_model: ImageRealESRGAN = None,
-    **kwargs,
-) -> trimesh.Trimesh:
-    args = parse_args()
-    for k, v in kwargs.items():
-        if hasattr(args, k) and v is not None:
-            setattr(args, k, v)
-
-    # Setup camera parameters.
-    camera_params = CameraSetting(
-        num_images=args.num_images,
-        elevation=args.elevation,
-        distance=args.distance,
-        resolution_hw=args.resolution_hw,
-        fov=math.radians(args.fov),
-        device=args.device,
-    )
-    view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]
-
-    color_grid = Image.open(args.color_path)
-    if args.delight:
-        if delight_model is None:
-            delight_model = DelightingModel(
-                model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
-            )
-        save_dir = os.path.dirname(args.output_path)
-        os.makedirs(save_dir, exist_ok=True)
-        color_grid.save(f"{save_dir}/color_grid.png")
-        color_grid = delight_model(color_grid)
-        color_grid.save(f"{save_dir}/color_grid_delight.png")
-
-    multiviews = get_images_from_grid(color_grid, img_size=512)
-
-    # Use RealESRGAN_x4plus for x4 (512->2048) image super resolution.
-    if imagesr_model is None:
-        imagesr_model = ImageRealESRGAN(outscale=4)
-    multiviews = [imagesr_model(img.convert("RGB")) for img in multiviews]
-    multiviews = [img.resize(args.target_hw[::-1]) for img in multiviews]
-
-    mesh = trimesh.load(args.mesh_path)
-    if isinstance(mesh, trimesh.Scene):
-        mesh = mesh.dump(concatenate=True)
-
-    if not args.skip_fix_mesh:
-        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-        mesh_fixer = MeshFixer(mesh.vertices, mesh.faces, args.device)
-        mesh.vertices, mesh.faces = mesh_fixer(
-            filter_ratio=args.mesh_sipmlify_ratio,
-            max_hole_size=0.04,
-            resolution=1024,
-            num_views=1000,
-            norm_mesh_ratio=0.5,
-        )
-        # Restore scale.
-        mesh.vertices = mesh.vertices / scale
-        mesh.vertices = mesh.vertices + center
-
-    # Baking texture to mesh.
-    import time
-    start = time.time()
-    texture_backer = TextureBacker(
-        camera_params=camera_params,
-        view_weights=view_weights,
-        render_wh=args.target_hw,
-        texture_wh=args.texture_wh,
-    )
-    print(time.time()-start)
-    start = time.time()
-    textured_mesh = texture_backer(multiviews, mesh, args.output_path)
-    print(f"Texture backproject time: {time.time() - start:.2f}s")
-
-    return textured_mesh
-
-
-if __name__ == "__main__":
-    entrypoint()

asset3d_gen/data/backup/backproject_v2.py

@@ -1,700 +0,0 @@
-import logging
-import math
-from typing import Union
-
-import custom_rasterizer as cr
-import cv2
-import numpy as np
-import torch
-import torch.nn.functional as F
-import trimesh
-import xatlas
-from PIL import Image
-from asset3d_gen.data.utils import (
-    get_images_from_file,
-    normalize_vertices_array,
-    post_process_texture,
-    save_mesh_with_mtl,
-)
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-__all__ = ["TextureBacker", "Image_Super_Net", "Image_GANNet"]
-
-
-import math
-import numpy as np
-
-
-def get_perspective_projection(
-    fov: float, aspect_wh: float, near: float = 0.01, far: float = 100
-) -> np.ndarray:
-    """Compute the perspective projection matrix for 3D rendering."""
-    fov_rad = math.radians(fov)
-    tan_half_fov = math.tan(fov_rad / 2.0)
-
-    return np.array(
-        [
-            [1.0 / (tan_half_fov * aspect_wh), 0.0, 0.0, 0.0],
-            [0.0, 1.0 / tan_half_fov, 0.0, 0.0],
-            [
-                0.0,
-                0.0,
-                -(far + near) / (far - near),
-                -(2.0 * far * near) / (far - near),
-            ],
-            [0.0, 0.0, -1.0, 0.0],
-        ],
-        dtype=np.float32,
-    )
-
-
-def transform_vertices(
-    mtx: torch.Tensor, pos: torch.Tensor, keepdim: bool = False
-) -> torch.Tensor:
-    """Transform 3D vertices using a projection matrix."""
-    t_mtx = torch.as_tensor(mtx, device=pos.device, dtype=pos.dtype)
-    if pos.size(-1) == 3:
-        pos = torch.cat([pos, torch.ones_like(pos[..., :1])], dim=-1)
-
-    result = pos @ t_mtx.T
-
-    return result if keepdim else result.unsqueeze(0)
-
-
-def compute_w2c_matrix(
-    elev_deg: float, azim_deg: float, cam_dist: float
-) -> np.ndarray:
-    """Compute w2c 4x4 transformation matrix from spherical coordinates."""
-
-    elev_rad = math.radians(-elev_deg)
-    azim_rad = math.radians(azim_deg)
-
-    sin_elev = math.sin(elev_rad)
-    cos_elev = math.cos(elev_rad)
-    sin_azim = math.sin(azim_rad)
-    cos_azim = math.cos(azim_rad)
-
-    cam_pos = np.array(
-        [
-            cam_dist * cos_elev * cos_azim,
-            cam_dist * cos_elev * sin_azim,
-            cam_dist * sin_elev,
-        ]
-    )
-
-    look_dir = -cam_pos / np.linalg.norm(cam_pos)
-    right_dir = np.cross(look_dir, [0, 0, 1])
-    right_dir /= np.linalg.norm(right_dir)
-    up_dir = np.cross(right_dir, look_dir)
-
-    c2w = np.eye(4)
-    c2w[:3, 0] = right_dir
-    c2w[:3, 1] = up_dir
-    c2w[:3, 2] = -look_dir
-    c2w[:3, 3] = cam_pos
-
-    try:
-        w2c = np.linalg.inv(c2w)
-    except np.linalg.LinAlgError as e:
-        raise ArithmeticError("Failed to invert camera-to-world matrix") from e
-
-    return w2c.astype(np.float32)
-
-
-def _bilinear_interpolation_scattering(
-    image_h: int, image_w: int, coords: torch.Tensor, values: torch.Tensor
-) -> torch.Tensor:
-    """Bilinear interpolation scattering for grid-based value accumulation."""
-    device = values.device
-    dtype = values.dtype
-    C = values.shape[-1]
-
-    indices = coords * torch.tensor(
-        [image_h - 1, image_w - 1], dtype=dtype, device=device
-    )
-    i, j = indices.unbind(-1)
-
-    i0, j0 = (
-        indices.floor()
-        .long()
-        .clamp(0, image_h - 2)
-        .clamp(0, image_w - 2)
-        .unbind(-1)
-    )
-    i1, j1 = i0 + 1, j0 + 1
-
-    w_i = i - i0.float()
-    w_j = j - j0.float()
-    weights = torch.stack(
-        [(1 - w_i) * (1 - w_j), (1 - w_i) * w_j, w_i * (1 - w_j), w_i * w_j],
-        dim=1,
-    )
-
-    indices_comb = torch.stack(
-        [
-            torch.stack([i0, j0], dim=1),
-            torch.stack([i0, j1], dim=1),
-            torch.stack([i1, j0], dim=1),
-            torch.stack([i1, j1], dim=1),
-        ],
-        dim=1,
-    )
-
-    grid = torch.zeros(image_h, image_w, C, device=device, dtype=dtype)
-    cnt = torch.zeros(image_h, image_w, 1, device=device, dtype=dtype)
-
-    for k in range(4):
-        idx = indices_comb[:, k]
-        w = weights[:, k].unsqueeze(-1)
-
-        stride = torch.tensor([image_w, 1], device=device, dtype=torch.long)
-        flat_idx = (idx * stride).sum(-1)
-
-        grid.view(-1, C).scatter_add_(
-            0, flat_idx.unsqueeze(-1).expand(-1, C), values * w
-        )
-        cnt.view(-1, 1).scatter_add_(0, flat_idx.unsqueeze(-1), w)
-
-    mask = cnt.squeeze(-1) > 0
-    grid[mask] = grid[mask] / cnt[mask].repeat(1, C)
-
-    return grid
-
-
-def _texture_inpaint_smooth(
-    texture: np.ndarray,
-    mask: np.ndarray,
-    vertices: np.ndarray,
-    faces: np.ndarray,
-    uv_map: np.ndarray,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Perform texture inpainting using vertex-based color propagation."""
-    image_h, image_w, C = texture.shape
-    N = vertices.shape[0]
-
-    # Initialize vertex data structures
-    vtx_mask = np.zeros(N, dtype=np.float32)
-    vtx_colors = np.zeros((N, C), dtype=np.float32)
-    unprocessed = []
-    adjacency = [[] for _ in range(N)]
-
-    # Build adjacency graph and initial color assignment
-    for face_idx in range(faces.shape[0]):
-        for k in range(3):
-            uv_idx_k = faces[face_idx, k]
-            v_idx = faces[face_idx, k]
-
-            # Convert UV to pixel coordinates with boundary clamping
-            u = np.clip(
-                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
-            )
-            v = np.clip(
-                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
-                0,
-                image_h - 1,
-            )
-
-            if mask[v, u]:
-                vtx_mask[v_idx] = 1.0
-                vtx_colors[v_idx] = texture[v, u]
-            elif v_idx not in unprocessed:
-                unprocessed.append(v_idx)
-
-            # Build undirected adjacency graph
-            neighbor = faces[face_idx, (k + 1) % 3]
-            if neighbor not in adjacency[v_idx]:
-                adjacency[v_idx].append(neighbor)
-            if v_idx not in adjacency[neighbor]:
-                adjacency[neighbor].append(v_idx)
-
-    # Color propagation with dynamic stopping
-    remaining_iters, prev_count = 2, 0
-    while remaining_iters > 0:
-        current_unprocessed = []
-
-        for v_idx in unprocessed:
-            valid_neighbors = [n for n in adjacency[v_idx] if vtx_mask[n] > 0]
-            if not valid_neighbors:
-                current_unprocessed.append(v_idx)
-                continue
-
-            # Calculate inverse square distance weights
-            neighbors_pos = vertices[valid_neighbors]
-            dist_sq = np.sum((vertices[v_idx] - neighbors_pos) ** 2, axis=1)
-            weights = 1 / np.maximum(dist_sq, 1e-8)
-
-            vtx_colors[v_idx] = np.average(
-                vtx_colors[valid_neighbors], weights=weights, axis=0
-            )
-            vtx_mask[v_idx] = 1.0
-
-        # Update iteration control
-        if len(current_unprocessed) == prev_count:
-            remaining_iters -= 1
-        else:
-            remaining_iters = min(remaining_iters + 1, 2)
-        prev_count = len(current_unprocessed)
-        unprocessed = current_unprocessed
-
-    # Generate output texture
-    inpainted_texture, updated_mask = texture.copy(), mask.copy()
-    for face_idx in range(faces.shape[0]):
-        for k in range(3):
-            v_idx = faces[face_idx, k]
-            if not vtx_mask[v_idx]:
-                continue
-
-            # UV coordinate conversion
-            uv_idx_k = faces[face_idx, k]
-            u = np.clip(
-                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
-            )
-            v = np.clip(
-                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
-                0,
-                image_h - 1,
-            )
-
-            inpainted_texture[v, u] = vtx_colors[v_idx]
-            updated_mask[v, u] = 255
-
-    return inpainted_texture, updated_mask
-
-
-class TextureBacker:
-    """Texture baking pipeline for multi-view projection and fusion."""
-
-    def __init__(
-        self,
-        camera_elevs: list[float],
-        camera_azims: list[float],
-        camera_distance: int,
-        camera_fov: float,
-        view_weights: list[float] = None,
-        render_wh: tuple[int, int] = (2048, 2048),
-        texture_wh: tuple[int, int] = (2048, 2048),
-        use_antialias: bool = True,
-        bake_angle_thres: int = 75,
-        device="cuda",
-    ):
-        self.camera_elevs = camera_elevs
-        self.camera_azims = camera_azims
-        self.view_weights = (
-            view_weights
-            if view_weights is not None
-            else [1] * len(camera_elevs)
-        )
-        self.device = device
-        self.render_wh = render_wh
-        self.texture_wh = texture_wh
-
-        self.camera_distance = camera_distance
-        self.use_antialias = use_antialias
-
-        self.bake_angle_thres = bake_angle_thres
-        self.bake_unreliable_kernel_size = int(
-            (2 / 512) * max(self.render_wh[0], self.render_wh[1])
-        )
-
-        self.camera_proj_mat = get_perspective_projection(
-            camera_fov,
-            self.render_wh[1] / self.render_wh[0],
-        )
-        self.cnt = 0
-
-    def rasterize_mesh(
-        self,
-        vertex: torch.Tensor,
-        face: torch.Tensor,
-        resolution: tuple[int, int],
-    ) -> torch.Tensor:
-        vertex = vertex[None] if vertex.ndim == 2 else vertex
-        indices, weights = cr.rasterize(vertex, face, resolution)
-
-        return torch.cat(
-            [weights, indices.unsqueeze(-1).to(weights.dtype)], dim=-1
-        ).unsqueeze(0)
-
-    def raster_interpolate(
-        self, uv: torch.Tensor, rast_out: torch.Tensor, faces: torch.Tensor
-    ) -> torch.Tensor:
-        barycentric = rast_out[0, ..., :-1]
-        findices = rast_out[0, ..., -1]
-        if uv.dim() == 2:
-            uv = uv.unsqueeze(0)
-
-        return cr.interpolate(uv, findices, barycentric, faces)[0]
-
-    def load_mesh(self, mesh_path: str) -> None:
-        mesh = trimesh.load(mesh_path)
-        if isinstance(mesh, trimesh.Scene):
-            mesh = mesh.dump(concatenate=True)
-
-        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-        self.scale, self.center = scale, center
-
-        vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
-        mesh.vertices = mesh.vertices[vmapping]
-        mesh.faces = indices
-        mesh.visual.uv = uvs
-
-        self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
-        self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
-        self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
-
-        # Transformation of coordinate system
-        self.vertices[:, [0, 1]] = -self.vertices[:, [0, 1]]
-        self.vertices[:, [1, 2]] = self.vertices[:, [2, 1]]
-        self.uv_map[:, 1] = 1 - self.uv_map[:, 1]
-
-    def get_mesh_attrs(
-        self,
-        scale: float = None,
-        center: np.ndarray = None,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        vertices = self.vertices.cpu().numpy()
-        faces = self.faces.cpu().numpy()
-        uv_map = self.uv_map.cpu().numpy()
-
-        # Inverse transformation of coordinate system
-        vertices[:, [1, 2]] = vertices[:, [2, 1]]
-        vertices[:, [0, 1]] = -vertices[:, [0, 1]]
-        uv_map[:, 1] = 1.0 - uv_map[:, 1]
-
-        if scale is not None:
-            vertices = vertices / scale
-        if center is not None:
-            vertices = vertices + center
-
-        return vertices, faces, uv_map
-
-    def _render_depth_edges(self, depth_image: torch.Tensor) -> torch.Tensor:
-        depth_image_np = depth_image.cpu().numpy()
-        depth_image_np = (depth_image_np * 255).astype(np.uint8)
-        depth_edges = cv2.Canny(depth_image_np, 30, 80)
-        combined_edges = depth_edges
-        sketch_image = (
-            torch.from_numpy(combined_edges).to(depth_image.device).float()
-            / 255
-        )
-        sketch_image = sketch_image.unsqueeze(-1)
-
-        return sketch_image
-
-    def back_project(
-        self, image: Image.Image, elev: float, azim: float
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        image = torch.as_tensor(image, device=self.device, dtype=torch.float32)
-        if image.ndim == 2:
-            image = image.unsqueeze(-1)
-        image = image / 255.0
-
-        view_mat = compute_w2c_matrix(elev, azim, self.camera_distance)
-        import pdb
-
-        pdb.set_trace()
-        pos_cam = transform_vertices(view_mat, self.vertices, keepdim=True)
-        pos_clip = transform_vertices(self.camera_proj_mat, pos_cam)
-        pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
-
-        v0, v1, v2 = (pos_cam[self.faces[:, i]] for i in range(3))
-        face_norm = F.normalize(torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1)
-        vertex_norm = (
-            torch.from_numpy(
-                trimesh.geometry.mean_vertex_normals(
-                    len(pos_cam), self.faces.cpu(), face_norm.cpu()
-                )
-            )
-            .to(self.device)
-            .contiguous()
-        )
-
-        rast_out = self.rasterize_mesh(pos_clip, self.faces, image.shape[:2])
-        vis_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0]
-
-        interp_data = {
-            "normal": self.raster_interpolate(
-                vertex_norm[None], rast_out, self.faces
-            ),
-            "uv": self.raster_interpolate(
-                self.uv_map[None], rast_out, self.faces
-            ),
-            "depth": self.raster_interpolate(
-                pos_cam[:, 2].reshape(1, -1, 1), rast_out, self.faces
-            ),
-        }
-
-        valid_depth = interp_data["depth"][vis_mask > 0]
-        depth_norm = (interp_data["depth"] - valid_depth.min()) / (
-            valid_depth.max() - valid_depth.min()
-        )
-        # depth_norm[vis_mask <= 0] = 0
-        sketch_image = self._render_depth_edges(depth_norm * vis_mask)
-
-        # ddd = depth_norm * vis_mask
-        # cv2.imwrite(f"v2_depth_d{self.cnt}.png", (ddd.cpu().numpy() * 255).astype(np.uint8))
-
-        cv2.imwrite(
-            f"v2_vis_mask{self.cnt}.png",
-            (vis_mask.cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_normal{self.cnt}.png",
-            (interp_data["normal"].cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_depth{self.cnt}.png",
-            (depth_norm.cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_uv{self.cnt}.png",
-            (interp_data["uv"][..., 0].cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_sketch{self.cnt}.png",
-            (sketch_image.cpu().numpy() * 255).astype(np.uint8),
-        )
-
-        self.cnt += 1
-
-        cos = F.cosine_similarity(
-            torch.tensor([[0, 0, -1]], device=self.device),
-            interp_data["normal"].view(-1, 3),
-        ).view_as(interp_data["normal"][..., :1])
-        cos[cos < np.cos(np.radians(self.bake_angle_thres))] = 0
-
-        cv2.imwrite(
-            f"v2_cos{self.cnt}.png", (cos.cpu().numpy() * 255).astype(np.uint8)
-        )
-
-        k = self.bake_unreliable_kernel_size * 2 + 1
-        kernel = torch.ones((1, 1, k, k), device=self.device)
-
-        vis_mask = vis_mask.permute(2, 0, 1).unsqueeze(0).float()
-        vis_mask = F.conv2d(
-            1.0 - vis_mask,
-            kernel,
-            padding=k // 2,
-        )
-        vis_mask = 1.0 - (vis_mask > 0).float()
-        vis_mask = vis_mask.squeeze(0).permute(1, 2, 0)
-
-        sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)
-        sketch_image = F.conv2d(sketch_image, kernel, padding=k // 2)
-        sketch_image = (sketch_image > 0).float()
-        sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
-        vis_mask = vis_mask * (sketch_image < 0.5)
-
-        cos[vis_mask == 0] = 0
-
-        vis_mask = cv2.imread(
-            f"v3_db_mask{self.cnt}.png", cv2.IMREAD_GRAYSCALE
-        )
-        vis_mask = (
-            torch.from_numpy(vis_mask[..., None]).to(self.device).float() / 255
-        )
-        # cos2 = cv2.imread(f"v3_db_cos{self.cnt}.png", cv2.IMREAD_GRAYSCALE)
-        # cos2 = torch.from_numpy(cos2[..., None]).to(self.device).float() / 255
-        # cos = cos2
-
-        valid_pixels = (vis_mask != 0).view(-1)
-        # import pdb; pdb.set_trace()
-
-        cv2.imwrite(
-            f"v2_db_sketch{self.cnt}.png",
-            (sketch_image.cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_db_uv{self.cnt}.png",
-            (interp_data["uv"][..., 0].cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_db_uv2{self.cnt}.png",
-            (interp_data["uv"][..., 1].cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_db_color{self.cnt}.png",
-            (image.cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_db_cos{self.cnt}.png",
-            (cos.cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v2_db_mask{self.cnt}.png",
-            (vis_mask.cpu().numpy() * 255).astype(np.uint8),
-        )
-        # import pdb; pdb.set_trace()
-        return (
-            self._scatter_texture(interp_data["uv"], image, valid_pixels),
-            self._scatter_texture(interp_data["uv"], cos, valid_pixels),
-        )
-
-    def _scatter_texture(self, uv, data, mask):
-        def __filter_data(data, mask):
-            return data.view(-1, data.shape[-1])[mask]
-
-        return _bilinear_interpolation_scattering(
-            self.texture_wh[1],
-            self.texture_wh[0],
-            __filter_data(uv, mask)[..., [1, 0]],
-            __filter_data(data, mask),
-        )
-
-    @torch.no_grad()
-    def fast_bake_texture(
-        self, textures: list[torch.Tensor], confidence_maps: list[torch.Tensor]
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        channel = textures[0].shape[-1]
-        texture_merge = torch.zeros(self.texture_wh + (channel,)).to(
-            self.device
-        )
-        trust_map_merge = torch.zeros(self.texture_wh + (1,)).to(self.device)
-        for texture, cos_map in zip(textures, confidence_maps):
-            view_sum = (cos_map > 0).sum()
-            painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()
-            if painted_sum / view_sum > 0.99:
-                continue
-            texture_merge += texture * cos_map
-            trust_map_merge += cos_map
-        texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)
-
-        return texture_merge, trust_map_merge > 1e-8
-
-    def uv_inpaint(
-        self, texture: torch.Tensor, mask: torch.Tensor
-    ) -> np.ndarray:
-        texture_np = texture.cpu().numpy()
-        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
-        vertices, faces, uv_map = self.get_mesh_attrs()
-        # import pdb; pdb.set_trace()
-        texture_np, mask_np = _texture_inpaint_smooth(
-            texture_np, mask_np, vertices, faces, uv_map
-        )
-        texture_np = texture_np.clip(0, 1)
-        texture_np = cv2.inpaint(
-            (texture_np * 255).astype(np.uint8),
-            255 - mask_np,
-            3,
-            cv2.INPAINT_NS,
-        )
-
-        return texture_np
-
-    def __call__(
-        self, colors: list[Image.Image], input_mesh: str, output_path: str
-    ) -> trimesh.Trimesh:
-        self.load_mesh(input_mesh)
-
-        textures, weighted_cos_maps = [], []
-        for color, cam_elev, cam_azim, weight in zip(
-            colors, self.camera_elevs, self.camera_azims, self.view_weights
-        ):
-            texture, cos_map = self.back_project(color, cam_elev, cam_azim)
-            cv2.imwrite(
-                f"v2_texture{self.cnt}.png",
-                (texture.cpu().numpy() * 255).astype(np.uint8),
-            )
-            cv2.imwrite(
-                f"v2_texture_cos{self.cnt}.png",
-                (cos_map.cpu().numpy() * 255).astype(np.uint8),
-            )
-            # import pdb; pdb.set_trace()
-            textures.append(texture)
-            weighted_cos_maps.append(weight * (cos_map**4))
-
-        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
-        texture_np = self.uv_inpaint(texture, mask)
-        texture_np = post_process_texture(texture_np)
-        vertices, faces, uvs = self.get_mesh_attrs(self.scale, self.center)
-        # import pdb; pdb.set_trace()
-        cv2.imwrite("v2_texture_np.png", texture_np)
-
-        textured_mesh = save_mesh_with_mtl(
-            vertices, faces, uvs, texture_np, output_path
-        )
-
-        return textured_mesh
-
-
-class Image_Super_Net:
-    def __init__(self, device="cuda"):
-        from diffusers import StableDiffusionUpscalePipeline
-
-        self.up_pipeline_x4 = StableDiffusionUpscalePipeline.from_pretrained(
-            "stabilityai/stable-diffusion-x4-upscaler",
-            torch_dtype=torch.float16,
-        ).to(device)
-        self.up_pipeline_x4.set_progress_bar_config(disable=True)
-
-    def __call__(self, image, prompt=""):
-        with torch.no_grad():
-            upscaled_image = self.up_pipeline_x4(
-                prompt=[prompt],
-                image=image,
-                num_inference_steps=10,
-            ).images[0]
-
-        return upscaled_image
-
-
-class Image_GANNet:
-    def __init__(self, outscale: int):
-        from basicsr.archs.rrdbnet_arch import RRDBNet
-        from realesrgan import RealESRGANer
-
-        self.outscale = outscale
-        model = RRDBNet(
-            num_in_ch=3,
-            num_out_ch=3,
-            num_feat=64,
-            num_block=23,
-            num_grow_ch=32,
-            scale=4,
-        )
-        self.upsampler = RealESRGANer(
-            scale=4,
-            model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/super_resolution/RealESRGAN_x4plus.pth",  # noqa
-            model=model,
-            pre_pad=0,
-            half=True,
-        )
-
-    def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        output, _ = self.upsampler.enhance(image, outscale=self.outscale)
-
-        return Image.fromarray(output)
-
-
-if __name__ == "__main__":
-    device = "cuda"
-    color_path = "outputs/texture_mesh_gen/multi_view/color_sample0.png"
-    mesh_path = "outputs/texture_mesh_gen/texture_mesh/kettle_color.glb"
-    output_path = "robot_test_v2/robot.obj"
-    target_image_size = (2048, 2048)
-
-    super_model = Image_GANNet(outscale=4)
-    multiviews = get_images_from_file(color_path, img_size=512)
-
-    texture_backer = TextureBacker(
-        camera_elevs=[20, 20, 20, -10, -10, -10],
-        camera_azims=[-180, -60, 60, -120, 0, 120],
-        view_weights=[1, 0.2, 0.2, 0.2, 1, 0.2],
-        camera_distance=5,
-        camera_fov=30,
-        render_wh=(2048, 2048),
-        texture_wh=(2048, 2048),
-    )
-
-    multiviews = [super_model(img) for img in multiviews]
-    multiviews = [img.convert("RGB") for img in multiviews]
-    textured_mesh = texture_backer(multiviews, mesh_path, output_path)

asset3d_gen/data/backup/backproject_v3.py

@@ -1,866 +0,0 @@
-import logging
-import math
-from typing import Union
-
-import custom_rasterizer as cr
-import cv2
-import numpy as np
-import torch
-import torch.nn.functional as F
-import trimesh
-import xatlas
-from PIL import Image
-from asset3d_gen.data.utils import (
-    get_images_from_file,
-    normalize_vertices_array,
-    post_process_texture,
-    save_mesh_with_mtl,
-)
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-__all__ = ["TextureBacker", "Image_Super_Net", "Image_GANNet"]
-
-
-import math
-import numpy as np
-
-
-def get_perspective_projection(
-    fov: float, aspect_wh: float, near: float = 0.01, far: float = 100
-) -> np.ndarray:
-    """Compute the perspective projection matrix for 3D rendering."""
-    fov_rad = math.radians(fov)
-    tan_half_fov = math.tan(fov_rad / 2.0)
-
-    return np.array(
-        [
-            [1.0 / (tan_half_fov * aspect_wh), 0.0, 0.0, 0.0],
-            [0.0, 1.0 / tan_half_fov, 0.0, 0.0],
-            [
-                0.0,
-                0.0,
-                -(far + near) / (far - near),
-                -(2.0 * far * near) / (far - near),
-            ],
-            [0.0, 0.0, -1.0, 0.0],
-        ],
-        dtype=np.float32,
-    )
-
-
-def transform_vertices(
-    mtx: torch.Tensor, pos: torch.Tensor, keepdim: bool = False
-) -> torch.Tensor:
-    """Transform 3D vertices using a projection matrix."""
-    t_mtx = torch.as_tensor(mtx, device=pos.device, dtype=pos.dtype)
-    if pos.size(-1) == 3:
-        pos = torch.cat([pos, torch.ones_like(pos[..., :1])], dim=-1)
-
-    result = pos @ t_mtx.T
-
-    return result if keepdim else result.unsqueeze(0)
-
-
-def compute_w2c_matrix(
-    elev_deg: float, azim_deg: float, cam_dist: float
-) -> np.ndarray:
-    """Compute w2c 4x4 transformation matrix from spherical coordinates."""
-
-    elev_rad = math.radians(-elev_deg)
-    azim_rad = math.radians(azim_deg)
-
-    sin_elev = math.sin(elev_rad)
-    cos_elev = math.cos(elev_rad)
-    sin_azim = math.sin(azim_rad)
-    cos_azim = math.cos(azim_rad)
-
-    cam_pos = np.array(
-        [
-            cam_dist * cos_elev * cos_azim,
-            cam_dist * cos_elev * sin_azim,
-            cam_dist * sin_elev,
-        ]
-    )
-
-    look_dir = -cam_pos / np.linalg.norm(cam_pos)
-    right_dir = np.cross(look_dir, [0, 0, 1])
-    right_dir /= np.linalg.norm(right_dir)
-    up_dir = np.cross(right_dir, look_dir)
-
-    c2w = np.eye(4)
-    c2w[:3, 0] = right_dir
-    c2w[:3, 1] = up_dir
-    c2w[:3, 2] = -look_dir
-    c2w[:3, 3] = cam_pos
-
-    try:
-        w2c = np.linalg.inv(c2w)
-    except np.linalg.LinAlgError as e:
-        raise ArithmeticError("Failed to invert camera-to-world matrix") from e
-
-    return w2c.astype(np.float32)
-
-
-def _bilinear_interpolation_scattering(
-    image_h: int, image_w: int, coords: torch.Tensor, values: torch.Tensor
-) -> torch.Tensor:
-    """Bilinear interpolation scattering for grid-based value accumulation."""
-    device = values.device
-    dtype = values.dtype
-    C = values.shape[-1]
-
-    indices = coords * torch.tensor(
-        [image_h - 1, image_w - 1], dtype=dtype, device=device
-    )
-    i, j = indices.unbind(-1)
-
-    i0, j0 = (
-        indices.floor()
-        .long()
-        .clamp(0, image_h - 2)
-        .clamp(0, image_w - 2)
-        .unbind(-1)
-    )
-    i1, j1 = i0 + 1, j0 + 1
-
-    w_i = i - i0.float()
-    w_j = j - j0.float()
-    weights = torch.stack(
-        [(1 - w_i) * (1 - w_j), (1 - w_i) * w_j, w_i * (1 - w_j), w_i * w_j],
-        dim=1,
-    )
-
-    indices_comb = torch.stack(
-        [
-            torch.stack([i0, j0], dim=1),
-            torch.stack([i0, j1], dim=1),
-            torch.stack([i1, j0], dim=1),
-            torch.stack([i1, j1], dim=1),
-        ],
-        dim=1,
-    )
-
-    grid = torch.zeros(image_h, image_w, C, device=device, dtype=dtype)
-    cnt = torch.zeros(image_h, image_w, 1, device=device, dtype=dtype)
-
-    for k in range(4):
-        idx = indices_comb[:, k]
-        w = weights[:, k].unsqueeze(-1)
-
-        stride = torch.tensor([image_w, 1], device=device, dtype=torch.long)
-        flat_idx = (idx * stride).sum(-1)
-
-        grid.view(-1, C).scatter_add_(
-            0, flat_idx.unsqueeze(-1).expand(-1, C), values * w
-        )
-        cnt.view(-1, 1).scatter_add_(0, flat_idx.unsqueeze(-1), w)
-
-    mask = cnt.squeeze(-1) > 0
-    grid[mask] = grid[mask] / cnt[mask].repeat(1, C)
-
-    return grid
-
-
-def _texture_inpaint_smooth(
-    texture: np.ndarray,
-    mask: np.ndarray,
-    vertices: np.ndarray,
-    faces: np.ndarray,
-    uv_map: np.ndarray,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Perform texture inpainting using vertex-based color propagation."""
-    image_h, image_w, C = texture.shape
-    N = vertices.shape[0]
-
-    # Initialize vertex data structures
-    vtx_mask = np.zeros(N, dtype=np.float32)
-    vtx_colors = np.zeros((N, C), dtype=np.float32)
-    unprocessed = []
-    adjacency = [[] for _ in range(N)]
-
-    # Build adjacency graph and initial color assignment
-    for face_idx in range(faces.shape[0]):
-        for k in range(3):
-            uv_idx_k = faces[face_idx, k]
-            v_idx = faces[face_idx, k]
-
-            # Convert UV to pixel coordinates with boundary clamping
-            u = np.clip(
-                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
-            )
-            v = np.clip(
-                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
-                0,
-                image_h - 1,
-            )
-
-            if mask[v, u]:
-                vtx_mask[v_idx] = 1.0
-                vtx_colors[v_idx] = texture[v, u]
-            elif v_idx not in unprocessed:
-                unprocessed.append(v_idx)
-
-            # Build undirected adjacency graph
-            neighbor = faces[face_idx, (k + 1) % 3]
-            if neighbor not in adjacency[v_idx]:
-                adjacency[v_idx].append(neighbor)
-            if v_idx not in adjacency[neighbor]:
-                adjacency[neighbor].append(v_idx)
-
-    # Color propagation with dynamic stopping
-    remaining_iters, prev_count = 2, 0
-    while remaining_iters > 0:
-        current_unprocessed = []
-
-        for v_idx in unprocessed:
-            valid_neighbors = [n for n in adjacency[v_idx] if vtx_mask[n] > 0]
-            if not valid_neighbors:
-                current_unprocessed.append(v_idx)
-                continue
-
-            # Calculate inverse square distance weights
-            neighbors_pos = vertices[valid_neighbors]
-            dist_sq = np.sum((vertices[v_idx] - neighbors_pos) ** 2, axis=1)
-            weights = 1 / np.maximum(dist_sq, 1e-8)
-
-            vtx_colors[v_idx] = np.average(
-                vtx_colors[valid_neighbors], weights=weights, axis=0
-            )
-            vtx_mask[v_idx] = 1.0
-
-        # Update iteration control
-        if len(current_unprocessed) == prev_count:
-            remaining_iters -= 1
-        else:
-            remaining_iters = min(remaining_iters + 1, 2)
-        prev_count = len(current_unprocessed)
-        unprocessed = current_unprocessed
-
-    # Generate output texture
-    inpainted_texture, updated_mask = texture.copy(), mask.copy()
-    for face_idx in range(faces.shape[0]):
-        for k in range(3):
-            v_idx = faces[face_idx, k]
-            if not vtx_mask[v_idx]:
-                continue
-
-            # UV coordinate conversion
-            uv_idx_k = faces[face_idx, k]
-            u = np.clip(
-                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
-            )
-            v = np.clip(
-                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
-                0,
-                image_h - 1,
-            )
-
-            inpainted_texture[v, u] = vtx_colors[v_idx]
-            updated_mask[v, u] = 255
-
-    return inpainted_texture, updated_mask
-
-
-class TextureBacker:
-    """Texture baking pipeline for multi-view projection and fusion."""
-
-    def __init__(
-        self,
-        camera_elevs: list[float],
-        camera_azims: list[float],
-        camera_distance: int,
-        camera_fov: float,
-        view_weights: list[float] = None,
-        render_wh: tuple[int, int] = (2048, 2048),
-        texture_wh: tuple[int, int] = (2048, 2048),
-        use_antialias: bool = True,
-        bake_angle_thresh: int = 75,
-        device="cuda",
-    ):
-        self.camera_elevs = camera_elevs
-        self.camera_azims = camera_azims
-        self.view_weights = (
-            view_weights
-            if view_weights is not None
-            else [1] * len(camera_elevs)
-        )
-        self.device = device
-        self.render_wh = render_wh
-        self.texture_wh = texture_wh
-
-        self.camera_distance = camera_distance
-        self.use_antialias = use_antialias
-
-        self.bake_angle_thresh = bake_angle_thresh
-        self.bake_unreliable_kernel_size = int(
-            (2 / 512) * max(self.render_wh[0], self.render_wh[1])
-        )
-
-        self.camera_proj_mat = get_perspective_projection(
-            camera_fov,
-            self.render_wh[1] / self.render_wh[0],
-        )
-        self.cnt = 0
-
-    def rasterize_mesh(
-        self,
-        vertex: torch.Tensor,
-        face: torch.Tensor,
-        resolution: tuple[int, int],
-    ) -> torch.Tensor:
-        vertex = vertex[None] if vertex.ndim == 2 else vertex
-        indices, weights = cr.rasterize(vertex, face, resolution)
-
-        return torch.cat(
-            [weights, indices.unsqueeze(-1).to(weights.dtype)], dim=-1
-        ).unsqueeze(0)
-
-    def raster_interpolate(
-        self, uv: torch.Tensor, rast_out: torch.Tensor, faces: torch.Tensor
-    ) -> torch.Tensor:
-        barycentric = rast_out[0, ..., :-1]
-        findices = rast_out[0, ..., -1]
-        if uv.dim() == 2:
-            uv = uv.unsqueeze(0)
-
-        return cr.interpolate(uv, findices, barycentric, faces)[0]
-
-    def load_mesh(self, mesh_path: str) -> None:
-        mesh = trimesh.load(mesh_path)
-        if isinstance(mesh, trimesh.Scene):
-            mesh = mesh.dump(concatenate=True)
-
-        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-        self.scale, self.center = scale, center
-
-        vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
-        mesh.vertices = mesh.vertices[vmapping]
-        mesh.faces = indices
-        mesh.visual.uv = uvs
-
-        self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
-        self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
-        self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
-
-        # Transformation of coordinate system
-        self.vertices[:, [0, 1]] = -self.vertices[:, [0, 1]]
-        self.vertices[:, [1, 2]] = self.vertices[:, [2, 1]]
-        self.uv_map[:, 1] = 1 - self.uv_map[:, 1]
-
-    def get_mesh_attrs(
-        self,
-        scale: float = None,
-        center: np.ndarray = None,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        vertices = self.vertices.cpu().numpy()
-        faces = self.faces.cpu().numpy()
-        uv_map = self.uv_map.cpu().numpy()
-
-        if scale is not None:
-            vertices = vertices / scale
-        if center is not None:
-            vertices = vertices + center
-
-        return vertices, faces, uv_map
-
-    def _render_depth_edges(self, depth_image: torch.Tensor) -> torch.Tensor:
-        depth_image_np = depth_image.cpu().numpy()
-        depth_image_np = (depth_image_np * 255).astype(np.uint8)
-        depth_edges = cv2.Canny(depth_image_np, 30, 80)
-        sketch_image = (
-            torch.from_numpy(depth_edges).to(depth_image.device).float() / 255
-        )
-        sketch_image = sketch_image.unsqueeze(-1)
-
-        return sketch_image
-
-    def back_project(
-        self, image: Image.Image, elev: float, azim: float
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        image = torch.as_tensor(image, device=self.device, dtype=torch.float32)
-        if image.ndim == 2:
-            image = image.unsqueeze(-1)
-        image = image / 255.0
-
-        view_mat = compute_w2c_matrix(elev, azim, self.camera_distance)
-        pos_cam = transform_vertices(view_mat, self.vertices, keepdim=True)
-        pos_clip = transform_vertices(self.camera_proj_mat, pos_cam)
-        pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
-
-        v0, v1, v2 = (pos_cam[self.faces[:, i]] for i in range(3))
-        face_norm = F.normalize(torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1)
-        vertex_norm = (
-            torch.from_numpy(
-                trimesh.geometry.mean_vertex_normals(
-                    len(pos_cam), self.faces.cpu(), face_norm.cpu()
-                )
-            )
-            .to(self.device)
-            .contiguous()
-        )
-
-        rast_out = self.rasterize_mesh(pos_clip, self.faces, image.shape[:2])
-        vis_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0]
-
-        interp_data = {
-            "normal": self.raster_interpolate(
-                vertex_norm[None], rast_out, self.faces
-            ),
-            "uv": self.raster_interpolate(
-                self.uv_map[None], rast_out, self.faces
-            ),
-            "depth": self.raster_interpolate(
-                pos_cam[:, 2].reshape(1, -1, 1), rast_out, self.faces
-            ),
-        }
-
-        valid_depth = interp_data["depth"][vis_mask > 0]
-        depth_norm = (interp_data["depth"] - valid_depth.min()) / (
-            valid_depth.max() - valid_depth.min()
-        )
-        depth_norm[vis_mask <= 0] = 0
-        sketch_image = self._render_depth_edges(depth_norm * vis_mask)
-
-        # cv2.imwrite("vis_mask.png", (vis_mask.cpu().numpy() * 255).astype(np.uint8))
-        # cv2.imwrite("normal.png", (interp_data['normal'].cpu().numpy() * 255).astype(np.uint8))
-        # cv2.imwrite("depth.png", (depth_norm.cpu().numpy() * 255).astype(np.uint8))
-        # cv2.imwrite("uv.png", (interp_data['uv'][..., 0].cpu().numpy() * 255).astype(np.uint8))
-        # import pdb; pdb.set_trace()
-
-        cos = F.cosine_similarity(
-            torch.tensor([[0, 0, -1]], device=self.device),
-            interp_data["normal"].view(-1, 3),
-        ).view_as(interp_data["normal"][..., :1])
-        cos[cos < np.cos(np.radians(self.bake_angle_thresh))] = 0
-
-        k = self.bake_unreliable_kernel_size * 2 + 1
-        kernel = torch.ones((1, 1, k, k), device=self.device)
-
-        vis_mask = vis_mask.permute(2, 0, 1).unsqueeze(0).float()
-        vis_mask = F.conv2d(
-            1.0 - vis_mask,
-            kernel,
-            padding=k // 2,
-        )
-        vis_mask = 1.0 - (vis_mask > 0).float()
-        vis_mask = vis_mask.squeeze(0).permute(1, 2, 0)
-
-        sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)
-        sketch_image = F.conv2d(sketch_image, kernel, padding=k // 2)
-        sketch_image = (sketch_image > 0).float()
-        sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
-        vis_mask = vis_mask * (sketch_image < 0.5)
-
-        cos[vis_mask == 0] = 0
-        valid_pixels = (vis_mask != 0).view(-1)
-
-        return (
-            self._scatter_texture(interp_data["uv"], image, valid_pixels),
-            self._scatter_texture(interp_data["uv"], cos, valid_pixels),
-        )
-
-    def back_project2(
-        self, image, vis_mask, depth, normal, uv
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        image = torch.as_tensor(image, device=self.device, dtype=torch.float32)
-        if image.ndim == 2:
-            image = image.unsqueeze(-1)
-        image = image / 255.0
-
-        depth_inv = (1.0 - depth) * vis_mask
-        sketch_image = self._render_depth_edges(depth_inv)
-
-        cv2.imwrite(
-            f"v3_depth_inv{self.cnt}.png",
-            (depth_inv.cpu().numpy() * 255).astype(np.uint8),
-        )
-
-        cos = F.cosine_similarity(
-            torch.tensor([[0, 0, 1]], device=self.device),
-            normal.view(-1, 3),
-        ).view_as(normal[..., :1])
-        cos[cos < np.cos(np.radians(self.bake_angle_thresh))] = 0
-        # import pdb; pdb.set_trace()
-        # cv2.imwrite(f"v3_cos{self.cnt}.png", (cos.cpu().numpy() * 255).astype(np.uint8))
-        # cv2.imwrite(f"v3_sketch{self.cnt}.png", (sketch_image.cpu().numpy() * 255).astype(np.uint8))
-
-        # cos2 = cv2.imread(f"v2_cos{self.cnt+1}.png", cv2.IMREAD_GRAYSCALE)
-        # cos2 = torch.from_numpy(cos2[..., None]).to(self.device).float() / 255
-        # cos = cos2
-
-        self.cnt += 1
-
-        k = self.bake_unreliable_kernel_size * 2 + 1
-        kernel = torch.ones((1, 1, k, k), device=self.device)
-
-        vis_mask = vis_mask.permute(2, 0, 1).unsqueeze(0).float()
-        vis_mask = F.conv2d(
-            1.0 - vis_mask,
-            kernel,
-            padding=k // 2,
-        )
-        vis_mask = 1.0 - (vis_mask > 0).float()
-        vis_mask = vis_mask.squeeze(0).permute(1, 2, 0)
-
-        sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)
-        sketch_image = F.conv2d(sketch_image, kernel, padding=k // 2)
-        sketch_image = (sketch_image > 0).float()
-        sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
-        vis_mask = vis_mask * (sketch_image < 0.5)
-        # import pdb; pdb.set_trace()
-        cv2.imwrite(
-            f"v3_db_sketch{self.cnt}.png",
-            (sketch_image.cpu().numpy() * 255).astype(np.uint8),
-        )
-
-        cos[vis_mask == 0] = 0
-        # import pdb; pdb.set_trace()
-        # vis_mask = cv2.imread(f"v2_db_mask{self.cnt}.png", cv2.IMREAD_GRAYSCALE)
-        # vis_mask = torch.from_numpy(vis_mask[..., None]).to(self.device).float() / 255
-        # cos2 = cv2.imread(f"v2_db_cos{self.cnt}.png", cv2.IMREAD_GRAYSCALE)
-        # cos2 = torch.from_numpy(cos2[..., None]).to(self.device).float() / 255
-        # cos = cos2
-
-        valid_pixels = (vis_mask != 0).view(-1)
-        # import pdb; pdb.set_trace()
-        cv2.imwrite(
-            f"v3_db_uv{self.cnt}.png",
-            (uv[..., 0].cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v3_db_uv2{self.cnt}.png",
-            (uv[..., 1].cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v3_db_color{self.cnt}.png",
-            (image.cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v3_db_cos{self.cnt}.png",
-            (cos.cpu().numpy() * 255).astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v3_db_mask{self.cnt}.png",
-            (vis_mask.cpu().numpy() * 255).astype(np.uint8),
-        )
-
-        return (
-            self._scatter_texture(uv, image, valid_pixels),
-            self._scatter_texture(uv, cos, valid_pixels),
-        )
-
-    def _scatter_texture(self, uv, data, mask):
-        def __filter_data(data, mask):
-            return data.view(-1, data.shape[-1])[mask]
-
-        return _bilinear_interpolation_scattering(
-            self.texture_wh[1],
-            self.texture_wh[0],
-            __filter_data(uv, mask)[..., [1, 0]],
-            __filter_data(data, mask),
-        )
-
-    @torch.no_grad()
-    def fast_bake_texture(
-        self, textures: list[torch.Tensor], confidence_maps: list[torch.Tensor]
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        channel = textures[0].shape[-1]
-        texture_merge = torch.zeros(self.texture_wh + (channel,)).to(
-            self.device
-        )
-        trust_map_merge = torch.zeros(self.texture_wh + (1,)).to(self.device)
-        for texture, cos_map in zip(textures, confidence_maps):
-            view_sum = (cos_map > 0).sum()
-            painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()
-            if painted_sum / view_sum > 0.99:
-                continue
-            texture_merge += texture * cos_map
-            trust_map_merge += cos_map
-        texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)
-
-        return texture_merge, trust_map_merge > 1e-8
-
-    def uv_inpaint(
-        self, texture: torch.Tensor, mask: torch.Tensor
-    ) -> np.ndarray:
-        texture_np = texture.cpu().numpy()
-        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
-        vertices, faces, uv_map = self.get_mesh_attrs()
-        # import pdb; pdb.set_trace()
-        texture_np, mask_np = _texture_inpaint_smooth(
-            texture_np, mask_np, vertices, faces, uv_map
-        )
-        texture_np = texture_np.clip(0, 1)
-        texture_np = cv2.inpaint(
-            (texture_np * 255).astype(np.uint8),
-            255 - mask_np,
-            3,
-            cv2.INPAINT_NS,
-        )
-
-        return texture_np
-
-    def __call__(
-        self, colors: list[Image.Image], input_mesh: str, output_path: str
-    ) -> trimesh.Trimesh:
-        self.load_mesh(input_mesh)
-
-        textures, weighted_cos_maps = [], []
-        for color, cam_elev, cam_azim, weight in zip(
-            colors, self.camera_elevs, self.camera_azims, self.view_weights
-        ):
-            texture, cos_map = self.back_project(color, cam_elev, cam_azim)
-            textures.append(texture)
-            weighted_cos_maps.append(weight * (cos_map**4))
-
-        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
-        texture_np = self.uv_inpaint(texture, mask)
-        texture_np = post_process_texture(texture_np)
-        vertices, faces, uv_map = self.get_mesh_attrs(self.scale, self.center)
-        # import pdb; pdb.set_trace()
-        textured_mesh = save_mesh_with_mtl(
-            vertices, faces, uv_map, texture_np, output_path
-        )
-
-        return textured_mesh
-
-    def forward(
-        self,
-        colors: list[Image.Image],
-        masks,
-        depths,
-        normals,
-        uvs,
-    ) -> trimesh.Trimesh:
-        textures, weighted_cos_maps = [], []
-        for color, mask, depth, normal, uv, weight in zip(
-            colors, masks, depths, normals, uvs, self.view_weights
-        ):
-            texture, cos_map = self.back_project2(
-                color, mask, depth, normal, uv
-            )
-            cv2.imwrite(
-                f"v3_texture{self.cnt}.png",
-                (texture.cpu().numpy() * 255).astype(np.uint8),
-            )
-            cv2.imwrite(
-                f"v3_texture_cos{self.cnt}.png",
-                (cos_map.cpu().numpy() * 255).astype(np.uint8),
-            )
-
-            textures.append(texture)
-            weighted_cos_maps.append(weight * (cos_map**4))
-
-        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
-        texture_np = self.uv_inpaint(texture, mask)
-        texture_np = post_process_texture(texture_np)
-        vertices, faces, uv_map = self.get_mesh_attrs(self.scale, self.center)
-        # import pdb; pdb.set_trace()
-        cv2.imwrite("v3_texture_np.png", texture_np)
-        textured_mesh = save_mesh_with_mtl(
-            vertices, faces, uv_map, texture_np, output_path
-        )
-
-        return textured_mesh
-
-
-class Image_Super_Net:
-    def __init__(self, device="cuda"):
-        from diffusers import StableDiffusionUpscalePipeline
-
-        self.up_pipeline_x4 = StableDiffusionUpscalePipeline.from_pretrained(
-            "stabilityai/stable-diffusion-x4-upscaler",
-            torch_dtype=torch.float16,
-        ).to(device)
-        self.up_pipeline_x4.set_progress_bar_config(disable=True)
-
-    def __call__(self, image, prompt=""):
-        with torch.no_grad():
-            upscaled_image = self.up_pipeline_x4(
-                prompt=[prompt],
-                image=image,
-                num_inference_steps=10,
-            ).images[0]
-
-        return upscaled_image
-
-
-class Image_GANNet:
-    def __init__(self, outscale: int):
-        from basicsr.archs.rrdbnet_arch import RRDBNet
-        from realesrgan import RealESRGANer
-
-        self.outscale = outscale
-        model = RRDBNet(
-            num_in_ch=3,
-            num_out_ch=3,
-            num_feat=64,
-            num_block=23,
-            num_grow_ch=32,
-            scale=4,
-        )
-        self.upsampler = RealESRGANer(
-            scale=4,
-            model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/super_resolution/RealESRGAN_x4plus.pth",  # noqa
-            model=model,
-            pre_pad=0,
-            half=True,
-        )
-
-    def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        output, _ = self.upsampler.enhance(image, outscale=self.outscale)
-
-        return Image.fromarray(output)
-
-
-if __name__ == "__main__":
-    device = "cuda"
-    color_path = "outputs/texture_mesh_gen/multi_view/color_sample0.png"
-    mesh_path = "outputs/texture_mesh_gen/texture_mesh/kettle_color.glb"
-    output_path = "robot_test_v6/robot.obj"
-    target_image_size = (2048, 2048)
-
-    super_model = Image_GANNet(outscale=4)
-    multiviews = get_images_from_file(color_path, img_size=512)
-    multiviews = [super_model(img) for img in multiviews]
-    multiviews = [img.convert("RGB") for img in multiviews]
-
-    from asset3d_gen.data.utils import (
-        CameraSetting,
-        init_kal_camera,
-        DiffrastRender,
-    )
-    import nvdiffrast.torch as dr
-
-    camera_params = CameraSetting(
-        num_images=6,
-        elevation=[20.0, -10.0],
-        distance=5,
-        resolution_hw=(2048, 2048),
-        fov=math.radians(30),
-        device="cuda",
-    )
-    camera = init_kal_camera(camera_params)
-    mv = camera.view_matrix()  # (n 4 4) world2cam
-    p = camera.intrinsics.projection_matrix()
-    # NOTE: add a negative sign at P[0, 2] as the y axis is flipped in `nvdiffrast` output.  # noqa
-    p[:, 1, 1] = -p[:, 1, 1]
-    renderer = DiffrastRender(
-        p_matrix=p,
-        mv_matrix=mv,
-        resolution_hw=camera_params.resolution_hw,
-        context=dr.RasterizeCudaContext(),
-        mask_thresh=0.5,
-        grad_db=False,
-        device=camera_params.device,
-        antialias_mask=True,
-    )
-
-    mesh = trimesh.load(mesh_path)
-    if isinstance(mesh, trimesh.Scene):
-        mesh = mesh.dump(concatenate=True)
-
-    mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-    vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
-    uvs[:, 1] = 1 - uvs[:, 1]
-    mesh.vertices = mesh.vertices[vmapping]
-    mesh.faces = indices
-    mesh.visual.uv = uvs
-
-    vertices = torch.from_numpy(mesh.vertices).to(camera_params.device).float()
-    faces = (
-        torch.from_numpy(mesh.faces).to(camera_params.device).to(torch.int64)
-    )
-    uvs = torch.from_numpy(mesh.visual.uv).to(camera_params.device).float()
-
-    rendered_view_normals = []
-    rast, vertices_clip = renderer.compute_dr_raster(vertices, faces)
-    for idx in range(len(mv)):
-        pos_cam = transform_vertices(mv[idx], vertices, keepdim=True)
-        pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
-        v0, v1, v2 = (pos_cam[faces[:, i]] for i in range(3))
-        face_norm = F.normalize(torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1)
-        vertex_norm = (
-            torch.from_numpy(
-                trimesh.geometry.mean_vertex_normals(
-                    len(pos_cam), faces.cpu(), face_norm.cpu()
-                )
-            )
-            .to(camera_params.device)
-            .contiguous()
-        )
-        im_base_normals, _ = dr.interpolate(
-            vertex_norm[None, ...].float(),
-            rast[idx : idx + 1],
-            faces.to(torch.int32),
-        )
-        rendered_view_normals.append(im_base_normals)
-
-    rendered_view_normals = torch.cat(rendered_view_normals, dim=0)
-
-    rendered_depth, masks = renderer.render_depth(vertices, faces)
-    norm_depths = []
-    for idx in range(len(rendered_depth)):
-        norm_depth = renderer.normalize_map_by_mask(
-            rendered_depth[idx : idx + 1], masks[idx : idx + 1]
-        )
-        norm_depths.append(norm_depth)
-    norm_depths = torch.cat(norm_depths, dim=0)
-    render_uvs, _ = renderer.render_uv(vertices, faces, uvs)
-
-    for index in range(6):
-        cv2.imwrite(
-            f"v3_mask{index}.png",
-            (masks[index] * 255).cpu().numpy().astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v3_normalv2{index}.png",
-            (rendered_view_normals[index] * 255)
-            .cpu()
-            .numpy()
-            .astype(np.uint8)[..., ::-1],
-        )
-        cv2.imwrite(
-            f"v3_depth{index}.png",
-            (norm_depths[index] * 255).cpu().numpy().astype(np.uint8),
-        )
-        cv2.imwrite(
-            f"v3_uv{index}.png",
-            (render_uvs[index, ..., 0] * 255).cpu().numpy().astype(np.uint8),
-        )
-        multiviews[index].save(f"v3_color{index}.png")
-
-    texture_backer = TextureBacker(
-        camera_elevs=[20, 20, 20, -10, -10, -10],
-        camera_azims=[-180, -60, 60, -120, 0, 120],
-        view_weights=[1, 0.2, 0.2, 0.2, 1, 0.2],
-        camera_distance=5,
-        camera_fov=30,
-        render_wh=(2048, 2048),
-        texture_wh=(2048, 2048),
-    )
-    texture_backer.vertices = vertices
-    texture_backer.faces = faces
-    uvs[:, 1] = 1.0 - uvs[:, 1]
-    texture_backer.uv_map = uvs
-    texture_backer.center = center
-    texture_backer.scale = scale
-
-    textured_mesh = texture_backer.forward(
-        multiviews, masks, norm_depths, rendered_view_normals, render_uvs
-    )
-
-    # multiviews = [super_model(img) for img in multiviews]
-    # multiviews = [img.convert("RGB") for img in multiviews]
-    # textured_mesh = texture_backer(multiviews, mesh_path, output_path)

asset3d_gen/data/backup/backprojectv2.py

@@ -1,835 +0,0 @@
-from PIL import Image
-import torch
-import torch.nn.functional as F
-import numpy as np
-import math
-import trimesh
-import cv2
-import xatlas
-from typing import Union
-
-
-def get_perspective_projection_matrix(fovy, aspect_wh, near, far):
-    fovy_rad = math.radians(fovy)
-    return np.array(
-        [
-            [1.0 / (math.tan(fovy_rad / 2.0) * aspect_wh), 0, 0, 0],
-            [0, 1.0 / math.tan(fovy_rad / 2.0), 0, 0],
-            [
-                0,
-                0,
-                -(far + near) / (far - near),
-                -2.0 * far * near / (far - near),
-            ],
-            [0, 0, -1, 0],
-        ]
-    ).astype(np.float32)
-
-
-def load_mesh(mesh):
-    vtx_pos = mesh.vertices if hasattr(mesh, "vertices") else None
-    pos_idx = mesh.faces if hasattr(mesh, "faces") else None
-
-    vtx_uv = mesh.visual.uv if hasattr(mesh.visual, "uv") else None
-    uv_idx = mesh.faces if hasattr(mesh, "faces") else None
-
-    texture_data = None
-
-    return vtx_pos, pos_idx, vtx_uv, uv_idx, texture_data
-
-
-def save_mesh(mesh, texture_data):
-    material = trimesh.visual.texture.SimpleMaterial(
-        image=texture_data, diffuse=(255, 255, 255)
-    )
-    texture_visuals = trimesh.visual.TextureVisuals(
-        uv=mesh.visual.uv, image=texture_data, material=material
-    )
-    mesh.visual = texture_visuals
-    return mesh
-
-
-def transform_pos(mtx, pos, keepdim=False):
-    t_mtx = (
-        torch.from_numpy(mtx).to(pos.device)
-        if isinstance(mtx, np.ndarray)
-        else mtx
-    )
-    if pos.shape[-1] == 3:
-        posw = torch.cat(
-            [pos, torch.ones([pos.shape[0], 1]).to(pos.device)], axis=1
-        )
-    else:
-        posw = pos
-
-    if keepdim:
-        return torch.matmul(posw, t_mtx.t())[...]
-    else:
-        return torch.matmul(posw, t_mtx.t())[None, ...]
-
-
-def get_mv_matrix(elev, azim, camera_distance, center=None):
-    elev = -elev
-
-    elev_rad = math.radians(elev)
-    azim_rad = math.radians(azim)
-
-    camera_position = np.array(
-        [
-            camera_distance * math.cos(elev_rad) * math.cos(azim_rad),
-            camera_distance * math.cos(elev_rad) * math.sin(azim_rad),
-            camera_distance * math.sin(elev_rad),
-        ]
-    )
-
-    if center is None:
-        center = np.array([0, 0, 0])
-    else:
-        center = np.array(center)
-
-    lookat = center - camera_position
-    lookat = lookat / np.linalg.norm(lookat)
-
-    up = np.array([0, 0, 1.0])
-    right = np.cross(lookat, up)
-    right = right / np.linalg.norm(right)
-    up = np.cross(right, lookat)
-    up = up / np.linalg.norm(up)
-
-    c2w = np.concatenate(
-        [np.stack([right, up, -lookat], axis=-1), camera_position[:, None]],
-        axis=-1,
-    )
-
-    w2c = np.zeros((4, 4))
-    w2c[:3, :3] = np.transpose(c2w[:3, :3], (1, 0))
-    w2c[:3, 3:] = -np.matmul(np.transpose(c2w[:3, :3], (1, 0)), c2w[:3, 3:])
-    w2c[3, 3] = 1.0
-
-    return w2c.astype(np.float32)
-
-
-def stride_from_shape(shape):
-    stride = [1]
-    for x in reversed(shape[1:]):
-        stride.append(stride[-1] * x)
-    return list(reversed(stride))
-
-
-def scatter_add_nd_with_count(input, count, indices, values, weights=None):
-    # input: [..., C], D dimension + C channel
-    # count: [..., 1], D dimension
-    # indices: [N, D], long
-    # values: [N, C]
-
-    D = indices.shape[-1]
-    C = input.shape[-1]
-    size = input.shape[:-1]
-    stride = stride_from_shape(size)
-
-    assert len(size) == D
-
-    input = input.view(-1, C)  # [HW, C]
-    count = count.view(-1, 1)
-
-    flatten_indices = (
-        indices * torch.tensor(stride, dtype=torch.long, device=indices.device)
-    ).sum(
-        -1
-    )  # [N]
-
-    if weights is None:
-        weights = torch.ones_like(values[..., :1])
-
-    input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values)
-    count.scatter_add_(0, flatten_indices.unsqueeze(1), weights)
-
-    return input.view(*size, C), count.view(*size, 1)
-
-
-def linear_grid_put_2d(H, W, coords, values, return_count=False):
-    # coords: [N, 2], float in [0, 1]
-    # values: [N, C]
-
-    C = values.shape[-1]
-
-    indices = coords * torch.tensor(
-        [H - 1, W - 1], dtype=torch.float32, device=coords.device
-    )
-    indices_00 = indices.floor().long()  # [N, 2]
-    indices_00[:, 0].clamp_(0, H - 2)
-    indices_00[:, 1].clamp_(0, W - 2)
-    indices_01 = indices_00 + torch.tensor(
-        [0, 1], dtype=torch.long, device=indices.device
-    )
-    indices_10 = indices_00 + torch.tensor(
-        [1, 0], dtype=torch.long, device=indices.device
-    )
-    indices_11 = indices_00 + torch.tensor(
-        [1, 1], dtype=torch.long, device=indices.device
-    )
-
-    h = indices[..., 0] - indices_00[..., 0].float()
-    w = indices[..., 1] - indices_00[..., 1].float()
-    w_00 = (1 - h) * (1 - w)
-    w_01 = (1 - h) * w
-    w_10 = h * (1 - w)
-    w_11 = h * w
-
-    result = torch.zeros(
-        H, W, C, device=values.device, dtype=values.dtype
-    )  # [H, W, C]
-    count = torch.zeros(
-        H, W, 1, device=values.device, dtype=values.dtype
-    )  # [H, W, 1]
-    weights = torch.ones_like(values[..., :1])  # [N, 1]
-
-    result, count = scatter_add_nd_with_count(
-        result,
-        count,
-        indices_00,
-        values * w_00.unsqueeze(1),
-        weights * w_00.unsqueeze(1),
-    )
-    result, count = scatter_add_nd_with_count(
-        result,
-        count,
-        indices_01,
-        values * w_01.unsqueeze(1),
-        weights * w_01.unsqueeze(1),
-    )
-    result, count = scatter_add_nd_with_count(
-        result,
-        count,
-        indices_10,
-        values * w_10.unsqueeze(1),
-        weights * w_10.unsqueeze(1),
-    )
-    result, count = scatter_add_nd_with_count(
-        result,
-        count,
-        indices_11,
-        values * w_11.unsqueeze(1),
-        weights * w_11.unsqueeze(1),
-    )
-
-    if return_count:
-        return result, count
-
-    mask = count.squeeze(-1) > 0
-    result[mask] = result[mask] / count[mask].repeat(1, C)
-
-    return result
-
-
-def meshVerticeInpaint_smooth(texture, mask, vtx_pos, vtx_uv, pos_idx, uv_idx):
-    texture_height, texture_width, texture_channel = texture.shape
-    vtx_num = vtx_pos.shape[0]
-
-    vtx_mask = np.zeros(vtx_num, dtype=np.float32)
-    vtx_color = [
-        np.zeros(texture_channel, dtype=np.float32) for _ in range(vtx_num)
-    ]
-    uncolored_vtxs = []
-    G = [[] for _ in range(vtx_num)]
-
-    for i in range(uv_idx.shape[0]):
-        for k in range(3):
-            vtx_uv_idx = uv_idx[i, k]
-            vtx_idx = pos_idx[i, k]
-            uv_v = int(round(vtx_uv[vtx_uv_idx, 0] * (texture_width - 1)))
-            uv_u = int(
-                round((1.0 - vtx_uv[vtx_uv_idx, 1]) * (texture_height - 1))
-            )
-            if mask[uv_u, uv_v] > 0:
-                vtx_mask[vtx_idx] = 1.0
-                vtx_color[vtx_idx] = texture[uv_u, uv_v]
-            else:
-                uncolored_vtxs.append(vtx_idx)
-            G[pos_idx[i, k]].append(pos_idx[i, (k + 1) % 3])
-
-    smooth_count = 2
-    last_uncolored_vtx_count = 0
-    while smooth_count > 0:
-        uncolored_vtx_count = 0
-        for vtx_idx in uncolored_vtxs:
-            sum_color = np.zeros(texture_channel, dtype=np.float32)
-            total_weight = 0.0
-            vtx_0 = vtx_pos[vtx_idx]
-            for connected_idx in G[vtx_idx]:
-                if vtx_mask[connected_idx] > 0:
-                    vtx1 = vtx_pos[connected_idx]
-                    dist = np.sqrt(np.sum((vtx_0 - vtx1) ** 2))
-                    dist_weight = 1.0 / max(dist, 1e-4)
-                    dist_weight *= dist_weight
-                    sum_color += vtx_color[connected_idx] * dist_weight
-                    total_weight += dist_weight
-            if total_weight > 0:
-                vtx_color[vtx_idx] = sum_color / total_weight
-                vtx_mask[vtx_idx] = 1.0
-            else:
-                uncolored_vtx_count += 1
-
-        if last_uncolored_vtx_count == uncolored_vtx_count:
-            smooth_count -= 1
-        else:
-            smooth_count += 1
-        last_uncolored_vtx_count = uncolored_vtx_count
-
-    new_texture = texture.copy()
-    new_mask = mask.copy()
-    for face_idx in range(uv_idx.shape[0]):
-        for k in range(3):
-            vtx_uv_idx = uv_idx[face_idx, k]
-            vtx_idx = pos_idx[face_idx, k]
-            if vtx_mask[vtx_idx] == 1.0:
-                uv_v = int(round(vtx_uv[vtx_uv_idx, 0] * (texture_width - 1)))
-                uv_u = int(
-                    round((1.0 - vtx_uv[vtx_uv_idx, 1]) * (texture_height - 1))
-                )
-                new_texture[uv_u, uv_v] = vtx_color[vtx_idx]
-                new_mask[uv_u, uv_v] = 255
-
-    return new_texture, new_mask
-
-
-def mesh_uv_wrap(mesh):
-    if isinstance(mesh, trimesh.Scene):
-        mesh = mesh.dump(concatenate=True)
-
-    if len(mesh.faces) > 500000000:
-        raise ValueError(
-            "The mesh has more than 500,000,000 faces, which is not supported."
-        )
-
-    vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
-
-    mesh.vertices = mesh.vertices[vmapping]
-    mesh.faces = indices
-    mesh.visual.uv = uvs
-
-    return mesh
-
-
-class MeshRender:
-    def __init__(
-        self,
-        camera_distance=1.45,
-        default_resolution=1024,
-        texture_size=1024,
-        use_antialias=True,
-        max_mip_level=None,
-        filter_mode="linear",
-        bake_mode="linear",
-        raster_mode="cr",
-        device="cuda",
-    ):
-
-        self.device = device
-
-        self.set_default_render_resolution(default_resolution)
-        self.set_default_texture_resolution(texture_size)
-
-        self.camera_distance = camera_distance
-        self.use_antialias = use_antialias
-        self.max_mip_level = max_mip_level
-        self.filter_mode = filter_mode
-
-        self.bake_angle_thres = 75
-        self.bake_unreliable_kernel_size = int(
-            (2 / 512)
-            * max(self.default_resolution[0], self.default_resolution[1])
-        )
-        self.bake_mode = bake_mode
-
-        self.raster_mode = raster_mode
-        if self.raster_mode == "cr":
-            import custom_rasterizer as cr
-
-            self.raster = cr
-        else:
-            raise f"No raster named {self.raster_mode}"
-
-        fov = 30
-        self.camera_proj_mat = get_perspective_projection_matrix(
-            fov,
-            self.default_resolution[1] / self.default_resolution[0],
-            0.01,
-            100.0,
-        )
-
-    def raster_rasterize(
-        self, pos, tri, resolution, ranges=None, grad_db=True
-    ):
-
-        if self.raster_mode == "cr":
-            rast_out_db = None
-            if pos.dim() == 2:
-                pos = pos.unsqueeze(0)
-            findices, barycentric = self.raster.rasterize(pos, tri, resolution)
-            rast_out = torch.cat((barycentric, findices.unsqueeze(-1)), dim=-1)
-            rast_out = rast_out.unsqueeze(0)
-        else:
-            raise f"No raster named {self.raster_mode}"
-
-        return rast_out, rast_out_db
-
-    def raster_interpolate(
-        self, uv, rast_out, uv_idx, rast_db=None, diff_attrs=None
-    ):
-
-        if self.raster_mode == "cr":
-            textd = None
-            barycentric = rast_out[0, ..., :-1]
-            findices = rast_out[0, ..., -1]
-            if uv.dim() == 2:
-                uv = uv.unsqueeze(0)
-            textc = self.raster.interpolate(uv, findices, barycentric, uv_idx)
-        else:
-            raise f"No raster named {self.raster_mode}"
-
-        return textc, textd
-
-    def load_mesh(
-        self,
-        mesh,
-    ):
-        vtx_pos, pos_idx, vtx_uv, uv_idx, texture_data = load_mesh(mesh)
-        self.mesh_copy = mesh
-        self.set_mesh(
-            vtx_pos,
-            pos_idx,
-            vtx_uv=vtx_uv,
-            uv_idx=uv_idx,
-        )
-        if texture_data is not None:
-            self.set_texture(texture_data)
-
-    def save_mesh(self):
-        texture_data = self.get_texture()
-        texture_data = Image.fromarray((texture_data * 255).astype(np.uint8))
-        return save_mesh(self.mesh_copy, texture_data)
-
-    def set_mesh(
-        self,
-        vtx_pos,
-        pos_idx,
-        vtx_uv=None,
-        uv_idx=None,
-    ):
-
-        self.vtx_pos = torch.from_numpy(vtx_pos).to(self.device).float()
-        self.pos_idx = torch.from_numpy(pos_idx).to(self.device).to(torch.int)
-        if (vtx_uv is not None) and (uv_idx is not None):
-            self.vtx_uv = torch.from_numpy(vtx_uv).to(self.device).float()
-            self.uv_idx = (
-                torch.from_numpy(uv_idx).to(self.device).to(torch.int)
-            )
-        else:
-            self.vtx_uv = None
-            self.uv_idx = None
-
-        self.vtx_pos[:, [0, 1]] = -self.vtx_pos[:, [0, 1]]
-        self.vtx_pos[:, [1, 2]] = self.vtx_pos[:, [2, 1]]
-        if (vtx_uv is not None) and (uv_idx is not None):
-            self.vtx_uv[:, 1] = 1.0 - self.vtx_uv[:, 1]
-
-    def set_texture(self, tex):
-        if isinstance(tex, np.ndarray):
-            tex = Image.fromarray((tex * 255).astype(np.uint8))
-        elif isinstance(tex, torch.Tensor):
-            tex = tex.cpu().numpy()
-            tex = Image.fromarray((tex * 255).astype(np.uint8))
-
-        tex = tex.resize(self.texture_size).convert("RGB")
-        tex = np.array(tex) / 255.0
-        self.tex = torch.from_numpy(tex).to(self.device)
-        self.tex = self.tex.float()
-
-    def set_default_render_resolution(self, default_resolution):
-        if isinstance(default_resolution, int):
-            default_resolution = (default_resolution, default_resolution)
-        self.default_resolution = default_resolution
-
-    def set_default_texture_resolution(self, texture_size):
-        if isinstance(texture_size, int):
-            texture_size = (texture_size, texture_size)
-        self.texture_size = texture_size
-
-    def get_mesh(self):
-        vtx_pos = self.vtx_pos.cpu().numpy()
-        pos_idx = self.pos_idx.cpu().numpy()
-        vtx_uv = self.vtx_uv.cpu().numpy()
-        uv_idx = self.uv_idx.cpu().numpy()
-
-        # 坐标变换的逆变换
-        vtx_pos[:, [1, 2]] = vtx_pos[:, [2, 1]]
-        vtx_pos[:, [0, 1]] = -vtx_pos[:, [0, 1]]
-
-        vtx_uv[:, 1] = 1.0 - vtx_uv[:, 1]
-        return vtx_pos, pos_idx, vtx_uv, uv_idx
-
-    def get_texture(self):
-        return self.tex.cpu().numpy()
-
-    def render_sketch_from_depth(self, depth_image):
-        depth_image_np = depth_image.cpu().numpy()
-        depth_image_np = (depth_image_np * 255).astype(np.uint8)
-        depth_edges = cv2.Canny(depth_image_np, 30, 80)
-        combined_edges = depth_edges
-        sketch_image = (
-            torch.from_numpy(combined_edges).to(depth_image.device).float()
-            / 255.0
-        )
-        sketch_image = sketch_image.unsqueeze(-1)
-        return sketch_image
-
-    def back_project(
-        self, image, elev, azim, camera_distance=None, center=None, method=None
-    ):
-        if isinstance(image, Image.Image):
-            image = torch.tensor(np.array(image) / 255.0)
-        elif isinstance(image, np.ndarray):
-            image = torch.tensor(image)
-        if image.dim() == 2:
-            image = image.unsqueeze(-1)
-        image = image.float().to(self.device)
-        resolution = image.shape[:2]
-        channel = image.shape[-1]
-        texture = torch.zeros(self.texture_size + (channel,)).to(self.device)
-        cos_map = torch.zeros(self.texture_size + (1,)).to(self.device)
-
-        proj = self.camera_proj_mat
-        r_mv = get_mv_matrix(
-            elev=elev,
-            azim=azim,
-            camera_distance=(
-                self.camera_distance
-                if camera_distance is None
-                else camera_distance
-            ),
-            center=center,
-        )
-        pos_camera = transform_pos(r_mv, self.vtx_pos, keepdim=True)
-        pos_clip = transform_pos(proj, pos_camera)
-        pos_camera = pos_camera[:, :3] / pos_camera[:, 3:4]
-        v0 = pos_camera[self.pos_idx[:, 0], :]
-        v1 = pos_camera[self.pos_idx[:, 1], :]
-        v2 = pos_camera[self.pos_idx[:, 2], :]
-        face_normals = F.normalize(
-            torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1
-        )
-        vertex_normals = trimesh.geometry.mean_vertex_normals(
-            vertex_count=self.vtx_pos.shape[0],
-            faces=self.pos_idx.cpu(),
-            face_normals=face_normals.cpu(),
-        )
-        vertex_normals = (
-            torch.from_numpy(vertex_normals)
-            .float()
-            .to(self.device)
-            .contiguous()
-        )
-        tex_depth = pos_camera[:, 2].reshape(1, -1, 1).contiguous()
-        rast_out, rast_out_db = self.raster_rasterize(
-            pos_clip, self.pos_idx, resolution=resolution
-        )
-        visible_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0, ...]
-
-        normal, _ = self.raster_interpolate(
-            vertex_normals[None, ...], rast_out, self.pos_idx
-        )
-        normal = normal[0, ...]
-
-        uv, _ = self.raster_interpolate(
-            self.vtx_uv[None, ...], rast_out, self.uv_idx
-        )
-        depth, _ = self.raster_interpolate(tex_depth, rast_out, self.pos_idx)
-        depth = depth[0, ...]
-
-        depth_max, depth_min = (
-            depth[visible_mask > 0].max(),
-            depth[visible_mask > 0].min(),
-        )
-        depth_normalized = (depth - depth_min) / (depth_max - depth_min)
-        depth_image = depth_normalized * visible_mask  # Mask out background.
-
-        sketch_image = self.render_sketch_from_depth(depth_image)
-
-        cv2.imwrite("d_depth.png", depth_image.cpu().numpy() * 255)
-        cv2.imwrite("d_normal.png", normal.cpu().numpy() * 255)
-        cv2.imwrite(
-            "d_image.png", image.cpu().numpy()[..., :3][..., ::-1] * 255
-        )
-        cv2.imwrite("d_sketch_image.png", sketch_image.cpu().numpy() * 255)
-        cv2.imwrite("d_uv1.png", uv.cpu().numpy()[0, ..., 0] * 255)
-        cv2.imwrite("d_uv2.png", uv.cpu().numpy()[0, ..., 1] * 255)
-        # p uv[0,...,0].mean(axis=0)
-        # import pdb; pdb.set_trace()
-
-        # depth_image = None
-        # normal = None
-        # image = None
-
-        sketch_image = self.render_sketch_from_depth(depth_image)
-        channel = image.shape[-1]
-
-        lookat = torch.tensor([[0, 0, -1]], device=self.device)
-        cos_image = torch.nn.functional.cosine_similarity(
-            lookat, normal.view(-1, 3)
-        )
-        cos_image = cos_image.view(normal.shape[0], normal.shape[1], 1)
-
-        cos_thres = np.cos(self.bake_angle_thres / 180 * np.pi)
-        cos_image[cos_image < cos_thres] = 0
-
-        # shrink
-        kernel_size = self.bake_unreliable_kernel_size * 2 + 1
-        kernel = torch.ones(
-            (1, 1, kernel_size, kernel_size), dtype=torch.float32
-        ).to(sketch_image.device)
-
-        visible_mask = visible_mask.permute(2, 0, 1).unsqueeze(0).float()
-        visible_mask = F.conv2d(
-            1.0 - visible_mask, kernel, padding=kernel_size // 2
-        )
-        visible_mask = 1.0 - (visible_mask > 0).float()  # 二值化
-        visible_mask = visible_mask.squeeze(0).permute(1, 2, 0)
-
-        sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)
-        sketch_image = F.conv2d(sketch_image, kernel, padding=kernel_size // 2)
-        sketch_image = (sketch_image > 0).float()  # 二值化
-        sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
-        visible_mask = visible_mask * (sketch_image < 0.5)
-
-        cos_image[visible_mask == 0] = 0
-        proj_mask = (visible_mask != 0).view(-1)
-        uv = uv.squeeze(0).contiguous().view(-1, 2)[proj_mask]
-        image = image.squeeze(0).contiguous().view(-1, channel)[proj_mask]
-        cos_image = cos_image.contiguous().view(-1, 1)[proj_mask]
-        sketch_image = sketch_image.contiguous().view(-1, 1)[proj_mask]
-        import pdb
-
-        pdb.set_trace()
-        texture = linear_grid_put_2d(
-            self.texture_size[1], self.texture_size[0], uv[..., [1, 0]], image
-        )
-        cos_map = linear_grid_put_2d(
-            self.texture_size[1],
-            self.texture_size[0],
-            uv[..., [1, 0]],
-            cos_image,
-        )
-        boundary_map = linear_grid_put_2d(
-            self.texture_size[1],
-            self.texture_size[0],
-            uv[..., [1, 0]],
-            sketch_image,
-        )
-
-        return texture, cos_map, boundary_map
-
-    @torch.no_grad()
-    def fast_bake_texture(self, textures, cos_maps):
-
-        channel = textures[0].shape[-1]
-        texture_merge = torch.zeros(self.texture_size + (channel,)).to(
-            self.device
-        )
-        trust_map_merge = torch.zeros(self.texture_size + (1,)).to(self.device)
-        for texture, cos_map in zip(textures, cos_maps):
-            view_sum = (cos_map > 0).sum()
-            painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()
-            if painted_sum / view_sum > 0.99:
-                continue
-            texture_merge += texture * cos_map
-            trust_map_merge += cos_map
-        texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)
-
-        return texture_merge, trust_map_merge > 1e-8
-
-    def uv_inpaint(self, texture, mask):
-
-        if isinstance(texture, torch.Tensor):
-            texture_np = texture.cpu().numpy()
-        elif isinstance(texture, np.ndarray):
-            texture_np = texture
-        elif isinstance(texture, Image.Image):
-            texture_np = np.array(texture) / 255.0
-
-        vtx_pos, pos_idx, vtx_uv, uv_idx = self.get_mesh()
-
-        texture_np, mask = meshVerticeInpaint_smooth(
-            texture_np, mask, vtx_pos, vtx_uv, pos_idx, uv_idx
-        )
-
-        texture_np = cv2.inpaint(
-            (texture_np * 255).astype(np.uint8), 255 - mask, 3, cv2.INPAINT_NS
-        )
-
-        return texture_np
-
-
-def get_images_from_file(img_path: str, img_size: int) -> list[np.array]:
-    input_image = Image.open(img_path)
-    view_images = np.array(input_image)
-    view_images = np.concatenate(
-        [view_images[:img_size, ...], view_images[img_size:, ...]], axis=1
-    )
-    images = np.split(view_images, view_images.shape[1] // img_size, axis=1)
-
-    return images
-
-
-def bake_from_multiview(
-    render, views, camera_elevs, camera_azims, view_weights, method="fast"
-):
-    project_textures, project_weighted_cos_maps = [], []
-    project_boundary_maps = []
-    for view, camera_elev, camera_azim, weight in zip(
-        views, camera_elevs, camera_azims, view_weights
-    ):
-        project_texture, project_cos_map, project_boundary_map = (
-            render.back_project(view, camera_elev, camera_azim)
-        )
-        project_cos_map = weight * (project_cos_map**4)
-        project_textures.append(project_texture)
-        project_weighted_cos_maps.append(project_cos_map)
-        project_boundary_maps.append(project_boundary_map)
-
-    if method == "fast":
-        texture, ori_trust_map = render.fast_bake_texture(
-            project_textures, project_weighted_cos_maps
-        )
-    else:
-        raise f"no method {method}"
-
-    return texture, ori_trust_map > 1e-8
-
-
-def post_process(texture: np.ndarray, iter: int = 2) -> np.ndarray:
-    for _ in range(iter):
-        texture = cv2.fastNlMeansDenoisingColored(texture, None, 11, 11, 9, 25)
-        texture = cv2.bilateralFilter(
-            texture, d=7, sigmaColor=80, sigmaSpace=80
-        )
-
-    return texture
-
-
-class Image_Super_Net:
-    def __init__(self, device="cuda"):
-        from diffusers import StableDiffusionUpscalePipeline
-
-        self.up_pipeline_x4 = StableDiffusionUpscalePipeline.from_pretrained(
-            "stabilityai/stable-diffusion-x4-upscaler",
-            torch_dtype=torch.float16,
-        ).to(device)
-        self.up_pipeline_x4.set_progress_bar_config(disable=True)
-
-    def __call__(self, image, prompt=""):
-        with torch.no_grad():
-            upscaled_image = self.up_pipeline_x4(
-                prompt=[prompt],
-                image=image,
-                num_inference_steps=10,
-            ).images[0]
-
-        return upscaled_image
-
-
-class Image_GANNet:
-    def __init__(self, outscale: int):
-        from realesrgan import RealESRGANer
-        from basicsr.archs.rrdbnet_arch import RRDBNet
-
-        self.outscale = outscale
-        model = RRDBNet(
-            num_in_ch=3,
-            num_out_ch=3,
-            num_feat=64,
-            num_block=23,
-            num_grow_ch=32,
-            scale=4,
-        )
-        self.upsampler = RealESRGANer(
-            scale=4,
-            model_path="/home/users/xinjie.wang/xinjie/Real-ESRGAN/weights/RealESRGAN_x4plus.pth",
-            model=model,
-            pre_pad=0,
-            half=True,
-        )
-
-    def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        output, _ = self.upsampler.enhance(image, outscale=self.outscale)
-
-        return Image.fromarray(output)
-
-
-if __name__ == "__main__":
-    device = "cuda"
-
-    # super_model = Image_Super_Net(device)
-    super_model = Image_GANNet(outscale=4)
-
-    selected_camera_elevs = [20, 20, 20, -10, -10, -10]
-    selected_camera_azims = [-180, -60, 60, -120, 0, 120]
-    selected_view_weights = [1, 0.2, 0.2, 0.2, 1, 0.2]
-    # selected_view_weights = [1, 0.1, 0.5, 0.1, 0.05, 0.05]
-
-    multiviews = get_images_from_file(
-        "scripts/apps/texture_sessions/mfq4e7u4ko/multi_view/color_sample1.png",
-        512,
-    )
-    target_image_size = (2048, 2048)
-
-    render = MeshRender(
-        camera_distance=5,
-        default_resolution=2048,
-        texture_size=2048,
-    )
-
-    mesh = trimesh.load("scripts/apps/assets/example_texture/meshes/robot.obj")
-    from asset3d_gen.data.utils import normalize_vertices_array
-
-    mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-    mesh = mesh_uv_wrap(mesh)
-    render.load_mesh(mesh)
-
-    # multiviews = [Image.fromarray(img) for img in multiviews]
-    # multiviews = [Image.fromarray(img).convert("RGB") for img in multiviews]
-    # for idx, img in enumerate(multiviews):
-    #     img.save(f"robot/raw/res_{idx}.png")
-
-    multiviews = [super_model(img) for img in multiviews]
-    multiviews = [img.convert("RGB") for img in multiviews]
-    for idx, img in enumerate(multiviews):
-        img.save(f"robot/super_gan_res_{idx}.png")
-
-    texture, mask = bake_from_multiview(
-        render,
-        multiviews,
-        selected_camera_elevs,
-        selected_camera_azims,
-        selected_view_weights,
-    )
-
-    texture_np = (texture.cpu().numpy() * 255).astype(np.uint8)[..., :3][
-        ..., ::-1
-    ]
-    cv2.imwrite("robot/raw_texture.png", texture_np)
-    print("texture done.")
-
-    mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
-    texture_np = render.uv_inpaint(texture, mask_np)
-    cv2.imwrite("robot/inpaint_texture.png", texture_np[..., ::-1])
-    # texture_np = post_process(texture_np, 2)
-    # cv2.imwrite("robot/inpaint_conv_texture.png", texture_np[..., ::-1])
-    print("inpaint done.")
-
-    texture = torch.tensor(texture_np / 255).float().to(texture.device)
-    render.set_texture(texture)
-    textured_mesh = render.save_mesh()
-    _ = textured_mesh.export("robot/robot.obj")

asset3d_gen/data/backup/gpt_qwen.py

@@ -1,70 +0,0 @@
-import torch
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
-from qwen_vl_utils import process_vision_info
-import os 
-os.environ["https_proxy"] = "10.9.0.31:8838"
-
-
-# # default: Load the model on the available device(s)
-# model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-#     "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
-# )
-
-# We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2.5-VL-7B-Instruct",
-    torch_dtype=torch.bfloat16,
-    attn_implementation="flash_attention_2",
-    device_map="auto",
-)
-
-
-# default processer
-processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
-
-# The default range for the number of visual tokens per image in the model is 4-16384.
-# You can set min_pixels and max_pixels according to your needs, such as a token range of 256-1280, to balance performance and cost.
-# min_pixels = 256*28*28
-# max_pixels = 1280*28*28
-# processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
-
-messages = [
-    {
-        "role": "user",
-        "content": [
-            {
-                "type": "image",
-                "image": "outputs/text2image/demo_objects/bed/sample_0.jpg",
-            },
-            {
-                "type": "image",
-                "image": "outputs/imageto3d/v2/cups/sample_69/URDF_sample_69/qa_renders/image_color/003.png",
-            },
-            {"type": "text", "text": "Describe the secend image."},
-        ],
-    }
-]
-
-# Preparation for inference
-text = processor.apply_chat_template(
-    messages, tokenize=False, add_generation_prompt=True
-)
-image_inputs, video_inputs = process_vision_info(messages)
-inputs = processor(
-    text=[text],
-    images=image_inputs,
-    videos=video_inputs,
-    padding=True,
-    return_tensors="pt",
-)
-inputs = inputs.to("cuda")
-
-# Inference: Generation of the output
-generated_ids = model.generate(**inputs, max_new_tokens=128)
-generated_ids_trimmed = [
-    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-]
-output_text = processor.batch_decode(
-    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-)
-print(output_text)

asset3d_gen/data/backup/quat.py

@@ -1,49 +0,0 @@
-import numpy as np
-
-def quaternion_rotation_x_counterclockwise(angle_degrees):
-    angle_radians = np.radians(angle_degrees)
-    w = np.cos(angle_radians / 2)
-    x = np.sin(angle_radians / 2)
-    y, z = 0.0, 0.0
-    return np.array([x, y, z, w]).round(4).tolist()
-
-
-def quaternion_rotation_y_counterclockwise(angle_degrees):
-    angle_radians = np.radians(angle_degrees)
-    w = np.cos(angle_radians / 2)
-    y = np.sin(angle_radians / 2)
-    x, z = 0.0, 0.0
-    return np.array([x, y, z, w]).round(4).tolist()
-
-
-def quaternion_rotation_z_counterclockwise(angle_degrees):
-    angle_radians = np.radians(angle_degrees)
-    w = np.cos(angle_radians / 2)
-    z = np.sin(angle_radians / 2)
-    x, y = 0.0, 0.0
-    return np.array([x, y, z, w]).round(4).tolist()
-
-
-def quaternion_multiply(q1, q2):
-    x1, y1, z1, w1 = q1
-    x2, y2, z2, w2 = q2
-    w = w1*w2 - x1*x2 - y1*y2 - z1*z2
-    x = w1*x2 + x1*w2 + y1*z2 - z1*y2
-    y = w1*y2 - x1*z2 + y1*w2 + z1*x2
-    z = w1*z2 + x1*y2 - y1*x2 + z1*w2
-    return np.array([w, x, y, z])
-
-    
-    
-angle = 180
-
-print(f"X轴逆时针旋转{angle}度: {quaternion_rotation_x_counterclockwise(angle)}")
-print(f"Y轴逆时针旋转{angle}度: {quaternion_rotation_y_counterclockwise(angle)}")
-print(f"Z轴逆时针旋转{angle}度: {quaternion_rotation_z_counterclockwise(angle)}")
-
-
-q_1 = np.array([1.0, 0.0, 0.0, 0.0]) 
-q_2 = np.array([0.0, 0.0, 1.0, 0.0]) 
-
-q_total = quaternion_multiply(q_2, q_1)
-print(q_total.round(4).tolist())

asset3d_gen/data/datasets.py

@@ -1,239 +0,0 @@
-import json
-import logging
-import os
-import random
-from typing import Any, Callable, Dict, List, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from PIL import Image
-from torch import nn
-from torch.utils.data import Dataset
-from torchvision import transforms
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "Asset3dGenDataset",
-]
-
-
-class Asset3dGenDataset(Dataset):
-    def __init__(
-        self,
-        index_file: str,
-        target_hw: Tuple[int, int],
-        transform: Callable = None,
-        control_transform: Callable = None,
-        max_train_samples: int = None,
-        sub_idxs: List[List[int]] = None,
-        seed: int = 79,
-    ) -> None:
-        if not os.path.exists(index_file):
-            raise FileNotFoundError(f"{index_file} index_file not found.")
-
-        self.index_file = index_file
-        self.target_hw = target_hw
-        self.transform = transform
-        self.control_transform = control_transform
-        self.max_train_samples = max_train_samples
-        self.meta_info = self.prepare_data_index(index_file)
-        self.data_list = sorted(self.meta_info.keys())
-        self.sub_idxs = sub_idxs  # sub_idxs [[0,1,2], [3,4,5], [...], ...]
-        self.image_num = 6  # hardcode temp.
-        random.seed(seed)
-        logger.info(f"Trainset: {len(self)} asset3d instances.")
-
-    def __len__(self) -> int:
-        return len(self.meta_info)
-
-    def prepare_data_index(self, index_file: str) -> Dict[str, Any]:
-        with open(index_file, "r") as fin:
-            meta_info = json.load(fin)
-
-        meta_info_filtered = dict()
-        for idx, uid in enumerate(meta_info):
-            if "status" not in meta_info[uid]:
-                continue
-            if meta_info[uid]["status"] != "success":
-                continue
-            if self.max_train_samples and idx >= self.max_train_samples:
-                break
-
-            meta_info_filtered[uid] = meta_info[uid]
-
-        logger.info(
-            f"Load {len(meta_info)} assets, keep {len(meta_info_filtered)} valids."  # noqa
-        )
-
-        return meta_info_filtered
-
-    def fetch_sample_images(
-        self,
-        uid: str,
-        attrs: List[str],
-        sub_index: int = None,
-        transform: Callable = None,
-    ) -> torch.Tensor:
-        sample = self.meta_info[uid]
-        images = []
-        for attr in attrs:
-            item = sample[attr]
-            if sub_index is not None:
-                item = item[sub_index]
-            mode = "L" if attr == "image_mask" else "RGB"
-            image = Image.open(item).convert(mode)
-            if transform is not None:
-                image = transform(image)
-                if len(image.shape) == 2:
-                    image = image[..., None]
-            images.append(image)
-
-        images = torch.cat(images, dim=0)
-
-        return images
-
-    def fetch_sample_grid_images(
-        self,
-        uid: str,
-        attrs: List[str],
-        sub_idxs: List[List[int]],
-        transform: Callable = None,
-    ) -> torch.Tensor:
-        assert transform is not None
-
-        grid_image = []
-        for row_idxs in sub_idxs:
-            row_image = []
-            for row_idx in row_idxs:
-                image = self.fetch_sample_images(
-                    uid, attrs, row_idx, transform
-                )
-                row_image.append(image)
-            row_image = torch.cat(row_image, dim=2)  # (c h w)
-            grid_image.append(row_image)
-
-        grid_image = torch.cat(grid_image, dim=1)
-
-        return grid_image
-
-    def compute_text_embeddings(
-        self, embed_path: str, original_size: Tuple[int, int]
-    ) -> Dict[str, nn.Module]:
-        data_dict = torch.load(embed_path)
-        prompt_embeds = data_dict["prompt_embeds"][0]
-        add_text_embeds = data_dict["pooled_prompt_embeds"][0]
-
-        # Need changed if random crop, set as crop_top_left [y1, x1], center crop as [0, 0].  # noqa
-        crops_coords_top_left = (0, 0)
-        add_time_ids = list(
-            original_size + crops_coords_top_left + self.target_hw
-        )
-        add_time_ids = torch.tensor([add_time_ids])
-        # add_time_ids = add_time_ids.repeat((len(add_text_embeds), 1))
-
-        unet_added_cond_kwargs = {
-            "text_embeds": add_text_embeds,
-            "time_ids": add_time_ids,
-        }
-
-        return {"prompt_embeds": prompt_embeds, **unet_added_cond_kwargs}
-
-    def visualize_item(
-        self,
-        control: torch.Tensor,
-        color: torch.Tensor,
-        save_dir: str = None,
-    ) -> List[Image.Image]:
-        to_pil = transforms.ToPILImage()
-
-        color = (color + 1) / 2
-        color_pil = to_pil(color)
-        normal_pil = to_pil(control[0:3])
-        position_pil = to_pil(control[3:6])
-        mask_pil = to_pil(control[6:])
-
-        if save_dir is not None:
-            os.makedirs(save_dir, exist_ok=True)
-            color_pil.save(f"{save_dir}/rgb.jpg")
-            normal_pil.save(f"{save_dir}/normal.jpg")
-            position_pil.save(f"{save_dir}/position.jpg")
-            mask_pil.save(f"{save_dir}/mask.jpg")
-            logger.info(f"Visualization in {save_dir}")
-
-        return normal_pil, position_pil, mask_pil, color_pil
-
-    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
-        uid = self.data_list[index]
-
-        sub_idxs = self.sub_idxs
-        if sub_idxs is None:
-            sub_idxs = [[random.randint(0, self.image_num - 1)]]
-
-        input_image = self.fetch_sample_grid_images(
-            uid,
-            attrs=["image_view_normal", "image_position", "image_mask"],
-            sub_idxs=sub_idxs,
-            transform=self.control_transform,
-        )
-        assert input_image.shape[1:] == self.target_hw
-
-        output_image = self.fetch_sample_grid_images(
-            uid,
-            attrs=["image_color"],
-            sub_idxs=sub_idxs,
-            transform=self.transform,
-        )
-
-        sample = self.meta_info[uid]
-        text_feats = self.compute_text_embeddings(
-            sample["text_feat"], tuple(sample["image_hw"])
-        )
-
-        data = dict(
-            pixel_values=output_image,
-            conditioning_pixel_values=input_image,
-            prompt_embeds=text_feats["prompt_embeds"],
-            text_embeds=text_feats["text_embeds"],
-            time_ids=text_feats["time_ids"],
-        )
-
-        return data
-
-
-if __name__ == "__main__":
-    index_file = "/horizon-bucket/robot_lab/users/xinjie.wang/datasets/objaverse/v1.0/statistics_1.0_gobjaverse_filter/view6s_v4/meta_ac2e0ddea8909db26d102c8465b5bcb2.json"  # noqa
-    target_hw = (512, 512)
-    transform_list = [
-        transforms.Resize(
-            target_hw, interpolation=transforms.InterpolationMode.BILINEAR
-        ),
-        transforms.CenterCrop(target_hw),
-        transforms.ToTensor(),
-        transforms.Normalize([0.5], [0.5]),
-    ]
-    image_transform = transforms.Compose(transform_list)
-    control_transform = transforms.Compose(transform_list[:-1])
-
-    sub_idxs = [[0, 1, 2], [3, 4, 5]]  # None
-    if sub_idxs is not None:
-        target_hw = (
-            target_hw[0] * len(sub_idxs),
-            target_hw[1] * len(sub_idxs[0]),
-        )
-
-    dataset = Asset3dGenDataset(
-        index_file,
-        target_hw,
-        image_transform,
-        control_transform,
-        sub_idxs=sub_idxs,
-    )
-    data = dataset[0]
-    dataset.visualize_item(
-        data["conditioning_pixel_values"], data["pixel_values"], save_dir="./"
-    )

asset3d_gen/data/differentiable_render.py

@@ -1,520 +0,0 @@
-import argparse
-import json
-import logging
-import math
-import os
-from collections import defaultdict
-from typing import List, Union
-
-import cv2
-import imageio
-import numpy as np
-import nvdiffrast.torch as dr
-import torch
-from PIL import Image
-from tqdm import tqdm
-from asset3d_gen.data.utils import (
-    CameraSetting,
-    DiffrastRender,
-    RenderItems,
-    as_list,
-    calc_vertex_normals,
-    import_kaolin_mesh,
-    init_kal_camera,
-    normalize_vertices_array,
-    render_pbr,
-    save_images,
-)
-
-os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
-os.environ["TORCH_EXTENSIONS_DIR"] = os.path.expanduser(
-    "~/.cache/torch_extensions"
-)
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-def create_gif_from_images(images, output_path, fps=10):
-    pil_images = []
-    for image in images:
-        image = image.clip(min=0, max=1)
-        image = (255.0 * image).astype(np.uint8)
-        image = Image.fromarray(image, mode="RGBA")
-        pil_images.append(image.convert("RGB"))
-
-    duration = 1000 // fps
-    pil_images[0].save(
-        output_path,
-        save_all=True,
-        append_images=pil_images[1:],
-        duration=duration,
-        loop=0,
-    )
-
-    logger.info(f"GIF saved to {output_path}")
-
-
-def create_mp4_from_images(images, output_path, fps=10, prompt=None):
-    font = cv2.FONT_HERSHEY_SIMPLEX  # 字体样式
-    font_scale = 0.5  # 字体大小
-    font_thickness = 1  # 字体粗细
-    color = (255, 255, 255)  # 文字颜色（白色）
-    position = (20, 25)  # 左上角坐标 (x, y)
-
-    with imageio.get_writer(output_path, fps=fps) as writer:
-        for image in images:
-            image = image.clip(min=0, max=1)
-            image = (255.0 * image).astype(np.uint8)
-            image = image[..., :3]
-            if prompt is not None:
-                cv2.putText(
-                    image,
-                    prompt,
-                    position,
-                    font,
-                    font_scale,
-                    color,
-                    font_thickness,
-                )
-
-            writer.append_data(image)
-
-    logger.info(f"MP4 video saved to {output_path}")
-
-
-class ImageRender(object):
-    def __init__(
-        self,
-        render_items: list[RenderItems],
-        camera_params: CameraSetting,
-        recompute_vtx_normal: bool = True,
-        device: str = "cuda",
-        with_mtl: bool = False,
-        gen_color_gif: bool = False,
-        gen_color_mp4: bool = False,
-        gen_viewnormal_mp4: bool = False,
-        gen_glonormal_mp4: bool = False,
-        no_index_file: bool = False,
-        light_factor: float = 1.0,
-    ) -> None:
-        camera_params.device = device
-        camera = init_kal_camera(camera_params)
-        self.camera = camera
-
-        # Setup MVP matrix and renderer.
-        mv = camera.view_matrix()  # (n 4 4) world2cam
-        p = camera.intrinsics.projection_matrix()
-        # NOTE: add a negative sign at P[0, 2] as the y axis is flipped in `nvdiffrast` output.  # noqa
-        p[:, 1, 1] = -p[:, 1, 1]
-        # mvp = torch.bmm(p, mv) # camera.view_projection_matrix()
-        self.mv = mv
-        self.p = p
-
-        renderer = DiffrastRender(
-            p_matrix=p,
-            mv_matrix=mv,
-            resolution_hw=camera_params.resolution_hw,
-            context=dr.RasterizeCudaContext(),
-            mask_thresh=0.5,
-            grad_db=False,
-            device=camera_params.device,
-            antialias_mask=True,
-        )
-        self.renderer = renderer
-        self.recompute_vtx_normal = recompute_vtx_normal
-        self.render_items = render_items
-        self.device = device
-        self.with_mtl = with_mtl
-        self.gen_color_gif = gen_color_gif
-        self.gen_color_mp4 = gen_color_mp4
-        self.gen_viewnormal_mp4 = gen_viewnormal_mp4
-        self.gen_glonormal_mp4 = gen_glonormal_mp4
-        self.light_factor = light_factor
-        self.no_index_file = no_index_file
-
-    def render_mesh(
-        self,
-        mesh_path: Union[str, List[str]],
-        output_root: str,
-        uuid: Union[str, List[str]] = None,
-        prompts: List[str] = None,
-    ) -> None:
-        mesh_path = as_list(mesh_path)
-        if uuid is None:
-            uuid = [os.path.basename(p).split(".")[0] for p in mesh_path]
-        uuid = as_list(uuid)
-        assert len(mesh_path) == len(uuid)
-        os.makedirs(output_root, exist_ok=True)
-
-        meta_info = dict()
-        for idx, (path, uid) in tqdm(
-            enumerate(zip(mesh_path, uuid)), total=len(mesh_path)
-        ):
-            output_dir = os.path.join(output_root, uid)
-            os.makedirs(output_dir, exist_ok=True)
-            prompt = prompts[idx] if prompts else None
-            data_dict = self(path, output_dir, prompt)
-            meta_info[uid] = data_dict
-
-        if self.no_index_file:
-            return
-
-        index_file = os.path.join(output_root, "index.json")
-        with open(index_file, "w") as fout:
-            json.dump(meta_info, fout)
-
-        logger.info(f"Rendering meta info logged in {index_file}")
-
-    def __call__(
-        self, mesh_path: str, output_dir: str, prompt: str = None
-    ) -> dict[str, str]:
-        try:
-            mesh = import_kaolin_mesh(mesh_path, self.with_mtl)
-        except Exception as e:
-            logger.error(f"[ERROR MESH LOAD]: {e}, skip {mesh_path}")
-            return
-
-        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
-        if self.recompute_vtx_normal:
-            mesh.vertex_normals = calc_vertex_normals(
-                mesh.vertices, mesh.faces
-            )
-
-        mesh = mesh.to(self.device)
-        vertices, faces, vertex_normals = (
-            mesh.vertices,
-            mesh.faces,
-            mesh.vertex_normals,
-        )
-
-        # Perform rendering.
-        data_dict = defaultdict(list)
-        if RenderItems.ALPHA.value in self.render_items:
-            masks, _ = self.renderer.render_rast_alpha(vertices, faces)
-            render_paths = save_images(
-                masks, f"{output_dir}/{RenderItems.ALPHA}"
-            )
-            data_dict[RenderItems.ALPHA.value] = render_paths
-
-        if RenderItems.GLOBAL_NORMAL.value in self.render_items:
-            rendered_normals, masks = self.renderer.render_global_normal(
-                vertices, faces, vertex_normals
-            )
-            if self.gen_glonormal_mp4:
-                if isinstance(rendered_normals, torch.Tensor):
-                    rendered_normals = rendered_normals.detach().cpu().numpy()
-                create_mp4_from_images(
-                    rendered_normals,
-                    output_path=f"{output_dir}/normal.mp4",
-                    fps=15,
-                    prompt=prompt,
-                )
-            else:
-                render_paths = save_images(
-                    rendered_normals,
-                    f"{output_dir}/{RenderItems.GLOBAL_NORMAL}",
-                    cvt_color=cv2.COLOR_BGR2RGB,
-                )
-                data_dict[RenderItems.GLOBAL_NORMAL.value] = render_paths
-
-            if RenderItems.VIEW_NORMAL.value in self.render_items:
-                assert (
-                    RenderItems.GLOBAL_NORMAL in self.render_items
-                ), f"Must render global normal firstly, got render_items: {self.render_items}."  # noqa
-                rendered_view_normals = self.renderer.transform_normal(
-                    rendered_normals, self.mv, masks, to_view=True
-                )
-                # rendered_inv_view_normals = renderer.transform_normal(rendered_view_normals, torch.linalg.inv(mv), masks, to_view=False)  # noqa
-                if self.gen_viewnormal_mp4:
-                    create_mp4_from_images(
-                        rendered_view_normals,
-                        output_path=f"{output_dir}/view_normal.mp4",
-                        fps=15,
-                        prompt=prompt,
-                    )
-                else:
-                    render_paths = save_images(
-                        rendered_view_normals,
-                        f"{output_dir}/{RenderItems.VIEW_NORMAL}",
-                        cvt_color=cv2.COLOR_BGR2RGB,
-                    )
-                    data_dict[RenderItems.VIEW_NORMAL.value] = render_paths
-
-        if RenderItems.POSITION_MAP.value in self.render_items:
-            rendered_position, masks = self.renderer.render_position(
-                vertices, faces
-            )
-            norm_position = self.renderer.normalize_map_by_mask(
-                rendered_position, masks
-            )
-            render_paths = save_images(
-                norm_position,
-                f"{output_dir}/{RenderItems.POSITION_MAP}",
-                cvt_color=cv2.COLOR_BGR2RGB,
-            )
-            data_dict[RenderItems.POSITION_MAP.value] = render_paths
-
-        if RenderItems.DEPTH.value in self.render_items:
-            rendered_depth, masks = self.renderer.render_depth(vertices, faces)
-            norm_depth = self.renderer.normalize_map_by_mask(
-                rendered_depth, masks
-            )
-            render_paths = save_images(
-                norm_depth,
-                f"{output_dir}/{RenderItems.DEPTH}",
-            )
-            data_dict[RenderItems.DEPTH.value] = render_paths
-
-            render_paths = save_images(
-                rendered_depth,
-                f"{output_dir}/{RenderItems.DEPTH}_exr",
-                to_uint8=False,
-                format=".exr",
-            )
-            data_dict[f"{RenderItems.DEPTH.value}_exr"] = render_paths
-
-        if RenderItems.IMAGE.value in self.render_items:
-            images = []
-            albedos = []
-            diffuses = []
-            masks, _ = self.renderer.render_rast_alpha(vertices, faces)
-            try:
-                for idx, cam in enumerate(self.camera):
-                    image, albedo, diffuse, _ = render_pbr(
-                        mesh, cam, light_factor=self.light_factor
-                    )
-                    image = torch.cat([image[0], masks[idx]], axis=-1)
-                    images.append(image.detach().cpu().numpy())
-
-                    if RenderItems.ALBEDO.value in self.render_items:
-                        albedo = torch.cat([albedo[0], masks[idx]], axis=-1)
-                        albedos.append(albedo.detach().cpu().numpy())
-
-                    if RenderItems.DIFFUSE.value in self.render_items:
-                        diffuse = torch.cat([diffuse[0], masks[idx]], axis=-1)
-                        diffuses.append(diffuse.detach().cpu().numpy())
-
-            except Exception as e:
-                logger.error(f"[ERROR pbr render]: {e}, skip {mesh_path}")
-                return
-
-            if self.gen_color_gif:
-                create_gif_from_images(
-                    images,
-                    output_path=f"{output_dir}/color.gif",
-                    fps=15,
-                )
-
-            if self.gen_color_mp4:
-                create_mp4_from_images(
-                    images,
-                    output_path=f"{output_dir}/color.mp4",
-                    fps=15,
-                    prompt=prompt,
-                )
-
-            if self.gen_color_mp4 or self.gen_color_gif:
-                return data_dict
-
-            render_paths = save_images(
-                images,
-                f"{output_dir}/{RenderItems.IMAGE}",
-                cvt_color=cv2.COLOR_BGRA2RGBA,
-            )
-            data_dict[RenderItems.IMAGE.value] = render_paths
-
-            render_paths = save_images(
-                albedos,
-                f"{output_dir}/{RenderItems.ALBEDO}",
-                cvt_color=cv2.COLOR_BGRA2RGBA,
-            )
-            data_dict[RenderItems.ALBEDO.value] = render_paths
-
-            render_paths = save_images(
-                diffuses,
-                f"{output_dir}/{RenderItems.DIFFUSE}",
-                cvt_color=cv2.COLOR_BGRA2RGBA,
-            )
-            data_dict[RenderItems.DIFFUSE.value] = render_paths
-
-        data_dict["status"] = "success"
-
-        logger.info(f"Finish rendering in {output_dir}")
-
-        return data_dict
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Render settings")
-
-    parser.add_argument(
-        "--mesh_path",
-        type=str,
-        nargs="+",
-        required=True,
-        help="Paths to the mesh files for rendering.",
-    )
-    parser.add_argument(
-        "--output_root",
-        type=str,
-        required=True,
-        help="Root directory for output",
-    )
-    parser.add_argument(
-        "--uuid",
-        type=str,
-        nargs="+",
-        default=None,
-        help="uuid for rendering saving.",
-    )
-    parser.add_argument(
-        "--num_images", type=int, default=6, help="Number of images to render."
-    )
-    parser.add_argument(
-        "--elevation",
-        type=float,
-        nargs="+",
-        default=[20.0, -10.0],
-        help="Elevation angles for the camera (default: [20.0, -10.0])",
-    )
-    parser.add_argument(
-        "--distance",
-        type=float,
-        default=5,
-        help="Camera distance (default: 5)",
-    )
-    parser.add_argument(
-        "--resolution_hw",
-        type=int,
-        nargs=2,
-        default=(512, 512),
-        help="Resolution of the output images (default: (512, 512))",
-    )
-    parser.add_argument(
-        "--fov",
-        type=float,
-        default=30,
-        help="Field of view in degrees (default: 30)",
-    )
-    parser.add_argument(
-        "--pbr_light_factor",
-        type=float,
-        default=1.0,
-        help="Light factor for mesh PBR rendering (default: 2.)",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        choices=["cpu", "cuda"],
-        default="cuda",
-        help="Device to run on (default: 'cuda')",
-    )
-    parser.add_argument(
-        "--with_mtl",
-        action="store_true",
-        help="Whether to render with mesh material.",
-    )
-    parser.add_argument(
-        "--gen_color_gif",
-        action="store_true",
-        help="Whether to generate color .gif rendering file.",
-    )
-    parser.add_argument(
-        "--gen_color_mp4",
-        action="store_true",
-        help="Whether to generate color .mp4 rendering file.",
-    )
-    parser.add_argument(
-        "--gen_viewnormal_mp4",
-        action="store_true",
-        help="Whether to generate view normal .mp4 rendering file.",
-    )
-    parser.add_argument(
-        "--gen_glonormal_mp4",
-        action="store_true",
-        help="Whether to generate global normal .mp4 rendering file.",
-    )
-    parser.add_argument(
-        "--prompts",
-        type=str,
-        nargs="+",
-        default=None,
-        help="Text prompts for the rendering.",
-    )
-
-    args = parser.parse_args()
-
-    if args.uuid is None:
-        args.uuid = []
-        for path in args.mesh_path:
-            uuid = os.path.basename(path).split(".")[0]
-            args.uuid.append(uuid)
-
-    return args
-
-
-def entrypoint() -> None:
-    args = parse_args()
-
-    camera_settings = CameraSetting(
-        num_images=args.num_images,
-        elevation=args.elevation,
-        distance=args.distance,
-        resolution_hw=args.resolution_hw,
-        fov=math.radians(args.fov),
-        device=args.device,
-    )
-
-    render_items = [
-        RenderItems.ALPHA.value,
-        RenderItems.GLOBAL_NORMAL.value,
-        RenderItems.VIEW_NORMAL.value,
-        RenderItems.POSITION_MAP.value,
-        RenderItems.IMAGE.value,
-        RenderItems.DEPTH.value,
-        # RenderItems.ALBEDO.value,
-        # RenderItems.DIFFUSE.value,
-    ]
-
-    gen_video = (
-        args.gen_color_gif
-        or args.gen_color_mp4
-        or args.gen_viewnormal_mp4
-        or args.gen_glonormal_mp4
-    )
-    if gen_video:
-        render_items = []
-        if args.gen_color_gif or args.gen_color_mp4:
-            render_items.append(RenderItems.IMAGE.value)
-        if args.gen_glonormal_mp4:
-            render_items.append(RenderItems.GLOBAL_NORMAL.value)
-        if args.gen_viewnormal_mp4:
-            render_items.append(RenderItems.VIEW_NORMAL.value)
-            if RenderItems.GLOBAL_NORMAL.value not in render_items:
-                render_items.append(RenderItems.GLOBAL_NORMAL.value)
-
-    image_render = ImageRender(
-        render_items=render_items,
-        camera_params=camera_settings,
-        with_mtl=args.with_mtl,
-        gen_color_gif=args.gen_color_gif,
-        gen_color_mp4=args.gen_color_mp4,
-        gen_viewnormal_mp4=args.gen_viewnormal_mp4,
-        gen_glonormal_mp4=args.gen_glonormal_mp4,
-        light_factor=args.pbr_light_factor,
-        no_index_file=gen_video,
-    )
-    image_render.render_mesh(
-        mesh_path=args.mesh_path,
-        output_root=args.output_root,
-        uuid=args.uuid,
-        prompts=args.prompts,
-    )
-
-    return
-
-
-if __name__ == "__main__":
-    entrypoint()

asset3d_gen/data/mesh_operator.py

@@ -1,425 +0,0 @@
-import logging
-from typing import Tuple, Union
-
-import igraph
-import numpy as np
-import pyvista as pv
-import torch
-import utils3d
-from pymeshfix import _meshfix
-from tqdm import tqdm
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-__all__ = ["MeshFixer"]
-
-
-def radical_inverse(base, n):
-    val = 0
-    inv_base = 1.0 / base
-    inv_base_n = inv_base
-    while n > 0:
-        digit = n % base
-        val += digit * inv_base_n
-        n //= base
-        inv_base_n *= inv_base
-    return val
-
-
-def halton_sequence(dim, n):
-    PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
-    return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
-
-
-def hammersley_sequence(dim, n, num_samples):
-    return [n / num_samples] + halton_sequence(dim - 1, n)
-
-
-def sphere_hammersley_sequence(n, num_samples, offset=(0, 0), remap=False):
-    """Generate a point on a unit sphere using the Hammersley sequence.
-
-    Args:
-        n (int): The index of the sample.
-        num_samples (int): The total number of samples.
-        offset (tuple, optional): Offset for the u and v coordinates.
-        remap (bool, optional): Whether to remap the u coordinate.
-
-    Returns:
-        list: A list containing the spherical coordinates [phi, theta].
-    """
-    u, v = hammersley_sequence(2, n, num_samples)
-    u += offset[0] / num_samples
-    v += offset[1]
-
-    if remap:
-        u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
-
-    theta = np.arccos(1 - 2 * u) - np.pi / 2
-    phi = v * 2 * np.pi
-    return [phi, theta]
-
-
-class MeshFixer(object):
-    """Reduce and postprocess 3D meshes, simplifying and filling holes."""
-
-    def __init__(
-        self,
-        vertices: Union[torch.Tensor, np.ndarray],
-        faces: Union[torch.Tensor, np.ndarray],
-        device: str = "cuda",
-    ) -> None:
-        self.device = device
-        self.vertices = (
-            torch.tensor(vertices, device=device)
-            if isinstance(vertices, np.ndarray)
-            else vertices.to(device)
-        )
-        self.faces = (
-            torch.tensor(faces.astype(np.int32), device=device)
-            if isinstance(faces, np.ndarray)
-            else faces.to(device)
-        )
-
-    @staticmethod
-    def log_mesh_changes(method):
-        def wrapper(self, *args, **kwargs):
-            logger.info(
-                f"Before {method.__name__}: {self.vertices.shape[0]} vertices, {self.faces.shape[0]} faces"  # noqa
-            )
-            result = method(self, *args, **kwargs)
-            logger.info(
-                f"After {method.__name__}: {self.vertices.shape[0]} vertices, {self.faces.shape[0]} faces"  # noqa
-            )
-            return result
-
-        return wrapper
-
-    @log_mesh_changes
-    def fill_holes(
-        self,
-        max_hole_size: float,
-        max_hole_nbe: int,
-        resolution: int,
-        num_views: int,
-        norm_mesh_ratio: float = 1.0,
-    ) -> None:
-        self.vertices = self.vertices * norm_mesh_ratio
-        vertices, self.faces = self._fill_holes(
-            self.vertices,
-            self.faces,
-            max_hole_size,
-            max_hole_nbe,
-            resolution,
-            num_views,
-        )
-        self.vertices = vertices / norm_mesh_ratio
-
-    @staticmethod
-    @torch.no_grad()
-    def _fill_holes(
-        vertices: torch.Tensor,
-        faces: torch.Tensor,
-        max_hole_size: float,
-        max_hole_nbe: int,
-        resolution: int,
-        num_views: int,
-    ) -> Union[torch.Tensor, torch.Tensor]:
-        yaws, pitchs = [], []
-        for i in range(num_views):
-            y, p = sphere_hammersley_sequence(i, num_views)
-            yaws.append(y)
-            pitchs.append(p)
-
-        yaws, pitchs = torch.tensor(yaws).to(vertices), torch.tensor(
-            pitchs
-        ).to(vertices)
-        radius, fov = 2.0, torch.deg2rad(torch.tensor(40)).to(vertices)
-        projection = utils3d.torch.perspective_from_fov_xy(fov, fov, 1, 3)
-
-        views = []
-        for yaw, pitch in zip(yaws, pitchs):
-            orig = (
-                torch.tensor(
-                    [
-                        torch.sin(yaw) * torch.cos(pitch),
-                        torch.cos(yaw) * torch.cos(pitch),
-                        torch.sin(pitch),
-                    ]
-                ).to(vertices)
-                * radius
-            )
-            view = utils3d.torch.view_look_at(
-                orig,
-                torch.tensor([0, 0, 0]).to(vertices),
-                torch.tensor([0, 0, 1]).to(vertices),
-            )
-            views.append(view)
-        views = torch.stack(views, dim=0)
-
-        # Rasterize the mesh
-        visibility = torch.zeros(
-            faces.shape[0], dtype=torch.int32, device=faces.device
-        )
-        rastctx = utils3d.torch.RastContext(backend="cuda")
-
-        for i in tqdm(
-            range(views.shape[0]), total=views.shape[0], desc="Rasterizing"
-        ):
-            view = views[i]
-            buffers = utils3d.torch.rasterize_triangle_faces(
-                rastctx,
-                vertices[None],
-                faces,
-                resolution,
-                resolution,
-                view=view,
-                projection=projection,
-            )
-            face_id = buffers["face_id"][0][buffers["mask"][0] > 0.95] - 1
-            face_id = torch.unique(face_id).long()
-            visibility[face_id] += 1
-
-        # Normalize visibility by the number of views
-        visibility = visibility.float() / num_views
-
-        # Mincut: Identify outer and inner faces
-        edges, face2edge, edge_degrees = utils3d.torch.compute_edges(faces)
-        boundary_edge_indices = torch.nonzero(edge_degrees == 1).reshape(-1)
-        connected_components = utils3d.torch.compute_connected_components(
-            faces, edges, face2edge
-        )
-
-        outer_face_indices = torch.zeros(
-            faces.shape[0], dtype=torch.bool, device=faces.device
-        )
-        for i in range(len(connected_components)):
-            outer_face_indices[connected_components[i]] = visibility[
-                connected_components[i]
-            ] > min(
-                max(
-                    visibility[connected_components[i]].quantile(0.75).item(),
-                    0.25,
-                ),
-                0.5,
-            )
-
-        outer_face_indices = outer_face_indices.nonzero().reshape(-1)
-        inner_face_indices = torch.nonzero(visibility == 0).reshape(-1)
-
-        if inner_face_indices.shape[0] == 0:
-            return vertices, faces
-
-        # Construct dual graph (faces as nodes, edges as edges)
-        dual_edges, dual_edge2edge = utils3d.torch.compute_dual_graph(
-            face2edge
-        )
-        dual_edge2edge = edges[dual_edge2edge]
-        dual_edges_weights = torch.norm(
-            vertices[dual_edge2edge[:, 0]] - vertices[dual_edge2edge[:, 1]],
-            dim=1,
-        )
-
-        # Mincut: Construct main graph and solve the mincut problem
-        g = igraph.Graph()
-        g.add_vertices(faces.shape[0])
-        g.add_edges(dual_edges.cpu().numpy())
-        g.es["weight"] = dual_edges_weights.cpu().numpy()
-
-        g.add_vertex("s")  # source
-        g.add_vertex("t")  # target
-
-        g.add_edges(
-            [(f, "s") for f in inner_face_indices],
-            attributes={
-                "weight": torch.ones(
-                    inner_face_indices.shape[0], dtype=torch.float32
-                )
-                .cpu()
-                .numpy()
-            },
-        )
-        g.add_edges(
-            [(f, "t") for f in outer_face_indices],
-            attributes={
-                "weight": torch.ones(
-                    outer_face_indices.shape[0], dtype=torch.float32
-                )
-                .cpu()
-                .numpy()
-            },
-        )
-
-        cut = g.mincut("s", "t", (np.array(g.es["weight"]) * 1000).tolist())
-        remove_face_indices = torch.tensor(
-            [v for v in cut.partition[0] if v < faces.shape[0]],
-            dtype=torch.long,
-            device=faces.device,
-        )
-
-        # Check if the cut is valid with each connected component
-        to_remove_cc = utils3d.torch.compute_connected_components(
-            faces[remove_face_indices]
-        )
-        valid_remove_cc = []
-        cutting_edges = []
-        for cc in to_remove_cc:
-            # Check visibility median for connected component
-            visibility_median = visibility[remove_face_indices[cc]].median()
-            if visibility_median > 0.25:
-                continue
-
-            # Check if the cutting loop is small enough
-            cc_edge_indices, cc_edges_degree = torch.unique(
-                face2edge[remove_face_indices[cc]], return_counts=True
-            )
-            cc_boundary_edge_indices = cc_edge_indices[cc_edges_degree == 1]
-            cc_new_boundary_edge_indices = cc_boundary_edge_indices[
-                ~torch.isin(cc_boundary_edge_indices, boundary_edge_indices)
-            ]
-            if len(cc_new_boundary_edge_indices) > 0:
-                cc_new_boundary_edge_cc = (
-                    utils3d.torch.compute_edge_connected_components(
-                        edges[cc_new_boundary_edge_indices]
-                    )
-                )
-                cc_new_boundary_edges_cc_center = [
-                    vertices[edges[cc_new_boundary_edge_indices[edge_cc]]]
-                    .mean(dim=1)
-                    .mean(dim=0)
-                    for edge_cc in cc_new_boundary_edge_cc
-                ]
-                cc_new_boundary_edges_cc_area = []
-                for i, edge_cc in enumerate(cc_new_boundary_edge_cc):
-                    _e1 = (
-                        vertices[
-                            edges[cc_new_boundary_edge_indices[edge_cc]][:, 0]
-                        ]
-                        - cc_new_boundary_edges_cc_center[i]
-                    )
-                    _e2 = (
-                        vertices[
-                            edges[cc_new_boundary_edge_indices[edge_cc]][:, 1]
-                        ]
-                        - cc_new_boundary_edges_cc_center[i]
-                    )
-                    cc_new_boundary_edges_cc_area.append(
-                        torch.norm(torch.cross(_e1, _e2, dim=-1), dim=1).sum()
-                        * 0.5
-                    )
-                cutting_edges.append(cc_new_boundary_edge_indices)
-                if any(
-                    [
-                        _l > max_hole_size
-                        for _l in cc_new_boundary_edges_cc_area
-                    ]
-                ):
-                    continue
-
-            valid_remove_cc.append(cc)
-
-        if len(valid_remove_cc) > 0:
-            remove_face_indices = remove_face_indices[
-                torch.cat(valid_remove_cc)
-            ]
-            mask = torch.ones(
-                faces.shape[0], dtype=torch.bool, device=faces.device
-            )
-            mask[remove_face_indices] = 0
-            faces = faces[mask]
-            faces, vertices = utils3d.torch.remove_unreferenced_vertices(
-                faces, vertices
-            )
-
-            tqdm.write(f"Removed {(~mask).sum()} faces by mincut")
-        else:
-            tqdm.write(f"Removed 0 faces by mincut")
-
-        # Fill small boundaries (holes)
-        mesh = _meshfix.PyTMesh()
-        mesh.load_array(vertices.cpu().numpy(), faces.cpu().numpy())
-        mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True)
-
-        _vertices, _faces = mesh.return_arrays()
-        vertices = torch.tensor(_vertices).to(vertices)
-        faces = torch.tensor(_faces).to(faces)
-
-        return vertices, faces
-
-    @property
-    def vertices_np(self) -> np.ndarray:
-        return self.vertices.cpu().numpy()
-
-    @property
-    def faces_np(self) -> np.ndarray:
-        return self.faces.cpu().numpy()
-
-    @log_mesh_changes
-    def simplify(self, ratio: float) -> None:
-        """Simplify the mesh using quadric edge collapse decimation.
-
-        Args:
-            ratio (float): Ratio of faces to filter out.
-        """
-        if ratio <= 0 or ratio >= 1:
-            raise ValueError("Simplify ratio must be between 0 and 1.")
-
-        # Convert to PyVista format for simplification
-        mesh = pv.PolyData(
-            self.vertices_np,
-            np.hstack([np.full((self.faces.shape[0], 1), 3), self.faces_np]),
-        )
-        mesh = mesh.decimate(ratio, progress_bar=True)
-
-        # Update vertices and faces
-        self.vertices = torch.tensor(
-            mesh.points, device=self.device, dtype=torch.float32
-        )
-        self.faces = torch.tensor(
-            mesh.faces.reshape(-1, 4)[:, 1:],
-            device=self.device,
-            dtype=torch.int32,
-        )
-
-    def __call__(
-        self,
-        filter_ratio: float,
-        max_hole_size: float,
-        resolution: int,
-        num_views: int,
-        norm_mesh_ratio: float = 1.0,
-    ) -> Tuple[np.ndarray, np.ndarray]:
-        """Post-process the mesh by simplifying and filling holes.
-
-        This method performs a two-step process:
-        1. Simplifies mesh by reducing faces using quadric edge decimation.
-        2. Fills holes by removing invisible faces, repairing small boundaries.
-
-        Args:
-            filter_ratio (float): Ratio of faces to simplify out.
-                Must be in the range (0, 1).
-            max_hole_size (float): Maximum area of a hole to fill. Connected
-                components of holes larger than this size will not be repaired.
-            resolution (int): Resolution of the rasterization buffer.
-            num_views (int): Number of viewpoints to sample for rasterization.
-            norm_mesh_ratio (float, optional): A scaling factor applied to the
-                vertices of the mesh during processing.
-
-        Returns:
-            Tuple[np.ndarray, np.ndarray]:
-                - vertices: Simplified and repaired vertex array of (V, 3).
-                - faces: Simplified and repaired face array of (F, 3).
-        """
-        self.simplify(ratio=filter_ratio)
-        self.fill_holes(
-            max_hole_size=max_hole_size,
-            max_hole_nbe=int(250 * np.sqrt(1 - filter_ratio)),
-            resolution=resolution,
-            num_views=num_views,
-            norm_mesh_ratio=norm_mesh_ratio,
-        )
-
-        return self.vertices_np, self.faces_np

asset3d_gen/data/utils.py

@@ -1,943 +0,0 @@
-import math
-import os
-import random
-from glob import glob
-from typing import List, Tuple, Union
-
-import cv2
-import kaolin as kal
-import numpy as np
-import nvdiffrast.torch as dr
-import torch
-import torch.nn.functional as F
-from PIL import Image
-
-try:
-    from kolors.models.modeling_chatglm import ChatGLMModel
-    from kolors.models.tokenization_chatglm import ChatGLMTokenizer
-except ImportError:
-    ChatGLMTokenizer = None
-    ChatGLMModel = None
-import logging
-from dataclasses import dataclass, field
-from enum import Enum
-
-import trimesh
-from kaolin.render.camera import Camera
-from torch import nn
-
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "center_points",
-    "get_points_stat",
-    "DiffrastRender",
-    "compute_cam_pts_by_az_el",
-    "compute_cam_pts_by_views",
-    "save_images",
-    "render_pbr",
-    "load_llm_models",
-    "prelabel_text_feature",
-    "calc_vertex_normals",
-    "normalize_vertices_array",
-    "load_mesh_to_unit_cube",
-    "as_list",
-    "CameraSetting",
-    "RenderItems",
-    "import_kaolin_mesh",
-    "save_mesh_with_mtl",
-    "get_images_from_grid",
-    "post_process_texture",
-]
-
-
-def get_points_stat(
-    points: torch.FloatTensor, eps: float = 1e-6
-) -> torch.FloatTensor:
-    assert (
-        len(points.shape) == 3
-    ), f"Points have unexpected shape {points.shape}"
-
-    vmin = points.min(dim=1, keepdim=True)[0]
-    vmax = points.max(dim=1, keepdim=True)[0]
-    pts_center = (vmin + vmax) / 2
-
-    pts_dim = (vmax - vmin).max(dim=-1, keepdim=True)[0].clip(min=eps)
-
-    return pts_center, pts_dim
-
-
-def center_points(
-    points: torch.FloatTensor, normalize: bool = False, eps: float = 1e-6
-) -> torch.FloatTensor:
-    vmid, den = get_points_stat(points)
-
-    res = points - vmid
-
-    if normalize:
-        res = res / den
-
-    return res
-
-
-class DiffrastRender(object):
-    """A class to handle differentiable rendering using nvdiffrast.
-
-    This class provides methods to render position, depth, and normal maps
-    with optional anti-aliasing and gradient disabling for rasterization.
-
-    Attributes:
-        p_mtx (torch.Tensor): Projection matrix.
-        mv_mtx (torch.Tensor): Model-view matrix.
-        mvp_mtx (torch.Tensor): Model-view-projection matrix, calculated as
-            p_mtx @ mv_mtx if not provided.
-        resolution_hw (Tuple[int, int]): Height and width of the rendering resolution.  # noqa
-        _ctx (Union[dr.RasterizeCudaContext, dr.RasterizeGLContext]): Rasterization context.  # noqa
-        mask_thresh (float): Threshold for mask creation.
-        grad_db (bool): Whether to disable gradients during rasterization.
-        antialias_mask (bool): Whether to apply anti-aliasing to the mask.
-        device (str): Device used for rendering ('cuda' or 'cpu').
-
-    Methods:
-        _warmup(glctx): Warmup rasterization by rendering a simple triangle.
-        compute_dr_raster(vertices, faces): Rasterizes the mesh and returns
-            rasterized outputs and transformed vertices.
-        transform_vertices(vertices, matrix): Transforms the vertices using
-            the provided transformation matrix.
-        normalize_map_by_mask_separately(map, mask): Normalizes each map in
-            the batch separately using the mask.
-        normalize_map_by_mask(map, mask): Normalizes the entire map using the
-            mask, keeping the output in the range [0, 1].
-        render_position(vertices, faces): Renders the position map and
-            alpha mask from the given vertices and faces.
-        render_depth(vertices, faces): Renders the depth map and alpha
-            mask from the given vertices and faces.
-        _compute_mask(rast, vertices_clip, faces): Computes the mask from the
-            rasterization output.
-        render_global_normal(vertices, faces, vertice_normals): Renders the
-            normal map and alpha mask from the given vertices, faces, and
-            vertex normals.
-        transform_normal_to_view(normals, mat_w2c, masks): Transforms the normals
-            to the view space using the world-to-camera matrix.
-    """
-
-    def __init__(
-        self,
-        p_matrix: torch.Tensor,
-        mv_matrix: torch.Tensor,
-        resolution_hw: Tuple[int, int],
-        context: Union[dr.RasterizeCudaContext, dr.RasterizeGLContext] = None,
-        mvp_matrix: torch.Tensor = None,
-        mask_thresh: float = 0.5,
-        grad_db: bool = False,
-        antialias_mask: bool = True,
-        align_coordinate: bool = True,
-        device: str = "cuda",
-    ) -> None:
-        self.p_mtx = p_matrix
-        self.mv_mtx = mv_matrix
-        if mvp_matrix is None:
-            self.mvp_mtx = torch.bmm(p_matrix, mv_matrix)
-
-        self.resolution_hw = resolution_hw
-        if context is None:
-            context = dr.RasterizeCudaContext(device=device)
-        self._ctx = context
-        self.mask_thresh = mask_thresh
-        self.grad_db = grad_db
-        self.antialias_mask = antialias_mask
-        self.align_coordinate = align_coordinate
-        self.device = device
-        # self._warmup(self._ctx)
-
-    def _warmup(self, glctx):
-        # Seem solved. https://github.com/NVlabs/nvdiffrast/issues/59
-        def tensor(*args, **kwargs):
-            return torch.tensor(*args, device=self.device, **kwargs)
-
-        pos = tensor(
-            [[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]],
-            dtype=torch.float32,
-        )
-        tri = tensor([[0, 1, 2]], dtype=torch.int32)
-        dr.rasterize(glctx, pos, tri, resolution=[256, 256])
-
-    def compute_dr_raster(
-        self,
-        vertices: torch.Tensor,
-        faces: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        vertices_clip = self.transform_vertices(vertices, matrix=self.mvp_mtx)
-        rast, _ = dr.rasterize(
-            self._ctx,
-            vertices_clip,
-            faces.int(),
-            resolution=self.resolution_hw,
-            grad_db=self.grad_db,
-        )
-
-        return rast, vertices_clip
-
-    def transform_vertices(
-        self,
-        vertices: torch.Tensor,
-        matrix: torch.Tensor,
-    ) -> torch.Tensor:
-        verts_ones = torch.ones((len(vertices), 1)).to(vertices)
-        verts_homo = torch.cat([vertices, verts_ones], dim=-1)
-        trans_vertices = torch.matmul(verts_homo, matrix.permute(0, 2, 1))
-
-        return trans_vertices
-
-    def normalize_map_by_mask_separately(
-        self, map: torch.Tensor, mask: torch.Tensor
-    ) -> torch.Tensor:
-        # Normalize each map separately by mask, normalized map in [0, 1].
-        normalized_maps = []
-        for map_item, mask_item in zip(map, mask):
-            normalized_map = self.normalize_map_by_mask(map_item, mask_item)
-            normalized_maps.append(normalized_map)
-
-        normalized_maps = torch.stack(normalized_maps, dim=0)
-
-        return normalized_maps
-
-    def normalize_map_by_mask(
-        self, map: torch.Tensor, mask: torch.Tensor
-    ) -> torch.Tensor:
-        # Normalize all maps in total by mask, normalized map in [0, 1].
-        foreground = (mask == 1).squeeze(dim=-1)
-        foreground_elements = map[foreground]
-        if len(foreground_elements) == 0:
-            return map
-
-        min_val, _ = foreground_elements.min(dim=0)
-        max_val, _ = foreground_elements.max(dim=0)
-        val_range = (max_val - min_val).clip(min=1e-6)
-
-        normalized_map = (map - min_val) / val_range
-        normalized_map = torch.lerp(
-            torch.zeros_like(normalized_map), normalized_map, mask
-        )
-        normalized_map[normalized_map < 0] = 0
-
-        return normalized_map
-
-    def _compute_mask(
-        self,
-        rast: torch.Tensor,
-        vertices_clip: torch.Tensor,
-        faces: torch.Tensor,
-    ) -> torch.Tensor:
-        mask = (rast[..., 3:] > 0).float()
-        mask = mask.clip(min=0, max=1)
-
-        if self.antialias_mask is True:
-            mask = dr.antialias(mask, rast, vertices_clip, faces)
-        else:
-            foreground = mask > self.mask_thresh
-            mask[foreground] = 1
-            mask[~foreground] = 0
-
-        return mask
-
-    def render_rast_alpha(
-        self,
-        vertices: torch.Tensor,
-        faces: torch.Tensor,
-    ):
-        faces = faces.to(torch.int32)
-        rast, vertices_clip = self.compute_dr_raster(vertices, faces)
-        mask = self._compute_mask(rast, vertices_clip, faces)
-
-        return mask, rast
-
-    def render_position(
-        self,
-        vertices: torch.Tensor,
-        faces: torch.Tensor,
-    ) -> Union[torch.Tensor, torch.Tensor]:
-        # Vertices in model coordinate system, real position coordinate number.
-        faces = faces.to(torch.int32)
-        mask, rast = self.render_rast_alpha(vertices, faces)
-
-        vertices_model = vertices[None, ...].contiguous().float()
-        position_map, _ = dr.interpolate(vertices_model, rast, faces)
-        # Align with blender.
-        if self.align_coordinate:
-            position_map = position_map[..., [0, 2, 1]]
-            position_map[..., 1] = -position_map[..., 1]
-
-        position_map = torch.lerp(
-            torch.zeros_like(position_map), position_map, mask
-        )
-
-        return position_map, mask
-
-    def render_uv(
-        self,
-        vertices: torch.Tensor,
-        faces: torch.Tensor,
-        vtx_uv: torch.Tensor,
-    ) -> Union[torch.Tensor, torch.Tensor]:
-        faces = faces.to(torch.int32)
-        mask, rast = self.render_rast_alpha(vertices, faces)
-        uv_map, _ = dr.interpolate(vtx_uv, rast, faces)
-        uv_map = torch.lerp(torch.zeros_like(uv_map), uv_map, mask)
-
-        return uv_map, mask
-
-    def render_depth(
-        self,
-        vertices: torch.Tensor,
-        faces: torch.Tensor,
-    ) -> Union[torch.Tensor, torch.Tensor]:
-        # Vertices in model coordinate system, real depth coordinate number.
-        faces = faces.to(torch.int32)
-        mask, rast = self.render_rast_alpha(vertices, faces)
-
-        vertices_camera = self.transform_vertices(vertices, matrix=self.mv_mtx)
-        vertices_camera = vertices_camera[..., 2:3].contiguous().float()
-        depth_map, _ = dr.interpolate(vertices_camera, rast, faces)
-        # Change camera depth minus to positive.
-        if self.align_coordinate:
-            depth_map = -depth_map
-        depth_map = torch.lerp(torch.zeros_like(depth_map), depth_map, mask)
-
-        return depth_map, mask
-
-    def render_global_normal(
-        self,
-        vertices: torch.Tensor,
-        faces: torch.Tensor,
-        vertice_normals: torch.Tensor,
-    ) -> Union[torch.Tensor, torch.Tensor]:
-        # NOTE: vertice_normals in [-1, 1],  return normal in [0, 1].
-        # vertices / vertice_normals in model coordinate system.
-        faces = faces.to(torch.int32)
-        mask, rast = self.render_rast_alpha(vertices, faces)
-        im_base_normals, _ = dr.interpolate(
-            vertice_normals[None, ...].float(), rast, faces
-        )
-
-        if im_base_normals is not None:
-            faces = faces.to(torch.int64)
-            vertices_cam = self.transform_vertices(
-                vertices, matrix=self.mv_mtx
-            )
-            face_vertices_ndc = kal.ops.mesh.index_vertices_by_faces(
-                vertices_cam[..., :3], faces
-            )
-            face_normal_sign = kal.ops.mesh.face_normals(face_vertices_ndc)[
-                ..., 2
-            ]
-            for idx in range(len(im_base_normals)):
-                face_idx = (rast[idx, ..., -1].long() - 1).contiguous()
-                im_normal_sign = torch.sign(face_normal_sign[idx, face_idx])
-                im_normal_sign[face_idx == -1] = 0
-                im_base_normals[idx] *= im_normal_sign.unsqueeze(-1)
-
-        normal = (im_base_normals + 1) / 2
-        normal = normal.clip(min=0, max=1)
-        normal = torch.lerp(torch.zeros_like(normal), normal, mask)
-
-        return normal, mask
-
-    def transform_normal(
-        self,
-        normals: torch.Tensor,
-        trans_matrix: torch.Tensor,
-        masks: torch.Tensor,
-        to_view: bool,
-    ) -> torch.Tensor:
-        # NOTE: input normals in [0, 1], output normals in [0, 1].
-        normals = normals.clone()
-        assert len(normals) == len(trans_matrix)
-
-        if not to_view:
-            # Flip the sign on the x-axis to match inv bae system for global transformation.  # noqa
-            normals[..., 0] = 1 - normals[..., 0]
-
-        normals = 2 * normals - 1
-        b, h, w, c = normals.shape
-
-        transformed_normals = []
-        for normal, matrix in zip(normals, trans_matrix):
-            # Transform normals using the transformation matrix (4x4).
-            reshaped_normals = normal.view(-1, c)  # (h w 3) -> (hw 3)
-            padded_vectors = torch.nn.functional.pad(
-                reshaped_normals, pad=(0, 1), mode="constant", value=0.0
-            )
-            transformed_normal = torch.matmul(
-                padded_vectors, matrix.transpose(0, 1)
-            )[..., :3]
-
-            # Normalize and clip the normals to [0, 1] range.
-            transformed_normal = F.normalize(transformed_normal, p=2, dim=-1)
-            transformed_normal = (transformed_normal + 1) / 2
-
-            if to_view:
-                # Flip the sign on the x-axis to match bae system for view transformation.  # noqa
-                transformed_normal[..., 0] = 1 - transformed_normal[..., 0]
-
-            transformed_normals.append(transformed_normal.view(h, w, c))
-
-        transformed_normals = torch.stack(transformed_normals, dim=0)
-
-        if masks is not None:
-            transformed_normals = torch.lerp(
-                torch.zeros_like(transformed_normals),
-                transformed_normals,
-                masks,
-            )
-
-        return transformed_normals
-
-
-def az_el_to_points(
-    azimuths: np.ndarray, elevations: np.ndarray
-) -> np.ndarray:
-    x = np.cos(azimuths) * np.cos(elevations)
-    y = np.sin(azimuths) * np.cos(elevations)
-    z = np.sin(elevations)
-
-    return np.stack([x, y, z], axis=-1)
-
-
-def compute_az_el_by_views(
-    num_view: int, el: float
-) -> Tuple[np.ndarray, np.ndarray]:
-    azimuths = np.arange(num_view) / num_view * np.pi * 2
-    elevations = np.deg2rad(np.array([el] * num_view))
-
-    return azimuths, elevations
-
-
-def compute_cam_pts_by_az_el(
-    azs: np.ndarray,
-    els: np.ndarray,
-    distance: float,
-    extra_pts: np.ndarray = None,
-) -> np.ndarray:
-    distances = np.array([distance for _ in range(len(azs))])
-    cam_pts = az_el_to_points(azs, els) * distances[:, None]
-
-    if extra_pts is not None:
-        cam_pts = np.concatenate([cam_pts, extra_pts], axis=0)
-
-    # Align coordinate system.
-    cam_pts = cam_pts[:, [0, 2, 1]]  # xyz -> xzy
-    cam_pts[..., 2] = -cam_pts[..., 2]
-
-    return cam_pts
-
-
-def compute_cam_pts_by_views(
-    num_view: int, el: float, distance: float, extra_pts: np.ndarray = None
-) -> torch.Tensor:
-    """Computes object-center camera points for a given number of views.
-
-    Args:
-        num_view (int): The number of views (camera positions) to compute.
-        el (float): The elevation angle in degrees.
-        distance (float): The distance from the origin to the camera.
-        extra_pts (np.ndarray): Extra camera points postion.
-
-    Returns:
-        torch.Tensor: A tensor containing the camera points for each view, with shape `(num_view, 3)`. # noqa
-    """
-    azimuths, elevations = compute_az_el_by_views(num_view, el)
-    cam_pts = compute_cam_pts_by_az_el(
-        azimuths, elevations, distance, extra_pts
-    )
-
-    return cam_pts
-
-
-def save_images(
-    images: Union[list[np.ndarray], list[torch.Tensor]],
-    output_dir: str,
-    cvt_color: str = None,
-    format: str = ".png",
-    to_uint8: bool = True,
-    verbose: bool = False,
-) -> List[str]:
-    # NOTE: images in [0, 1]
-    os.makedirs(output_dir, exist_ok=True)
-    save_paths = []
-    for idx, image in enumerate(images):
-        if isinstance(image, torch.Tensor):
-            image = image.detach().cpu().numpy()
-        if to_uint8:
-            image = image.clip(min=0, max=1)
-            image = (255.0 * image).astype(np.uint8)
-        if cvt_color is not None:
-            image = cv2.cvtColor(image, cvt_color)
-        save_path = os.path.join(output_dir, f"{idx:04d}{format}")
-        save_paths.append(save_path)
-
-        cv2.imwrite(save_path, image)
-
-    if verbose:
-        logger.info(f"Images saved in {output_dir}")
-
-    return save_paths
-
-
-def current_lighting(
-    azimuths: List[float],
-    elevations: List[float],
-    light_factor: float = 1.0,
-    device: str = "cuda",
-):
-    # azimuths, elevations in degress.
-    directions = []
-    for az, el in zip(azimuths, elevations):
-        az, el = math.radians(az), math.radians(el)
-        direction = kal.render.lighting.sg_direction_from_azimuth_elevation(
-            az, el
-        )
-        directions.append(direction)
-    directions = torch.cat(directions, dim=0)
-
-    amplitude = torch.ones_like(directions) * light_factor
-    light_condition = kal.render.lighting.SgLightingParameters(
-        amplitude=amplitude,
-        direction=directions,
-        sharpness=3,
-    ).to(device)
-
-    # light_condition = kal.render.lighting.SgLightingParameters.from_sun(
-    #     directions, strength=1, angle=90, color=None
-    # ).to(device)
-
-    return light_condition
-
-
-def render_pbr(
-    mesh,
-    camera,
-    device="cuda",
-    cxt=None,
-    custom_materials=None,
-    light_factor=1.0,
-):
-    if cxt is None:
-        cxt = dr.RasterizeCudaContext()
-
-    light_condition = current_lighting(
-        azimuths=[0, 90, 180, 270],
-        elevations=[90, 60, 30, 20],
-        light_factor=light_factor,
-        device=device,
-    )
-    render_res = kal.render.easy_render.render_mesh(
-        camera,
-        mesh,
-        lighting=light_condition,
-        nvdiffrast_context=cxt,
-        custom_materials=custom_materials,
-    )
-
-    image = render_res[kal.render.easy_render.RenderPass.render]
-    image = image.clip(0, 1)
-
-    albedo = render_res[kal.render.easy_render.RenderPass.albedo]
-    albedo = albedo.clip(0, 1)
-
-    diffuse = render_res[kal.render.easy_render.RenderPass.diffuse]
-    diffuse = diffuse.clip(0, 1)
-
-    normal = render_res[kal.render.easy_render.RenderPass.normals]
-    normal = normal.clip(-1, 1)
-
-    return image, albedo, diffuse, normal
-
-
-def load_saved_normal(path: str) -> np.ndarray:
-    image_paths = glob(os.path.join(path, "*.jpg"))
-    images = []
-    for path in sorted(image_paths):
-        image = cv2.imread(path)
-        image = image[..., ::-1]  # rgb -> bgr
-        images.append(image)
-    images = np.stack(images, axis=0)
-
-    return images
-
-
-def _move_to_target_device(data, device: str):
-    if isinstance(data, dict):
-        for key, value in data.items():
-            data[key] = _move_to_target_device(value, device)
-    elif isinstance(data, torch.Tensor):
-        return data.to(device)
-
-    return data
-
-
-def _encode_prompt(
-    prompt_batch,
-    text_encoders,
-    tokenizers,
-    proportion_empty_prompts=0,
-    is_train=True,
-):
-    prompt_embeds_list = []
-
-    captions = []
-    for caption in prompt_batch:
-        if random.random() < proportion_empty_prompts:
-            captions.append("")
-        elif isinstance(caption, str):
-            captions.append(caption)
-        elif isinstance(caption, (list, np.ndarray)):
-            captions.append(random.choice(caption) if is_train else caption[0])
-
-    with torch.no_grad():
-        for tokenizer, text_encoder in zip(tokenizers, text_encoders):
-            text_inputs = tokenizer(
-                captions,
-                padding="max_length",
-                max_length=256,
-                truncation=True,
-                return_tensors="pt",
-            ).to(text_encoder.device)
-
-            output = text_encoder(
-                input_ids=text_inputs.input_ids,
-                attention_mask=text_inputs.attention_mask,
-                position_ids=text_inputs.position_ids,
-                output_hidden_states=True,
-            )
-
-            # We are only interested in the pooled output of the text encoder.
-            prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone()
-            pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone()
-            bs_embed, seq_len, _ = prompt_embeds.shape
-            prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
-            prompt_embeds_list.append(prompt_embeds)
-
-    prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
-    pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
-
-    return prompt_embeds, pooled_prompt_embeds
-
-
-def load_llm_models(pretrained_model_name_or_path: str, device: str):
-    tokenizer = ChatGLMTokenizer.from_pretrained(
-        pretrained_model_name_or_path,
-        subfolder="text_encoder",
-    )
-    text_encoder = ChatGLMModel.from_pretrained(
-        pretrained_model_name_or_path,
-        subfolder="text_encoder",
-    ).to(device)
-
-    text_encoders = [
-        text_encoder,
-    ]
-    tokenizers = [
-        tokenizer,
-    ]
-
-    logger.info(f"Load model from {pretrained_model_name_or_path} done.")
-
-    return tokenizers, text_encoders
-
-
-def prelabel_text_feature(
-    prompt_batch: List[str],
-    output_dir: str,
-    tokenizers: nn.Module,
-    text_encoders: nn.Module,
-) -> List[str]:
-    os.makedirs(output_dir, exist_ok=True)
-
-    # prompt_batch ["text..."]
-    prompt_embeds, pooled_prompt_embeds = _encode_prompt(
-        prompt_batch, text_encoders, tokenizers
-    )
-
-    prompt_embeds = _move_to_target_device(prompt_embeds, device="cpu")
-    pooled_prompt_embeds = _move_to_target_device(
-        pooled_prompt_embeds, device="cpu"
-    )
-
-    data_dict = dict(
-        prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds
-    )
-
-    save_path = os.path.join(output_dir, "text_feat.pth")
-    torch.save(data_dict, save_path)
-
-    return save_path
-
-
-def calc_face_normals(
-    vertices: torch.Tensor,  # V,3 first vertex may be unreferenced
-    faces: torch.Tensor,  # F,3 long, first face may be all zero
-    normalize: bool = False,
-) -> torch.Tensor:  # F,3
-    full_vertices = vertices[faces]  # F,C=3,3
-    v0, v1, v2 = full_vertices.unbind(dim=1)  # F,3
-    face_normals = torch.cross(v1 - v0, v2 - v0, dim=1)  # F,3
-    if normalize:
-        face_normals = F.normalize(
-            face_normals, eps=1e-6, dim=1
-        )  # TODO inplace?
-    return face_normals  # F,3
-
-
-def calc_vertex_normals(
-    vertices: torch.Tensor,  # V,3 first vertex may be unreferenced
-    faces: torch.Tensor,  # F,3 long, first face may be all zero
-    face_normals: torch.Tensor = None,  # F,3, not normalized
-) -> torch.Tensor:  # F,3
-    _F = faces.shape[0]
-
-    if face_normals is None:
-        face_normals = calc_face_normals(vertices, faces)
-
-    vertex_normals = torch.zeros(
-        (vertices.shape[0], 3, 3), dtype=vertices.dtype, device=vertices.device
-    )  # V,C=3,3
-    vertex_normals.scatter_add_(
-        dim=0,
-        index=faces[:, :, None].expand(_F, 3, 3),
-        src=face_normals[:, None, :].expand(_F, 3, 3),
-    )
-    vertex_normals = vertex_normals.sum(dim=1)  # V,3
-    return F.normalize(vertex_normals, eps=1e-6, dim=1)
-
-
-def normalize_vertices_array(
-    vertices: Union[torch.Tensor, np.ndarray],
-    mesh_scale: float = 1.0,
-    exec_norm: bool = True,
-):
-    if isinstance(vertices, torch.Tensor):
-        bbmin, bbmax = vertices.min(0)[0], vertices.max(0)[0]
-    else:
-        bbmin, bbmax = vertices.min(0), vertices.max(0)  # (3,)
-    center = (bbmin + bbmax) * 0.5
-    bbsize = bbmax - bbmin
-    scale = 2 * mesh_scale / bbsize.max()
-    if exec_norm:
-        vertices = (vertices - center) * scale
-
-    return vertices, scale, center
-
-
-def load_mesh_to_unit_cube(
-    mesh_file: str,
-    mesh_scale: float = 1.0,
-) -> tuple[trimesh.Trimesh, float, list[float]]:
-    if not os.path.exists(mesh_file):
-        raise FileNotFoundError(f"mesh_file path {mesh_file} not exists.")
-
-    mesh = trimesh.load(mesh_file)
-    if isinstance(mesh, trimesh.Scene):
-        mesh = trimesh.utils.concatenate(mesh)
-
-    vertices, scale, center = normalize_vertices_array(
-        mesh.vertices, mesh_scale
-    )
-    mesh.vertices = vertices
-
-    return mesh, scale, center
-
-
-def as_list(obj):
-    if isinstance(obj, (list, tuple)):
-        return obj
-    elif isinstance(obj, set):
-        return list(obj)
-    else:
-        return [obj]
-
-
-@dataclass
-class CameraSetting:
-    """Camera settings for images rendering."""
-
-    num_images: int
-    elevation: list[float]
-    distance: float
-    resolution_hw: tuple[int, int]
-    fov: float
-    at: tuple[float, float, float] = field(
-        default_factory=lambda: (0.0, 0.0, 0.0)
-    )
-    up: tuple[float, float, float] = field(
-        default_factory=lambda: (0.0, 1.0, 0.0)
-    )
-    device: str = "cuda"
-    near: float = 1e-2
-    far: float = 1e2
-
-    def __post_init__(
-        self,
-    ):
-        h = self.resolution_hw[0]
-        f = (h / 2) / math.tan(self.fov / 2)
-        cx = self.resolution_hw[1] / 2
-        cy = self.resolution_hw[0] / 2
-        Ks = [
-            [f, 0, cx],
-            [0, f, cy],
-            [0, 0, 1],
-        ]
-
-        self.Ks = Ks
-
-
-@dataclass
-class RenderItems(str, Enum):
-    IMAGE = "image_color"
-    ALPHA = "image_mask"
-    VIEW_NORMAL = "image_view_normal"
-    GLOBAL_NORMAL = "image_global_normal"
-    POSITION_MAP = "image_position"
-    DEPTH = "image_depth"
-    ALBEDO = "image_albedo"
-    DIFFUSE = "image_diffuse"
-
-
-def compute_az_el_by_camera_params(
-    camera_params: CameraSetting, flip_az: bool = False
-):
-    num_view = camera_params.num_images // len(camera_params.elevation)
-    view_interval = 2 * np.pi / num_view / 2
-    azimuths = []
-    elevations = []
-    for idx, el in enumerate(camera_params.elevation):
-        azs = np.arange(num_view) / num_view * np.pi * 2 + idx * view_interval
-        if flip_az:
-            azs *= -1
-        els = np.deg2rad(np.array([el] * num_view))
-        azimuths.append(azs)
-        elevations.append(els)
-
-    azimuths = np.concatenate(azimuths, axis=0)
-    elevations = np.concatenate(elevations, axis=0)
-
-    return azimuths, elevations
-
-
-def init_kal_camera(camera_params: CameraSetting) -> Camera:
-    azimuths, elevations = compute_az_el_by_camera_params(camera_params)
-    cam_pts = compute_cam_pts_by_az_el(
-        azimuths, elevations, camera_params.distance
-    )
-
-    up = torch.cat(
-        [
-            torch.tensor(camera_params.up).repeat(camera_params.num_images, 1),
-        ],
-        dim=0,
-    )
-
-    camera = Camera.from_args(
-        eye=torch.tensor(cam_pts),
-        at=torch.tensor(camera_params.at),
-        up=up,
-        fov=camera_params.fov,
-        height=camera_params.resolution_hw[0],
-        width=camera_params.resolution_hw[1],
-        near=camera_params.near,
-        far=camera_params.far,
-        device=camera_params.device,
-    )
-
-    return camera
-
-
-def import_kaolin_mesh(mesh_path: str, with_mtl: bool = False):
-    if mesh_path.endswith(".glb"):
-        mesh = kal.io.gltf.import_mesh(mesh_path)
-    elif mesh_path.endswith(".obj"):
-        with_material = True if with_mtl else False
-        mesh = kal.io.obj.import_mesh(mesh_path, with_materials=with_material)
-        if with_mtl and mesh.materials and len(mesh.materials) > 0:
-            material = kal.render.materials.PBRMaterial()
-            assert (
-                "map_Kd" in mesh.materials[0]
-            ), "'map_Kd' not found in materials."
-            material.diffuse_texture = mesh.materials[0]["map_Kd"] / 255.0
-            mesh.materials = [material]
-    elif mesh_path.endswith(".ply"):
-        mesh = trimesh.load(mesh_path)
-        mesh_path = mesh_path.replace(".ply", ".obj")
-        mesh.export(mesh_path)
-        mesh = kal.io.obj.import_mesh(mesh_path)
-    elif mesh_path.endswith(".off"):
-        mesh = kal.io.off.import_mesh(mesh_path)
-    else:
-        raise RuntimeError(
-            f"{mesh_path} mesh type not supported, "
-            "supported mesh type `.glb`, `.obj`, `.ply`, `.off`."
-        )
-
-    return mesh
-
-
-def save_mesh_with_mtl(
-    vertices: np.ndarray,
-    faces: np.ndarray,
-    uvs: np.ndarray,
-    texture: Union[Image.Image, np.ndarray],
-    output_path: str,
-    material_base=(250, 250, 250, 255),
-) -> trimesh.Trimesh:
-    if isinstance(texture, np.ndarray):
-        texture = Image.fromarray(texture)
-
-    mesh = trimesh.Trimesh(
-        vertices,
-        faces,
-        visual=trimesh.visual.TextureVisuals(uv=uvs, image=texture),
-    )
-    mesh.visual.material = trimesh.visual.material.SimpleMaterial(
-        image=texture,
-        diffuse=material_base,
-        ambient=material_base,
-        specular=material_base,
-    )
-
-    dir_name = os.path.dirname(output_path)
-    os.makedirs(dir_name, exist_ok=True)
-
-    _ = mesh.export(output_path)
-    # texture.save(os.path.join(dir_name, f"{file_name}_texture.png"))
-
-    logger.info(f"Saved mesh with texture to {output_path}")
-
-    return mesh
-
-
-def get_images_from_grid(
-    image: Union[str, Image.Image], img_size: int
-) -> list[Image.Image]:
-    if isinstance(image, str):
-        image = Image.open(image)
-
-    view_images = np.array(image)
-    view_images = np.concatenate(
-        [view_images[:img_size, ...], view_images[img_size:, ...]], axis=1
-    )
-    images = np.split(view_images, view_images.shape[1] // img_size, axis=1)
-    images = [Image.fromarray(img) for img in images]
-
-    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

asset3d_gen/models/delight.py

@@ -1,162 +0,0 @@
-import os
-from typing import Union
-
-import cv2
-import numpy as np
-import torch
-from diffusers import (
-    EulerAncestralDiscreteScheduler,
-    StableDiffusionInstructPix2PixPipeline,
-)
-from huggingface_hub import snapshot_download
-from PIL import Image
-from asset3d_gen.models.segment import RembgRemover
-
-__all__ = [
-    "DelightingModel",
-]
-
-
-class DelightingModel(object):
-    def __init__(
-        self,
-        model_path: str = None,
-        num_infer_step: int = 50,
-        mask_erosion_size: int = 3,
-        image_guide_scale: float = 1.5,
-        text_guide_scale: float = 1.0,
-        device: str = "cuda",
-        seed: int = 0,
-    ) -> None:
-        self.image_guide_scale = image_guide_scale
-        self.text_guide_scale = text_guide_scale
-        self.num_infer_step = num_infer_step
-        self.mask_erosion_size = mask_erosion_size
-        self.kernel = np.ones(
-            (self.mask_erosion_size, self.mask_erosion_size), np.uint8
-        )
-        self.seed = seed
-        self.device = device
-        self.bg_remover = RembgRemover()
-
-        if model_path is None:
-            suffix = "hunyuan3d-delight-v2-0"
-            model_path = snapshot_download(
-                repo_id="tencent/Hunyuan3D-2", allow_patterns=f"{suffix}/*"
-            )
-            model_path = os.path.join(model_path, suffix)
-
-        pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
-            model_path,
-            torch_dtype=torch.float16,
-            safety_checker=None,
-        )
-        pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
-            pipeline.scheduler.config
-        )
-        pipeline.set_progress_bar_config(disable=True)
-
-        pipeline.to(self.device, torch.float16)
-        # pipeline.enable_model_cpu_offload()
-        # pipeline.enable_xformers_memory_efficient_attention()
-        self.pipeline = pipeline
-
-    def recenter_image(
-        self, image: Image.Image, border_ratio: float = 0.2
-    ) -> Image.Image:
-        if image.mode == "RGB":
-            return image
-        elif image.mode == "L":
-            image = image.convert("RGB")
-            return image
-
-        alpha_channel = np.array(image)[:, :, 3]
-        non_zero_indices = np.argwhere(alpha_channel > 0)
-        if non_zero_indices.size == 0:
-            raise ValueError("Image is fully transparent")
-
-        min_row, min_col = non_zero_indices.min(axis=0)
-        max_row, max_col = non_zero_indices.max(axis=0)
-
-        cropped_image = image.crop(
-            (min_col, min_row, max_col + 1, max_row + 1)
-        )
-
-        width, height = cropped_image.size
-        border_width = int(width * border_ratio)
-        border_height = int(height * border_ratio)
-
-        new_width = width + 2 * border_width
-        new_height = height + 2 * border_height
-
-        square_size = max(new_width, new_height)
-
-        new_image = Image.new(
-            "RGBA", (square_size, square_size), (255, 255, 255, 0)
-        )
-
-        paste_x = (square_size - new_width) // 2 + border_width
-        paste_y = (square_size - new_height) // 2 + border_height
-
-        new_image.paste(cropped_image, (paste_x, paste_y))
-
-        return new_image
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        image: Union[str, np.ndarray, Image.Image],
-        preprocess: bool = False,
-        target_wh: tuple[int, int] = None,
-    ) -> Image.Image:
-        if isinstance(image, str):
-            image = Image.open(image)
-        elif isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-
-        if preprocess:
-            image = self.bg_remover(image)
-            image = self.recenter_image(image)
-
-        if target_wh is not None:
-            image = image.resize(target_wh)
-        else:
-            target_wh = image.size
-
-        image_array = np.array(image)
-        assert image_array.shape[-1] == 4, "Image must have alpha channel"
-
-        raw_alpha_channel = image_array[:, :, 3]
-        alpha_channel = cv2.erode(raw_alpha_channel, self.kernel, iterations=1)
-        image_array[alpha_channel == 0, :3] = 255  # must be white background
-        image_array[:, :, 3] = alpha_channel
-
-        image = self.pipeline(
-            prompt="",
-            image=Image.fromarray(image_array).convert("RGB"),
-            generator=torch.manual_seed(self.seed),
-            num_inference_steps=self.num_infer_step,
-            image_guidance_scale=self.image_guide_scale,
-            guidance_scale=self.text_guide_scale,
-        ).images[0]
-
-        alpha_channel = Image.fromarray(alpha_channel)
-        rgba_image = image.convert("RGBA").resize(target_wh)
-        rgba_image.putalpha(alpha_channel)
-
-        return rgba_image
-
-
-if __name__ == "__main__":
-    delighting_model = DelightingModel(
-        # model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
-    )
-    image_path = "scripts/apps/assets/example_image/room_bottle_002.jpeg"
-    image = delighting_model(
-        image_path, preprocess=True, target_wh=(512, 512)
-    )  # noqa
-    image.save("delight.png")
-
-    # image_path = "asset3d_gen/scripts/test_robot.png"
-    # image = delighting_model(image_path)
-    # image.save("delighting_image_a2.png")

asset3d_gen/models/gs_model.py

@@ -1,540 +0,0 @@
-import logging
-import os
-import struct
-from dataclasses import dataclass, field
-from typing import Optional, Union
-
-import cv2
-import numpy as np
-import torch
-from gsplat.cuda._wrapper import spherical_harmonics
-from gsplat.rendering import rasterization
-from plyfile import PlyData
-from scipy.spatial.transform import Rotation
-from torch.nn import functional as F
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "RenderResult",
-    "GaussianOperator",
-]
-
-
-def quat_mult(q1, q2):
-    # NOTE:
-    # Q1 is the quaternion that rotates the vector from the original position to the final position  # noqa
-    # Q2 is the quaternion that been rotated
-    w1, x1, y1, z1 = q1.T
-    w2, x2, y2, z2 = q2.T
-    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
-    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
-    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
-    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
-    return torch.stack([w, x, y, z]).T
-
-
-def quat_to_rotmat(quats: torch.Tensor, mode="wxyz") -> torch.Tensor:
-    """Convert quaternion to rotation matrix."""
-    quats = F.normalize(quats, p=2, dim=-1)
-
-    if mode == "xyzw":
-        x, y, z, w = torch.unbind(quats, dim=-1)
-    elif mode == "wxyz":
-        w, x, y, z = torch.unbind(quats, dim=-1)
-    else:
-        raise ValueError(f"Invalid mode: {mode}.")
-
-    R = torch.stack(
-        [
-            1 - 2 * (y**2 + z**2),
-            2 * (x * y - w * z),
-            2 * (x * z + w * y),
-            2 * (x * y + w * z),
-            1 - 2 * (x**2 + z**2),
-            2 * (y * z - w * x),
-            2 * (x * z - w * y),
-            2 * (y * z + w * x),
-            1 - 2 * (x**2 + y**2),
-        ],
-        dim=-1,
-    )
-
-    return R.reshape(quats.shape[:-1] + (3, 3))
-
-
-def gamma_shs(shs: torch.Tensor, gamma: float) -> torch.Tensor:
-    C0 = 0.28209479177387814  # Constant for normalization in spherical harmonics  # noqa
-    # Clip to the range [0.0, 1.0], apply gamma correction, and then un-clip back  # noqa
-    new_shs = torch.clip(shs * C0 + 0.5, 0.0, 1.0)
-    new_shs = (torch.pow(new_shs, gamma) - 0.5) / C0
-    return new_shs
-
-
-@dataclass
-class RenderResult:
-    rgb: np.ndarray
-    depth: np.ndarray
-    opacity: np.ndarray
-    mask_threshold: float = 10
-    mask: Optional[np.ndarray] = None
-    rgba: Optional[np.ndarray] = None
-
-    def __post_init__(self):
-        if isinstance(self.rgb, torch.Tensor):
-            rgb = self.rgb.detach().cpu().numpy()
-            rgb = (rgb * 255).astype(np.uint8)
-            self.rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
-        if isinstance(self.depth, torch.Tensor):
-            self.depth = self.depth.detach().cpu().numpy()
-        if isinstance(self.opacity, torch.Tensor):
-            opacity = self.opacity.detach().cpu().numpy()
-            opacity = (opacity * 255).astype(np.uint8)
-            self.opacity = cv2.cvtColor(opacity, cv2.COLOR_GRAY2RGB)
-            mask = np.where(self.opacity > self.mask_threshold, 255, 0)
-            self.mask = mask[..., 0:1].astype(np.uint8)
-            self.rgba = np.concatenate([self.rgb, self.mask], axis=-1)
-
-
-@dataclass
-class GaussianBase:
-    _opacities: torch.Tensor
-    _means: torch.Tensor
-    _scales: torch.Tensor
-    _quats: torch.Tensor
-    _rgbs: Optional[torch.Tensor] = None
-    _features_dc: Optional[torch.Tensor] = None
-    _features_rest: Optional[torch.Tensor] = None
-    sh_degree: Optional[int] = 0
-    device: str = "cuda"
-
-    def __post_init__(self):
-        self.active_sh_degree: int = self.sh_degree
-        self.to(self.device)
-
-    def to(self, device: str) -> None:
-        for k, v in self.__dict__.items():
-            if not isinstance(v, torch.Tensor):
-                continue
-            self.__dict__[k] = v.to(device)
-
-    def get_numpy_data(self):
-        data = {}
-        for k, v in self.__dict__.items():
-            if not isinstance(v, torch.Tensor):
-                continue
-            data[k] = v.detach().cpu().numpy()
-
-        return data
-
-    def quat_norm(self, x: torch.Tensor) -> torch.Tensor:
-        return x / x.norm(dim=-1, keepdim=True)
-
-    @classmethod
-    def load_from_ply(
-        cls,
-        path: str,
-        gamma: float = 1.0,
-    ) -> "GaussianBase":
-        plydata = PlyData.read(path)
-        xyz = torch.stack(
-            (
-                torch.tensor(plydata.elements[0]["x"], dtype=torch.float32),
-                torch.tensor(plydata.elements[0]["y"], dtype=torch.float32),
-                torch.tensor(plydata.elements[0]["z"], dtype=torch.float32),
-            ),
-            dim=1,
-        )
-
-        opacities = torch.tensor(
-            plydata.elements[0]["opacity"], dtype=torch.float32
-        ).unsqueeze(-1)
-        features_dc = torch.zeros((xyz.shape[0], 3), dtype=torch.float32)
-        features_dc[:, 0] = torch.tensor(
-            plydata.elements[0]["f_dc_0"], dtype=torch.float32
-        )
-        features_dc[:, 1] = torch.tensor(
-            plydata.elements[0]["f_dc_1"], dtype=torch.float32
-        )
-        features_dc[:, 2] = torch.tensor(
-            plydata.elements[0]["f_dc_2"], dtype=torch.float32
-        )
-
-        scale_names = [
-            p.name
-            for p in plydata.elements[0].properties
-            if p.name.startswith("scale_")
-        ]
-        scale_names = sorted(scale_names, key=lambda x: int(x.split("_")[-1]))
-        scales = torch.zeros(
-            (xyz.shape[0], len(scale_names)), dtype=torch.float32
-        )
-        for idx, attr_name in enumerate(scale_names):
-            scales[:, idx] = torch.tensor(
-                plydata.elements[0][attr_name], dtype=torch.float32
-            )
-
-        rot_names = [
-            p.name
-            for p in plydata.elements[0].properties
-            if p.name.startswith("rot_")
-        ]
-        rot_names = sorted(rot_names, key=lambda x: int(x.split("_")[-1]))
-        rots = torch.zeros((xyz.shape[0], len(rot_names)), dtype=torch.float32)
-        for idx, attr_name in enumerate(rot_names):
-            rots[:, idx] = torch.tensor(
-                plydata.elements[0][attr_name], dtype=torch.float32
-            )
-
-        rots = rots / torch.norm(rots, dim=-1, keepdim=True)
-
-        # extra features
-        extra_f_names = [
-            p.name
-            for p in plydata.elements[0].properties
-            if p.name.startswith("f_rest_")
-        ]
-        extra_f_names = sorted(
-            extra_f_names, key=lambda x: int(x.split("_")[-1])
-        )
-
-        max_sh_degree = int(np.sqrt((len(extra_f_names) + 3) / 3) - 1)
-        if max_sh_degree != 0:
-            features_extra = torch.zeros(
-                (xyz.shape[0], len(extra_f_names)), dtype=torch.float32
-            )
-            for idx, attr_name in enumerate(extra_f_names):
-                features_extra[:, idx] = torch.tensor(
-                    plydata.elements[0][attr_name], dtype=torch.float32
-                )
-
-            features_extra = features_extra.view(
-                (features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1)
-            )
-            features_extra = features_extra.permute(0, 2, 1)
-
-            if abs(gamma - 1.0) > 1e-3:
-                features_dc = gamma_shs(features_dc, gamma)
-                features_extra[..., :] = 0.0
-                opacities *= 0.8
-
-            shs = torch.cat(
-                [
-                    features_dc.reshape(-1, 3),
-                    features_extra.reshape(len(features_dc), -1),
-                ],
-                dim=-1,
-            )
-        else:
-            # sh_dim is 0, only dc features
-            shs = features_dc
-            features_extra = None
-
-        return cls(
-            sh_degree=max_sh_degree,
-            _means=xyz,
-            _opacities=opacities,
-            _rgbs=shs,
-            _scales=scales,
-            _quats=rots,
-            _features_dc=features_dc,
-            _features_rest=features_extra,
-        )
-
-    def save_to_ply(
-        self, path: str, colors: torch.Tensor = None, enable_mask: bool = False
-    ):
-        os.makedirs(os.path.dirname(path), exist_ok=True)
-        numpy_data = self.get_numpy_data()
-        means = numpy_data["_means"]
-        scales = numpy_data["_scales"]
-        quats = numpy_data["_quats"]
-        opacities = numpy_data["_opacities"]
-        sh0 = numpy_data["_features_dc"]
-        shN = numpy_data.get("_features_rest", np.zeros((means.shape[0], 0)))
-
-        # Create a mask to identify rows with NaN or Inf in any of the numpy_data arrays  # noqa
-        if enable_mask:
-            invalid_mask = (
-                np.isnan(means).any(axis=1)
-                | np.isinf(means).any(axis=1)
-                | np.isnan(scales).any(axis=1)
-                | np.isinf(scales).any(axis=1)
-                | np.isnan(quats).any(axis=1)
-                | np.isinf(quats).any(axis=1)
-                | np.isnan(opacities).any(axis=0)
-                | np.isinf(opacities).any(axis=0)
-                | np.isnan(sh0).any(axis=1)
-                | np.isinf(sh0).any(axis=1)
-                | np.isnan(shN).any(axis=1)
-                | np.isinf(shN).any(axis=1)
-            )
-
-            # Filter out rows with NaNs or Infs from all data arrays
-            means = means[~invalid_mask]
-            scales = scales[~invalid_mask]
-            quats = quats[~invalid_mask]
-            opacities = opacities[~invalid_mask]
-            sh0 = sh0[~invalid_mask]
-            shN = shN[~invalid_mask]
-
-        num_points = means.shape[0]
-
-        with open(path, "wb") as f:
-            # Write PLY header
-            f.write(b"ply\n")
-            f.write(b"format binary_little_endian 1.0\n")
-            f.write(f"element vertex {num_points}\n".encode())
-            f.write(b"property float x\n")
-            f.write(b"property float y\n")
-            f.write(b"property float z\n")
-            f.write(b"property float nx\n")
-            f.write(b"property float ny\n")
-            f.write(b"property float nz\n")
-
-            if colors is not None:
-                for j in range(colors.shape[1]):
-                    f.write(f"property float f_dc_{j}\n".encode())
-            else:
-                for i, data in enumerate([sh0, shN]):
-                    prefix = "f_dc" if i == 0 else "f_rest"
-                    for j in range(data.shape[1]):
-                        f.write(f"property float {prefix}_{j}\n".encode())
-
-            f.write(b"property float opacity\n")
-
-            for i in range(scales.shape[1]):
-                f.write(f"property float scale_{i}\n".encode())
-            for i in range(quats.shape[1]):
-                f.write(f"property float rot_{i}\n".encode())
-
-            f.write(b"end_header\n")
-
-            # Write vertex data
-            for i in range(num_points):
-                f.write(struct.pack("<fff", *means[i]))  # x, y, z
-                f.write(struct.pack("<fff", 0, 0, 0))  # nx, ny, nz (zeros)
-
-                if colors is not None:
-                    color = colors.detach().cpu().numpy()
-                    for j in range(color.shape[1]):
-                        f_dc = (color[i, j] - 0.5) / 0.2820947917738781
-                        f.write(struct.pack("<f", f_dc))
-                else:
-                    for data in [sh0, shN]:
-                        for j in range(data.shape[1]):
-                            f.write(struct.pack("<f", data[i, j]))
-
-                f.write(struct.pack("<f", opacities[i]))  # opacity
-
-                for data in [scales, quats]:
-                    for j in range(data.shape[1]):
-                        f.write(struct.pack("<f", data[i, j]))
-
-
-@dataclass
-class GaussianOperator(GaussianBase):
-
-    def _compute_transform(
-        self,
-        means: torch.Tensor,
-        quats: torch.Tensor,
-        instance_pose: torch.Tensor,
-    ):
-        """Compute the transform of the GS models.
-
-        Args:
-            means: tensor of gs means.
-            quats: tensor of gs quaternions.
-            instance_pose: instances poses in [x y z qx qy qz qw] format.
-
-        """
-        # (x y z qx qy qz qw) -> (x y z qw qx qy qz)
-        instance_pose = instance_pose[[0, 1, 2, 6, 3, 4, 5]]
-        cur_instances_quats = self.quat_norm(instance_pose[3:])
-        rot_cur = quat_to_rotmat(cur_instances_quats, mode="wxyz")
-
-        # update the means
-        num_gs = means.shape[0]
-        trans_per_pts = torch.stack([instance_pose[:3]] * num_gs, dim=0)
-        quat_per_pts = torch.stack([instance_pose[3:]] * num_gs, dim=0)
-        rot_per_pts = torch.stack([rot_cur] * num_gs, dim=0)  # (num_gs, 3, 3)
-
-        # update the means
-        cur_means = (
-            torch.bmm(rot_per_pts, means.unsqueeze(-1)).squeeze(-1)
-            + trans_per_pts
-        )
-
-        # update the quats
-        _quats = self.quat_norm(quats)
-        cur_quats = quat_mult(quat_per_pts, _quats)
-
-        return cur_means, cur_quats
-
-    def get_gaussians(
-        self,
-        c2w: torch.Tensor = None,
-        instance_pose: torch.Tensor = None,
-        apply_activate: bool = False,
-    ) -> "GaussianBase":
-        """Get Gaussian data under the given instance_pose."""
-        if c2w is None:
-            c2w = torch.eye(4).to(self.device)
-
-        if instance_pose is not None:
-            # compute the transformed gs means and quats
-            world_means, world_quats = self._compute_transform(
-                self._means, self._quats, instance_pose.float().to(self.device)
-            )
-        else:
-            world_means, world_quats = self._means, self._quats
-
-        # get colors of gaussians
-        if self._features_rest is not None:
-            colors = torch.cat(
-                (self._features_dc[:, None, :], self._features_rest), dim=1
-            )
-        else:
-            colors = self._features_dc[:, None, :]
-
-        if self.sh_degree > 0:
-            viewdirs = world_means.detach() - c2w[..., :3, 3]  # (N, 3)
-            viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
-            rgbs = spherical_harmonics(self.sh_degree, viewdirs, colors)
-            rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0)
-        else:
-            rgbs = torch.sigmoid(colors[:, 0, :])
-
-        gs_dict = dict(
-            _means=world_means,
-            _opacities=(
-                torch.sigmoid(self._opacities)
-                if apply_activate
-                else self._opacities
-            ),
-            _rgbs=rgbs,
-            _scales=(
-                torch.exp(self._scales) if apply_activate else self._scales
-            ),
-            _quats=self.quat_norm(world_quats),
-            _features_dc=self._features_dc,
-            _features_rest=self._features_rest,
-            sh_degree=self.sh_degree,
-        )
-
-        return GaussianOperator(**gs_dict)
-
-    def rescale(self, scale: float):
-        if scale != 1.0:
-            self._means *= scale
-            self._scales += torch.log(self._scales.new_tensor(scale))
-
-    def set_scale_by_height(self, real_height: float) -> None:
-        def _ptp(tensor, dim):
-            val = tensor.max(dim=dim).values - tensor.min(dim=dim).values
-            return val.tolist()
-
-        xyz_scale = max(_ptp(self._means, dim=0))
-        self.rescale(1 / (xyz_scale + 1e-6))  # Normalize to [-0.5, 0.5]
-        raw_height = _ptp(self._means, dim=0)[1]
-        scale = real_height / raw_height
-
-        self.rescale(scale)
-
-        return
-
-    @staticmethod
-    def resave_ply(
-        in_ply: str,
-        out_ply: str,
-        real_height: float = None,
-        instance_pose: np.ndarray = None,
-        sh_degree: int = 0,
-    ) -> None:
-        gs_model = GaussianOperator.load_from_ply(in_ply, sh_degree)
-
-        if instance_pose is not None:
-            gs_model = gs_model.get_gaussians(instance_pose=instance_pose)
-
-        if real_height is not None:
-            gs_model.set_scale_by_height(real_height)
-
-        gs_model.save_to_ply(out_ply)
-
-        return
-
-    @staticmethod
-    def trans_to_quatpose(
-        rot_matrix: list[list[float]],
-        trans_matrix: list[float] = [0, 0, 0],
-    ) -> torch.Tensor:
-        if isinstance(rot_matrix, list):
-            rot_matrix = np.array(rot_matrix)
-
-        rot = Rotation.from_matrix(rot_matrix)
-        qx, qy, qz, qw = rot.as_quat()
-        instance_pose = torch.tensor([*trans_matrix, qx, qy, qz, qw])
-
-        return instance_pose
-
-    def render(
-        self,
-        c2w: torch.Tensor,
-        Ks: torch.Tensor,
-        image_width: int,
-        image_height: int,
-    ) -> RenderResult:
-        gs = self.get_gaussians(c2w, apply_activate=True)
-        renders, alphas, _ = rasterization(
-            means=gs._means,
-            quats=gs._quats,
-            scales=gs._scales,
-            opacities=gs._opacities.squeeze(),
-            colors=gs._rgbs,
-            viewmats=torch.linalg.inv(c2w)[None, ...],
-            Ks=Ks[None, ...],
-            width=image_width,
-            height=image_height,
-            packed=False,
-            absgrad=True,
-            sparse_grad=False,
-            # rasterize_mode="classic",
-            rasterize_mode="antialiased",
-            **{
-                "near_plane": 0.01,
-                "far_plane": 1000000000,
-                "radius_clip": 0.0,
-                "render_mode": "RGB+ED",
-            },
-        )
-        renders = renders[0]
-        alphas = alphas[0].squeeze(-1)
-
-        assert renders.shape[-1] == 4, f"Must render rgb, depth and alpha"
-        rendered_rgb, rendered_depth = torch.split(renders, [3, 1], dim=-1)
-
-        return RenderResult(
-            torch.clamp(rendered_rgb, min=0, max=1),
-            rendered_depth,
-            alphas[..., None],
-        )
-
-
-if __name__ == "__main__":
-    input_gs = "outputs/test/debug.ply"
-    output_gs = "./debug_v3.ply"
-    gs_model: GaussianOperator = GaussianOperator.load_from_ply(input_gs)
-
-    # 绕 x 轴旋转 180°
-    R_x = [[1, 0, 0], [0, -1, 0], [0, 0, -1]]
-    instance_pose = gs_model.trans_to_quatpose(R_x)
-    gs_model = gs_model.get_gaussians(instance_pose=instance_pose)
-
-    gs_model.rescale(2)
-
-    gs_model.set_scale_by_height(1.3)
-
-    gs_model.save_to_ply(output_gs)

asset3d_gen/models/segment.py

@@ -1,375 +0,0 @@
-import logging
-import os
-from typing import Literal, Union
-
-import cv2
-import numpy as np
-import rembg
-import torch
-from huggingface_hub import snapshot_download
-from PIL import Image
-from segment_anything import (
-    SamAutomaticMaskGenerator,
-    SamPredictor,
-    sam_model_registry,
-)
-from asset3d_gen.utils.process_media import filter_small_connected_components
-from asset3d_gen.validators.quality_checkers import ImageSegChecker
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "resize_pil",
-    "trellis_preprocess",
-    "SAMRemover",
-    "SAMPredictor",
-    "RembgRemover",
-    "get_segmented_image",
-]
-
-
-def resize_pil(image: Image.Image, max_size: int = 1024) -> Image.Image:
-    max_size = max(image.size)
-    scale = min(1, 1024 / max_size)
-    if scale < 1:
-        new_size = (int(image.width * scale), int(image.height * scale))
-        image = image.resize(new_size, Image.Resampling.LANCZOS)
-
-    return image
-
-
-def trellis_preprocess(image: Image.Image) -> Image.Image:
-    """Process the input image as trellis done."""
-    image_np = np.array(image)
-    alpha = image_np[:, :, 3]
-    bbox = np.argwhere(alpha > 0.8 * 255)
-    bbox = (
-        np.min(bbox[:, 1]),
-        np.min(bbox[:, 0]),
-        np.max(bbox[:, 1]),
-        np.max(bbox[:, 0]),
-    )
-    center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
-    size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
-    size = int(size * 1.2)
-    bbox = (
-        center[0] - size // 2,
-        center[1] - size // 2,
-        center[0] + size // 2,
-        center[1] + size // 2,
-    )
-    image = image.crop(bbox)
-    image = image.resize((518, 518), Image.Resampling.LANCZOS)
-    image = np.array(image).astype(np.float32) / 255
-    image = image[:, :, :3] * image[:, :, 3:4]
-    image = Image.fromarray((image * 255).astype(np.uint8))
-
-    return image
-
-
-class SAMRemover(object):
-    """Loading SAM models and performing background removal on images.
-
-    Attributes:
-        checkpoint (str): Path to the model checkpoint.
-        model_type (str): Type of the SAM model to load (default: "vit_h").
-        area_ratio (float): Area ratio filtering small connected components.
-    """
-
-    def __init__(
-        self,
-        checkpoint: str = None,
-        model_type: str = "vit_h",
-        area_ratio: float = 15,
-    ):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model_type = model_type
-        self.area_ratio = area_ratio
-
-        if checkpoint is None:
-            suffix = "sam"
-            model_path = snapshot_download(
-                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
-            )
-            checkpoint = os.path.join(
-                model_path, suffix, "sam_vit_h_4b8939.pth"
-            )
-
-        self.mask_generator = self._load_sam_model(checkpoint)
-
-    def _load_sam_model(self, checkpoint: str) -> SamAutomaticMaskGenerator:
-        sam = sam_model_registry[self.model_type](checkpoint=checkpoint)
-        sam.to(device=self.device)
-
-        return SamAutomaticMaskGenerator(sam)
-
-    def __call__(
-        self, image: Union[str, Image.Image, np.ndarray], save_path: str = None
-    ) -> Image.Image:
-        """Removes the background from an image using the SAM model.
-
-        Args:
-            image (Union[str, Image.Image, np.ndarray]): Input image,
-                can be a file path, PIL Image, or numpy array.
-            save_path (str): Path to save the output image (default: None).
-
-        Returns:
-            Image.Image: The image with background removed,
-                including an alpha channel.
-        """
-        # Convert input to numpy array
-        if isinstance(image, str):
-            image = Image.open(image)
-        elif isinstance(image, np.ndarray):
-            image = Image.fromarray(image).convert("RGB")
-        image = resize_pil(image)
-        image = np.array(image.convert("RGB"))
-
-        # Generate masks
-        masks = self.mask_generator.generate(image)
-        masks = sorted(masks, key=lambda x: x["area"], reverse=True)
-
-        if not masks:
-            logger.warning(
-                "Segmentation failed: No mask generated, return raw image."
-            )
-            output_image = Image.fromarray(image, mode="RGB")
-        else:
-            # Use the largest mask
-            best_mask = masks[0]["segmentation"]
-            mask = (best_mask * 255).astype(np.uint8)
-            mask = filter_small_connected_components(
-                mask, area_ratio=self.area_ratio
-            )
-            # Apply the mask to remove the background
-            background_removed = cv2.bitwise_and(image, image, mask=mask)
-            output_image = np.dstack((background_removed, mask))
-            output_image = Image.fromarray(output_image, mode="RGBA")
-
-        if save_path is not None:
-            os.makedirs(os.path.dirname(save_path), exist_ok=True)
-            output_image.save(save_path)
-
-        return output_image
-
-
-class SAMPredictor(object):
-    def __init__(
-        self,
-        checkpoint: str = None,
-        model_type: str = "vit_h",
-        binary_thresh: float = 0.1,
-        device: str = "cuda",
-    ):
-        self.device = device
-        self.model_type = model_type
-
-        if checkpoint is None:
-            suffix = "sam"
-            model_path = snapshot_download(
-                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
-            )
-            checkpoint = os.path.join(
-                model_path, suffix, "sam_vit_h_4b8939.pth"
-            )
-
-        self.predictor = self._load_sam_model(checkpoint)
-        self.binary_thresh = binary_thresh
-
-    def _load_sam_model(self, checkpoint: str) -> SamPredictor:
-        sam = sam_model_registry[self.model_type](checkpoint=checkpoint)
-        sam.to(device=self.device)
-
-        return SamPredictor(sam)
-
-    def preprocess_image(self, image: Image.Image) -> np.ndarray:
-        if isinstance(image, str):
-            image = Image.open(image)
-        elif isinstance(image, np.ndarray):
-            image = Image.fromarray(image).convert("RGB")
-
-        image = resize_pil(image)
-        image = np.array(image.convert("RGB"))
-
-        return image
-
-    def generate_masks(
-        self,
-        image: np.ndarray,
-        selected_points: list[list[int]],
-    ) -> np.ndarray:
-        if len(selected_points) == 0:
-            return []
-
-        points = (
-            torch.Tensor([p for p, _ in selected_points])
-            .to(self.predictor.device)
-            .unsqueeze(1)
-        )
-
-        labels = (
-            torch.Tensor([int(l) for _, l in selected_points])
-            .to(self.predictor.device)
-            .unsqueeze(1)
-        )
-
-        transformed_points = self.predictor.transform.apply_coords_torch(
-            points, image.shape[:2]
-        )
-
-        masks, scores, _ = self.predictor.predict_torch(
-            point_coords=transformed_points,
-            point_labels=labels,
-            multimask_output=True,
-        )
-        valid_mask = masks[:, torch.argmax(scores, dim=1)]
-        masks_pos = valid_mask[labels[:, 0] == 1, 0].cpu().detach().numpy()
-        masks_neg = valid_mask[labels[:, 0] == 0, 0].cpu().detach().numpy()
-        if len(masks_neg) == 0:
-            masks_neg = np.zeros_like(masks_pos)
-        if len(masks_pos) == 0:
-            masks_pos = np.zeros_like(masks_neg)
-        masks_neg = masks_neg.max(axis=0, keepdims=True)
-        masks_pos = masks_pos.max(axis=0, keepdims=True)
-        valid_mask = (masks_pos.astype(int) - masks_neg.astype(int)).clip(0, 1)
-
-        binary_mask = (valid_mask > self.binary_thresh).astype(np.int32)
-
-        return [(mask, f"mask_{i}") for i, mask in enumerate(binary_mask)]
-
-    def get_segmented_image(
-        self, image: np.ndarray, masks: list[tuple[np.ndarray, str]]
-    ) -> Image.Image:
-        seg_image = Image.fromarray(image, mode="RGB")
-        alpha_channel = np.zeros(
-            (seg_image.height, seg_image.width), dtype=np.uint8
-        )
-        for mask, _ in masks:
-            # Use the maximum to combine multiple masks
-            alpha_channel = np.maximum(alpha_channel, mask)
-
-        alpha_channel = np.clip(alpha_channel, 0, 1)
-        alpha_channel = (alpha_channel * 255).astype(np.uint8)
-        alpha_image = Image.fromarray(alpha_channel, mode="L")
-        r, g, b = seg_image.split()
-        seg_image = Image.merge("RGBA", (r, g, b, alpha_image))
-
-        return seg_image
-
-    def __call__(
-        self,
-        image: Union[str, Image.Image, np.ndarray],
-        selected_points: list[list[int]],
-    ) -> Image.Image:
-        image = self.preprocess_image(image)
-        self.predictor.set_image(image)
-        masks = self.generate_masks(image, selected_points)
-
-        return self.get_segmented_image(image, masks)
-
-
-class RembgRemover(object):
-    def __init__(self):
-        self.rembg_session = rembg.new_session("u2net")
-
-    def __call__(
-        self, image: Union[str, Image.Image, np.ndarray], save_path: str = None
-    ) -> Image.Image:
-        if isinstance(image, str):
-            image = Image.open(image)
-        elif isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-
-        image = resize_pil(image)
-        output_image = rembg.remove(image, session=self.rembg_session)
-
-        if save_path is not None:
-            os.makedirs(os.path.dirname(save_path), exist_ok=True)
-            output_image.save(save_path)
-
-        return output_image
-
-
-def invert_rgba_pil(
-    image: Image.Image, mask: Image.Image, save_path: str = None
-) -> Image.Image:
-    mask = (255 - np.array(mask))[..., None]
-    image_array = np.concatenate([np.array(image), mask], axis=-1)
-    inverted_image = Image.fromarray(image_array, "RGBA")
-
-    if save_path is not None:
-        os.makedirs(os.path.dirname(save_path), exist_ok=True)
-        inverted_image.save(save_path)
-
-    return inverted_image
-
-
-def get_segmented_image(
-    image: Image.Image,
-    sam_remover: SAMRemover,
-    rbg_remover: RembgRemover,
-    seg_checker: ImageSegChecker = None,
-    save_path: str = None,
-    mode: Literal["loose", "strict"] = "loose",
-) -> Image.Image:
-    def _is_valid_seg(raw_img: Image.Image, seg_img: Image.Image) -> bool:
-        if seg_checker is None:
-            return True
-        return raw_img.mode == "RGBA" and seg_checker([raw_img, seg_img])[0]
-
-    out_sam = f"{save_path}_sam.png" if save_path else None
-    out_sam_inv = f"{save_path}_sam_inv.png" if save_path else None
-    out_rbg = f"{save_path}_rbg.png" if save_path else None
-
-    seg_image = sam_remover(image, out_sam)
-    seg_image = seg_image.convert("RGBA")
-    _, _, _, alpha = seg_image.split()
-    seg_image_inv = invert_rgba_pil(image.convert("RGB"), alpha, out_sam_inv)
-    seg_image_rbg = rbg_remover(image, out_rbg)
-
-    final_image = None
-    if _is_valid_seg(image, seg_image):
-        final_image = seg_image
-    elif _is_valid_seg(image, seg_image_inv):
-        final_image = seg_image_inv
-    elif _is_valid_seg(image, seg_image_rbg):
-        logger.warning(f"Failed to segment by `SAM`, retry with `rembg`.")
-        final_image = seg_image_rbg
-    else:
-        if mode == "strict":
-            raise RuntimeError(
-                f"Failed to segment by `SAM` or `rembg`, abort."
-            )
-        logger.warning("Failed to segment by SAM or rembg, use raw image.")
-        final_image = image.convert("RGBA")
-
-    if save_path:
-        final_image.save(save_path)
-
-    final_image = trellis_preprocess(final_image)
-
-    return final_image
-
-
-if __name__ == "__main__":
-    input_image = "outputs/text2image/demo_objects/electrical/sample_0.jpg"
-    output_image = "sample_0_seg2.png"
-
-    # input_image = "outputs/text2image/tmp/coffee_machine.jpeg"
-    # output_image = "outputs/text2image/tmp/coffee_machine_seg.png"
-
-    # input_image = "outputs/text2image/tmp/bucket.jpeg"
-    # output_image = "outputs/text2image/tmp/bucket_seg.png"
-
-    remover = SAMRemover(
-        # checkpoint="/horizon-bucket/robot_lab/users/xinjie.wang/weights/sam/sam_vit_h_4b8939.pth",  # noqa
-        model_type="vit_h",
-    )
-    remover = RembgRemover()
-    # clean_image = remover(input_image)
-    # clean_image.save(output_image)
-    get_segmented_image(
-        Image.open(input_image), remover, remover, None, "./test_seg.png"
-    )

asset3d_gen/models/super_resolution.py

@@ -1,116 +0,0 @@
-import logging
-import os
-from typing import Union
-
-import numpy as np
-import torch
-from huggingface_hub import snapshot_download
-from PIL import Image
-from asset3d_gen.data.utils import get_images_from_grid
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "ImageStableSR",
-    "ImageRealESRGAN",
-]
-
-
-class ImageStableSR:
-    def __init__(
-        self,
-        model_path: str = "stabilityai/stable-diffusion-x4-upscaler",
-        device="cuda",
-    ) -> None:
-        from diffusers import StableDiffusionUpscalePipeline
-
-        self.up_pipeline_x4 = StableDiffusionUpscalePipeline.from_pretrained(
-            model_path,
-            torch_dtype=torch.float16,
-        ).to(device)
-        self.up_pipeline_x4.set_progress_bar_config(disable=True)
-        # self.up_pipeline_x4.enable_model_cpu_offload()
-
-    def __call__(
-        self,
-        image: Union[Image.Image, np.ndarray],
-        prompt: str = "",
-        infer_step: int = 20,
-    ) -> Image.Image:
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-
-        image = image.convert("RGB")
-
-        with torch.no_grad():
-            upscaled_image = self.up_pipeline_x4(
-                image=image,
-                prompt=[prompt],
-                num_inference_steps=infer_step,
-            ).images[0]
-
-        return upscaled_image
-
-
-class ImageRealESRGAN:
-    def __init__(self, outscale: int, model_path: str = None) -> None:
-        from basicsr.archs.rrdbnet_arch import RRDBNet
-        from realesrgan import RealESRGANer
-
-        self.outscale = outscale
-        model = RRDBNet(
-            num_in_ch=3,
-            num_out_ch=3,
-            num_feat=64,
-            num_block=23,
-            num_grow_ch=32,
-            scale=4,
-        )
-        if model_path is None:
-            suffix = "super_resolution"
-            model_path = snapshot_download(
-                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
-            )
-            model_path = os.path.join(
-                model_path, suffix, "RealESRGAN_x4plus.pth"
-            )
-
-        self.upsampler = RealESRGANer(
-            scale=4,
-            model_path=model_path,
-            model=model,
-            pre_pad=0,
-            half=True,
-        )
-
-    def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-
-        with torch.no_grad():
-            output, _ = self.upsampler.enhance(image, outscale=self.outscale)
-
-        return Image.fromarray(output)
-
-
-if __name__ == "__main__":
-    color_path = "outputs/texture_mesh_gen/multi_view/color_sample0.png"
-
-    # Use RealESRGAN_x4plus for x4 (512->2048) image super resolution.
-    # model_path = "/horizon-bucket/robot_lab/users/xinjie.wang/weights/super_resolution/RealESRGAN_x4plus.pth"  # noqa
-    super_model = ImageRealESRGAN(outscale=4)
-    multiviews = get_images_from_grid(color_path, img_size=512)
-    multiviews = [super_model(img.convert("RGB")) for img in multiviews]
-    for idx, img in enumerate(multiviews):
-        img.save(f"sr{idx}.png")
-
-    # # Use stable diffusion for x4 (512->2048) image super resolution.
-    # super_model = ImageStableSR()
-    # multiviews = get_images_from_grid(color_path, img_size=512)
-    # multiviews = [super_model(img) for img in multiviews]
-    # for idx, img in enumerate(multiviews):
-    #     img.save(f"sr_stable{idx}.png")

asset3d_gen/models/text_model.py

@@ -1,143 +0,0 @@
-import logging
-
-import torch
-from diffusers import (
-    AutoencoderKL,
-    EulerDiscreteScheduler,
-    UNet2DConditionModel,
-)
-from kolors.models.modeling_chatglm import ChatGLMModel
-from kolors.models.tokenization_chatglm import ChatGLMTokenizer
-from kolors.models.unet_2d_condition import (
-    UNet2DConditionModel as UNet2DConditionModelIP,
-)
-from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import (
-    StableDiffusionXLPipeline,
-)
-from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter import (  # noqa
-    StableDiffusionXLPipeline as StableDiffusionXLPipelineIP,
-)
-from PIL import Image
-from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "build_text2img_ip_pipeline",
-    "build_text2img_pipeline",
-    "text2img_gen",
-]
-
-
-def build_text2img_ip_pipeline(
-    ckpt_dir: str,
-    ref_scale: float,
-    device: str = "cuda",
-) -> StableDiffusionXLPipelineIP:
-    text_encoder = ChatGLMModel.from_pretrained(
-        f"{ckpt_dir}/text_encoder", torch_dtype=torch.float16
-    ).half()
-    tokenizer = ChatGLMTokenizer.from_pretrained(f"{ckpt_dir}/text_encoder")
-    vae = AutoencoderKL.from_pretrained(
-        f"{ckpt_dir}/vae", revision=None
-    ).half()
-    scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
-    unet = UNet2DConditionModelIP.from_pretrained(
-        f"{ckpt_dir}/unet", revision=None
-    ).half()
-    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
-        f"{ckpt_dir}/../Kolors-IP-Adapter-Plus/image_encoder",
-        ignore_mismatched_sizes=True,
-    ).to(dtype=torch.float16)
-    clip_image_processor = CLIPImageProcessor(size=336, crop_size=336)
-
-    pipe = StableDiffusionXLPipelineIP(
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        unet=unet,
-        scheduler=scheduler,
-        image_encoder=image_encoder,
-        feature_extractor=clip_image_processor,
-        force_zeros_for_empty_prompt=False,
-    )
-
-    if hasattr(pipe.unet, "encoder_hid_proj"):
-        pipe.unet.text_encoder_hid_proj = pipe.unet.encoder_hid_proj
-
-    pipe.load_ip_adapter(
-        f"{ckpt_dir}/../Kolors-IP-Adapter-Plus",
-        subfolder="",
-        weight_name=["ip_adapter_plus_general.bin"],
-    )
-    pipe.set_ip_adapter_scale([ref_scale])
-
-    pipe = pipe.to(device)
-    # pipe.enable_model_cpu_offload()
-    # pipe.enable_xformers_memory_efficient_attention()
-    # pipe.enable_vae_slicing()
-
-    return pipe
-
-
-def build_text2img_pipeline(
-    ckpt_dir: str,
-    device: str = "cuda",
-) -> StableDiffusionXLPipeline:
-    text_encoder = ChatGLMModel.from_pretrained(
-        f"{ckpt_dir}/text_encoder", torch_dtype=torch.float16
-    ).half()
-    tokenizer = ChatGLMTokenizer.from_pretrained(f"{ckpt_dir}/text_encoder")
-    vae = AutoencoderKL.from_pretrained(
-        f"{ckpt_dir}/vae", revision=None
-    ).half()
-    scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
-    unet = UNet2DConditionModel.from_pretrained(
-        f"{ckpt_dir}/unet", revision=None
-    ).half()
-    pipe = StableDiffusionXLPipeline(
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        unet=unet,
-        scheduler=scheduler,
-        force_zeros_for_empty_prompt=False,
-    )
-    pipe = pipe.to(device)
-    # pipe.enable_model_cpu_offload()
-    # pipe.enable_xformers_memory_efficient_attention()
-
-    return pipe
-
-
-def text2img_gen(
-    prompt: str,
-    n_sample: int,
-    guidance_scale: float,
-    pipeline: StableDiffusionXLPipeline | StableDiffusionXLPipelineIP,
-    ip_image: Image.Image | str = None,
-    image_wh: tuple[int, int] = [1024, 1024],
-    infer_step: int = 50,
-    ip_image_size: int = 512,
-) -> list[Image.Image]:
-    prompt = "Single " + prompt + ", in the center of the image"
-    prompt += ", high quality, high resolution, best quality, white background, 3D style,"  # noqa
-    logger.info(f"Processing prompt: {prompt}")
-
-    kwargs = dict(
-        prompt=prompt,
-        height=image_wh[1],
-        width=image_wh[0],
-        num_inference_steps=infer_step,
-        guidance_scale=guidance_scale,
-        num_images_per_prompt=n_sample,
-    )
-    if ip_image is not None:
-        if isinstance(ip_image, str):
-            ip_image = Image.open(ip_image)
-        ip_image = ip_image.resize((ip_image_size, ip_image_size))
-        kwargs.update(ip_adapter_image=[ip_image])
-
-    return pipeline(**kwargs).images

asset3d_gen/models/texture_model.py

@@ -1,91 +0,0 @@
-import os
-
-import torch
-from diffusers import AutoencoderKL, DiffusionPipeline, EulerDiscreteScheduler
-from huggingface_hub import snapshot_download
-from kolors.models.controlnet import ControlNetModel
-from kolors.models.modeling_chatglm import ChatGLMModel
-from kolors.models.tokenization_chatglm import ChatGLMTokenizer
-from kolors.models.unet_2d_condition import UNet2DConditionModel
-from kolors.pipelines.pipeline_controlnet_xl_kolors_img2img import (
-    StableDiffusionXLControlNetImg2ImgPipeline,
-)
-from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
-
-__all__ = [
-    "build_texture_gen_pipe",
-]
-
-
-def build_texture_gen_pipe(
-    base_ckpt_dir: str,
-    controlnet_ckpt: str = None,
-    ip_adapt_scale: float = 0,
-    device: str = "cuda",
-) -> DiffusionPipeline:
-    tokenizer = ChatGLMTokenizer.from_pretrained(
-        f"{base_ckpt_dir}/Kolors/text_encoder"
-    )
-    text_encoder = ChatGLMModel.from_pretrained(
-        f"{base_ckpt_dir}/Kolors/text_encoder", torch_dtype=torch.float16
-    ).half()
-    vae = AutoencoderKL.from_pretrained(
-        f"{base_ckpt_dir}/Kolors/vae", revision=None
-    ).half()
-    unet = UNet2DConditionModel.from_pretrained(
-        f"{base_ckpt_dir}/Kolors/unet", revision=None
-    ).half()
-    scheduler = EulerDiscreteScheduler.from_pretrained(
-        f"{base_ckpt_dir}/Kolors/scheduler"
-    )
-
-    if controlnet_ckpt is None:
-        suffix = "geo_cond_mv"
-        model_path = snapshot_download(
-            repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
-        )
-        controlnet_ckpt = os.path.join(model_path, suffix)
-
-    controlnet = ControlNetModel.from_pretrained(
-        controlnet_ckpt, use_safetensors=True
-    ).half()
-
-    # IP-Adapter model
-    image_encoder = None
-    clip_image_processor = None
-    if ip_adapt_scale > 0:
-        image_encoder = CLIPVisionModelWithProjection.from_pretrained(
-            f"{base_ckpt_dir}/Kolors-IP-Adapter-Plus/image_encoder",
-            # ignore_mismatched_sizes=True,
-        ).to(dtype=torch.float16)
-        ip_img_size = 336
-        clip_image_processor = CLIPImageProcessor(
-            size=ip_img_size, crop_size=ip_img_size
-        )
-
-    pipe = StableDiffusionXLControlNetImg2ImgPipeline(
-        vae=vae,
-        controlnet=controlnet,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        unet=unet,
-        scheduler=scheduler,
-        image_encoder=image_encoder,
-        feature_extractor=clip_image_processor,
-        force_zeros_for_empty_prompt=False,
-    )
-
-    if ip_adapt_scale > 0:
-        if hasattr(pipe.unet, "encoder_hid_proj"):
-            pipe.unet.text_encoder_hid_proj = pipe.unet.encoder_hid_proj
-        pipe.load_ip_adapter(
-            f"{base_ckpt_dir}/Kolors-IP-Adapter-Plus",
-            subfolder="",
-            weight_name=["ip_adapter_plus_general.bin"],
-        )
-        pipe.set_ip_adapter_scale([ip_adapt_scale])
-
-    pipe = pipe.to(device)
-    # pipe.enable_model_cpu_offload()
-
-    return pipe

asset3d_gen/scripts/render_gs.py

@@ -1,156 +0,0 @@
-import argparse
-import logging
-import math
-import os
-
-import cv2
-import numpy as np
-import torch
-from tqdm import tqdm
-from asset3d_gen.data.utils import (
-    CameraSetting,
-    init_kal_camera,
-    normalize_vertices_array,
-)
-from asset3d_gen.models.gs_model import GaussianOperator
-
-logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
-)
-logger = logging.getLogger(__name__)
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Render GS color images")
-
-    parser.add_argument(
-        "--input_gs", type=str, help="Input render GS.ply path."
-    )
-    parser.add_argument(
-        "--output_path",
-        type=str,
-        help="Output grid image path for rendered GS color images.",
-    )
-    parser.add_argument(
-        "--num_images", type=int, default=6, help="Number of images to render."
-    )
-    parser.add_argument(
-        "--elevation",
-        type=float,
-        nargs="+",
-        default=[20.0, -10.0],
-        help="Elevation angles for the camera (default: [20.0, -10.0])",
-    )
-    parser.add_argument(
-        "--distance",
-        type=float,
-        default=5,
-        help="Camera distance (default: 5)",
-    )
-    parser.add_argument(
-        "--resolution_hw",
-        type=int,
-        nargs=2,
-        default=(512, 512),
-        help="Resolution of the output images (default: (512, 512))",
-    )
-    parser.add_argument(
-        "--fov",
-        type=float,
-        default=30,
-        help="Field of view in degrees (default: 30)",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        choices=["cpu", "cuda"],
-        default="cuda",
-        help="Device to run on (default: `cuda`)",
-    )
-    parser.add_argument(
-        "--image_size",
-        type=int,
-        default=512,
-        help="Output image size for single view in color grid (default: 512)",
-    )
-
-    args = parser.parse_args()
-
-    return args
-
-
-def load_gs_model(
-    input_gs: str, pre_quat: list[float] = [0.0, 0.7071, 0.0, -0.7071]
-) -> GaussianOperator:
-    gs_model = GaussianOperator.load_from_ply(input_gs)
-    # Normalize vertices to [-1, 1], center to (0, 0, 0).
-    _, scale, center = normalize_vertices_array(gs_model._means)
-    scale, center = float(scale), center.tolist()
-    transpose = [*[-v for v in center], *pre_quat]
-    instance_pose = torch.tensor(transpose).to(gs_model.device)
-    gs_model = gs_model.get_gaussians(instance_pose=instance_pose)
-    gs_model.rescale(scale)
-
-    return gs_model
-
-
-def entrypoint(input_gs: str = None, output_path: str = None) -> None:
-    args = parse_args()
-    if isinstance(input_gs, str):
-        args.input_gs = input_gs
-    if isinstance(output_path, str):
-        args.output_path = output_path
-
-    # Setup camera parameters
-    camera_params = CameraSetting(
-        num_images=args.num_images,
-        elevation=args.elevation,
-        distance=args.distance,
-        resolution_hw=args.resolution_hw,
-        fov=math.radians(args.fov),
-        device=args.device,
-    )
-    camera = init_kal_camera(camera_params)
-    matrix_mv = camera.view_matrix()  # (n_cam 4 4) world2cam
-    matrix_mv[:, :3, 3] = -matrix_mv[:, :3, 3]
-    w2cs = matrix_mv.to(camera_params.device)
-    c2ws = [torch.linalg.inv(matrix) for matrix in w2cs]
-    Ks = torch.tensor(camera_params.Ks).to(camera_params.device)
-
-    # Load GS model and normalize.
-    gs_model = load_gs_model(args.input_gs, pre_quat=[0.0, 0.0, 1.0, 0.0])
-
-    # Render GS color images.
-    images = []
-    for idx in tqdm(range(len(c2ws)), desc="Rendering GS"):
-        result = gs_model.render(
-            c2ws[idx],
-            Ks=Ks,
-            image_width=camera_params.resolution_hw[1],
-            image_height=camera_params.resolution_hw[0],
-        )
-        color = cv2.resize(
-            result.rgba,
-            (args.image_size, args.image_size),
-            interpolation=cv2.INTER_AREA,
-        )
-        images.append(color)
-
-    # Cat color images into grid image and save.
-    select_idxs = [[0, 2, 1], [5, 4, 3]]  # fix order for 6 views
-    grid_image = []
-    for row_idxs in select_idxs:
-        row_image = []
-        for row_idx in row_idxs:
-            row_image.append(images[row_idx])
-        row_image = np.concatenate(row_image, axis=1)
-        grid_image.append(row_image)
-
-    grid_image = np.concatenate(grid_image, axis=0)
-    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
-    cv2.imwrite(args.output_path, grid_image)
-    logger.info(f"Saved grid image to {args.output_path}")
-
-
-if __name__ == "__main__":
-    entrypoint()

asset3d_gen/scripts/render_mv.py

@@ -1,183 +0,0 @@
-import logging
-import os
-import random
-from typing import List, Tuple
-
-import fire
-import numpy as np
-import torch
-from diffusers.utils import make_image_grid
-from kolors.pipelines.pipeline_controlnet_xl_kolors_img2img import (
-    StableDiffusionXLControlNetImg2ImgPipeline,
-)
-from PIL import Image, ImageEnhance, ImageFilter
-from torchvision import transforms
-from asset3d_gen.data.datasets import Asset3dGenDataset
-from asset3d_gen.models.texture_model import build_texture_gen_pipe
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-def get_init_noise_image(image: Image.Image) -> Image.Image:
-    blurred_image = image.convert("L").filter(
-        ImageFilter.GaussianBlur(radius=3)
-    )
-
-    enhancer = ImageEnhance.Contrast(blurred_image)
-    image_decreased_contrast = enhancer.enhance(factor=0.5)
-
-    return image_decreased_contrast
-
-
-def infer_pipe(
-    index_file: str,
-    controlnet_ckpt: str = None,
-    uid: str = None,
-    prompt: str = None,
-    controlnet_cond_scale: float = 0.4,
-    control_guidance_end: float = 0.9,
-    strength: float = 1.0,
-    num_inference_steps: int = 50,
-    guidance_scale: float = 10,
-    ip_adapt_scale: float = 0,
-    ip_img_path: str = None,
-    sub_idxs: List[List[int]] = None,
-    num_images_per_prompt: int = 3,  # increase if want similar images.
-    device: str = "cuda",
-    save_dir: str = "infer_vis",
-    seed: int = None,
-    target_hw: tuple[int, int] = (512, 512),
-    pipeline: StableDiffusionXLControlNetImg2ImgPipeline = None,
-) -> str:
-    # sub_idxs = [[0, 1, 2], [3, 4, 5]] # None for single image.
-    if sub_idxs is None:
-        sub_idxs = [[random.randint(0, 5)]]  # 6 views.
-        target_hw = [2 * size for size in target_hw]
-
-    transform_list = [
-        transforms.Resize(
-            target_hw, interpolation=transforms.InterpolationMode.BILINEAR
-        ),
-        transforms.CenterCrop(target_hw),
-        transforms.ToTensor(),
-        transforms.Normalize([0.5], [0.5]),
-    ]
-    image_transform = transforms.Compose(transform_list)
-    control_transform = transforms.Compose(transform_list[:-1])
-
-    grid_hw = (target_hw[0] * len(sub_idxs), target_hw[1] * len(sub_idxs[0]))
-    dataset = Asset3dGenDataset(
-        index_file, target_hw=grid_hw, sub_idxs=sub_idxs
-    )
-
-    if uid is None:
-        uid = random.choice(list(dataset.meta_info.keys()))
-    if prompt is None:
-        prompt = dataset.meta_info[uid]["capture"]
-    if isinstance(prompt, List) or isinstance(prompt, Tuple):
-        prompt = ", ".join(map(str, prompt))
-    # prompt += "high quality, ultra-clear, high resolution, best quality, 4k"
-    # prompt += "高品质,清晰,细节"
-    prompt += ", high quality, high resolution, best quality"
-    # prompt += ", with diffuse lighting, showing no reflections."
-    logger.info(f"Inference with prompt: {prompt}")
-
-    negative_prompt = (
-        "nsfw,脸部阴影,低分辨率,jpeg伪影、模糊、糟糕,黑脸,霓虹灯,高光,镜面反射"
-    )
-
-    control_image = dataset.fetch_sample_grid_images(
-        uid,
-        attrs=["image_view_normal", "image_position", "image_mask"],
-        sub_idxs=sub_idxs,
-        transform=control_transform,
-    )
-
-    color_image = dataset.fetch_sample_grid_images(
-        uid,
-        attrs=["image_color"],
-        sub_idxs=sub_idxs,
-        transform=image_transform,
-    )
-
-    normal_pil, position_pil, mask_pil, color_pil = dataset.visualize_item(
-        control_image,
-        color_image,
-        save_dir=save_dir,
-    )
-
-    if pipeline is None:
-        pipeline = build_texture_gen_pipe(
-            base_ckpt_dir="./weights",
-            controlnet_ckpt=controlnet_ckpt,
-            ip_adapt_scale=ip_adapt_scale,
-            device=device,
-        )
-
-    if ip_adapt_scale > 0 and ip_img_path is not None and len(ip_img_path) > 0:
-        ip_image = Image.open(ip_img_path).convert("RGB")
-        ip_image = ip_image.resize(target_hw[::-1])
-        ip_image = [ip_image]
-        pipeline.set_ip_adapter_scale([ip_adapt_scale])
-    else:
-        ip_image = None
-
-    generator = None
-    if seed is not None:
-        generator = torch.Generator(device).manual_seed(seed)
-        torch.manual_seed(seed)
-        np.random.seed(seed)
-        random.seed(seed)
-
-    init_image = get_init_noise_image(normal_pil)
-    # init_image = get_init_noise_image(color_pil)
-
-    images = []
-    row_num, col_num = 2, 3
-    img_save_paths = []
-    while len(images) < col_num:
-        image = pipeline(
-            prompt=prompt,
-            image=init_image,
-            controlnet_conditioning_scale=controlnet_cond_scale,
-            control_guidance_end=control_guidance_end,
-            strength=strength,
-            control_image=control_image[None, ...],
-            negative_prompt=negative_prompt,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            num_images_per_prompt=num_images_per_prompt,
-            ip_adapter_image=ip_image,
-            generator=generator,
-        ).images
-        images.extend(image)
-
-    grid_image = [normal_pil, position_pil, color_pil] + images[:col_num]
-    # save_dir = os.path.join(save_dir, uid)
-    os.makedirs(save_dir, exist_ok=True)
-
-    for idx in range(col_num):
-        rgba_image = Image.merge("RGBA", (*images[idx].split(), mask_pil))
-        img_save_path = os.path.join(save_dir, f"color_sample{idx}.png")
-        rgba_image.save(img_save_path)
-        img_save_paths.append(img_save_path)
-
-    sub_idxs = "_".join(
-        [str(item) for sublist in sub_idxs for item in sublist]
-    )
-    save_path = os.path.join(
-        save_dir, f"sample_idx{str(sub_idxs)}_ip{ip_adapt_scale}.jpg"
-    )
-    make_image_grid(grid_image, row_num, col_num).save(save_path)
-    logger.info(f"Visualize in {save_path}")
-
-    return img_save_paths
-
-
-def entrypoint() -> None:
-    fire.Fire(infer_pipe)
-
-
-if __name__ == "__main__":
-    entrypoint()

asset3d_gen/scripts/text2image.py

@@ -1,145 +0,0 @@
-import argparse
-import logging
-import os
-
-from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import (
-    StableDiffusionXLPipeline,
-)
-from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter import (  # noqa
-    StableDiffusionXLPipeline as StableDiffusionXLPipelineIP,
-)
-from tqdm import tqdm
-from asset3d_gen.models.text_model import (
-    build_text2img_ip_pipeline,
-    build_text2img_pipeline,
-    text2img_gen,
-)
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Text to Image.")
-    parser.add_argument(
-        "--prompts",
-        type=str,
-        nargs="+",
-        help="List of prompts (space-separated).",
-    )
-    parser.add_argument(
-        "--ref_image",
-        type=str,
-        nargs="+",
-        help="List of ref_image paths (space-separated).",
-    )
-    parser.add_argument(
-        "--output_root",
-        type=str,
-        help="Root directory for saving outputs.",
-    )
-    parser.add_argument(
-        "--guidance_scale",
-        type=float,
-        default=12.0,
-        help="Guidance scale for the diffusion model.",
-    )
-    parser.add_argument(
-        "--ref_scale",
-        type=float,
-        default=0.3,
-        help="Reference image scale for the IP adapter.",
-    )
-    parser.add_argument(
-        "--n_sample",
-        type=int,
-        default=1,
-    )
-    parser.add_argument(
-        "--resolution",
-        type=int,
-        default=1024,
-    )
-    parser.add_argument(
-        "--infer_step",
-        type=int,
-        default=50,
-    )
-    args = parser.parse_args()
-
-    return args
-
-
-def entrypoint(
-    pipeline: StableDiffusionXLPipeline | StableDiffusionXLPipelineIP = None,
-    **kwargs,
-) -> list[str]:
-    args = parse_args()
-    for k, v in kwargs.items():
-        if hasattr(args, k) and v is not None:
-            setattr(args, k, v)
-
-    prompts = args.prompts
-    if len(prompts) == 1 and prompts[0].endswith(".txt"):
-        with open(prompts[0], "r") as f:
-            prompts = f.readlines()
-            prompts = [
-                prompt.strip() for prompt in prompts if prompt.strip() != ""
-            ]
-
-    os.makedirs(args.output_root, exist_ok=True)
-
-    ip_img_paths = args.ref_image
-    if ip_img_paths is None or len(ip_img_paths) == 0:
-        args.ref_scale = 0
-        ip_img_paths = [None] * len(prompts)
-    elif isinstance(ip_img_paths, str):
-        ip_img_paths = [ip_img_paths] * len(prompts)
-    elif isinstance(ip_img_paths, list):
-        if len(ip_img_paths) == 1:
-            ip_img_paths = ip_img_paths * len(prompts)
-    else:
-        raise ValueError("Invalid ref_image paths.")
-    assert len(ip_img_paths) == len(
-        prompts
-    ), f"Number of ref images does not match prompts, {len(ip_img_paths)} != {len(prompts)}"  # noqa
-
-    if pipeline is None:
-        if args.ref_scale > 0:
-            pipeline = build_text2img_ip_pipeline(
-                "weights/Kolors",
-                ref_scale=args.ref_scale,
-            )
-        else:
-            pipeline = build_text2img_pipeline("weights/Kolors")
-
-    for idx, (prompt, ip_img_path) in tqdm(
-        enumerate(zip(prompts, ip_img_paths)),
-        desc="Generating images",
-        total=len(prompts),
-    ):
-        images = text2img_gen(
-            prompt=prompt,
-            n_sample=args.n_sample,
-            guidance_scale=args.guidance_scale,
-            pipeline=pipeline,
-            ip_image=ip_img_path,
-            image_wh=[args.resolution, args.resolution],
-            infer_step=args.infer_step,
-        )
-
-        save_paths = []
-        for sub_idx, image in enumerate(images):
-            save_path = (
-                f"{args.output_root}/sample_{idx*args.n_sample+sub_idx}.png"
-            )
-            image.save(save_path)
-            save_paths.append(save_path)
-
-        logger.info(f"Images saved to {args.output_root}")
-
-    return save_paths
-
-
-if __name__ == "__main__":
-    entrypoint()

asset3d_gen/utils/gpt_clients.py

@@ -1,191 +0,0 @@
-import base64
-import logging
-import os
-from io import BytesIO
-from typing import Optional
-
-from openai import AzureOpenAI, OpenAI  # pip install openai
-from PIL import Image
-from tenacity import (
-    retry,
-    stop_after_attempt,
-    stop_after_delay,
-    wait_random_exponential,
-)
-from asset3d_gen.utils.process_media import combine_images_to_base64
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class GPTclient:
-    """A client to interact with the GPT model via OpenAI or Azure API."""
-
-    def __init__(
-        self,
-        endpoint: str,
-        api_key: str,
-        model_name: str = "yfb-gpt-4o",
-        api_version: str = None,
-        verbose: bool = False,
-    ):
-        if api_version is not None:
-            self.client = AzureOpenAI(
-                azure_endpoint=endpoint,
-                api_key=api_key,
-                api_version=api_version,
-            )
-        else:
-            self.client = OpenAI(
-                base_url=endpoint,
-                api_key=api_key,
-            )
-
-        self.endpoint = endpoint
-        self.model_name = model_name
-        self.image_formats = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}
-        self.verbose = verbose
-
-    @retry(
-        wait=wait_random_exponential(min=1, max=20),
-        stop=(stop_after_attempt(10) | stop_after_delay(30)),
-    )
-    def completion_with_backoff(self, **kwargs):
-        return self.client.chat.completions.create(**kwargs)
-
-    def query(
-        self,
-        text_prompt: str,
-        image_base64: Optional[list[str | Image.Image]] = None,
-        system_role: Optional[str] = None,
-    ) -> Optional[str]:
-        """Queries the GPT model with a text and optional image prompts.
-
-        Args:
-            text_prompt (str): The main text input that the model responds to.
-            image_base64 (Optional[List[str]]): A list of image base64 strings
-                or local image paths or PIL.Image to accompany the text prompt.
-            system_role (Optional[str]): Optional system-level instructions
-                that specify the behavior of the assistant.
-
-        Returns:
-            Optional[str]: The response content generated by the model based on
-                the prompt. Returns `None` if an error occurs.
-        """
-        if system_role is None:
-            system_role = "You are a highly knowledgeable assistant specializing in physics, engineering, and object properties."  # noqa
-
-        content_user = [
-            {
-                "type": "text",
-                "text": text_prompt,
-            },
-        ]
-
-        # Process images if provided
-        if image_base64 is not None:
-            image_base64 = (
-                image_base64
-                if isinstance(image_base64, list)
-                else [image_base64]
-            )
-            for img in image_base64:
-                if isinstance(img, Image.Image):
-                    buffer = BytesIO()
-                    img.save(buffer, format=img.format or "PNG")
-                    buffer.seek(0)
-                    image_binary = buffer.read()
-                    img = base64.b64encode(image_binary).decode("utf-8")
-                elif (
-                    len(os.path.splitext(img)) > 1
-                    and os.path.splitext(img)[-1].lower() in self.image_formats
-                ):
-                    if not os.path.exists(img):
-                        raise FileNotFoundError(f"Image file not found: {img}")
-                    with open(img, "rb") as f:
-                        img = base64.b64encode(f.read()).decode("utf-8")
-
-                content_user.append(
-                    {
-                        "type": "image_url",
-                        "image_url": {"url": f"data:image/png;base64,{img}"},
-                    }
-                )
-
-        payload = {
-            "messages": [
-                {"role": "system", "content": system_role},
-                {"role": "user", "content": content_user},
-            ],
-            "temperature": 0.1,
-            "max_tokens": 500,
-            "top_p": 0.1,
-            "frequency_penalty": 0,
-            "presence_penalty": 0,
-            "stop": None,
-        }
-        payload.update({"model": self.model_name})
-
-        response = None
-        try:
-            response = self.completion_with_backoff(**payload)
-            response = response.choices[0].message.content
-        except Exception as e:
-            logger.error(f"Error GPTclint {self.endpoint} API call: {e}")
-            response = None
-
-        if self.verbose:
-            logger.info(f"Prompt: {text_prompt}")
-            logger.info(f"Response: {response}")
-
-        return response
-
-
-endpoint = os.environ.get("endpoint", None)
-api_key = os.environ.get("api_key", None)
-api_version = os.environ.get("api_version", None)
-if endpoint and api_key and api_version:
-    GPT_CLIENT = GPTclient(
-        endpoint=endpoint,
-        api_key=api_key,
-        api_version=api_version,
-        model_name="yfb-gpt-4o-sweden" if "sweden" in endpoint else None,
-    )
-else:
-    GPT_CLIENT = GPTclient(
-        endpoint="https://openrouter.ai/api/v1",
-        api_key="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",  # noqa
-        model_name="qwen/qwen2.5-vl-72b-instruct:free",
-    )
-
-
-if __name__ == "__main__":
-    if "openrouter" in GPT_CLIENT.endpoint:
-        response = GPT_CLIENT.query(
-            text_prompt="What is the content in each image?",
-            image_base64=combine_images_to_base64(
-                [
-                    "outputs/text2image/demo_objects/bed/sample_0.jpg",
-                    "outputs/imageto3d/v2/cups/sample_69/URDF_sample_69/qa_renders/image_color/003.png",  # noqa
-                    "outputs/text2image/demo_objects/cardboard/sample_1.jpg",
-                ]
-            ),  # input raw image_path if only one image
-        )
-        print(response)
-    else:
-        response = GPT_CLIENT.query(
-            text_prompt="What is the content in the images?",
-            image_base64=[
-                Image.open("outputs/text2image/demo_objects/bed/sample_0.jpg"),
-                Image.open(
-                    "outputs/imageto3d/v2/cups/sample_69/URDF_sample_69/qa_renders/image_color/003.png"  # noqa
-                ),
-            ],
-        )
-        print(response)
-
-        # test2: text prompt
-        response = GPT_CLIENT.query(
-            text_prompt="What is the capital of China?"
-        )
-        print(response)

asset3d_gen/utils/process_media.py

@@ -1,194 +0,0 @@
-import base64
-import logging
-import math
-import os
-import subprocess
-from glob import glob
-from io import BytesIO
-from typing import Union
-
-import cv2
-import imageio
-import numpy as np
-import PIL.Image as Image
-from moviepy.editor import VideoFileClip, clips_array
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-__all__ = [
-    "render_asset3d",
-    "merge_images_video",
-    "filter_small_connected_components",
-    "filter_image_small_connected_components",
-    "combine_images_to_base64",
-]
-
-
-def render_asset3d(
-    mesh_path: str,
-    output_root: str,
-    distance: float = 5.0,
-    num_images: int = 1,
-    elevation: list[float] = (0.0,),
-    pbr_light_factor: float = 1.5,
-    return_key: str = "image_color/*",
-    output_subdir: str = "renders",
-    gen_color_mp4: bool = False,
-    gen_viewnormal_mp4: bool = False,
-    gen_glonormal_mp4: bool = False,
-    device: str = "cpu",
-) -> list[str]:
-    command = [
-        "python3",
-        "asset3d_gen/data/differentiable_render.py",
-        "--mesh_path",
-        mesh_path,
-        "--output_root",
-        output_root,
-        "--uuid",
-        output_subdir,
-        "--distance",
-        str(distance),
-        "--num_images",
-        str(num_images),
-        "--elevation",
-        *map(str, elevation),
-        "--pbr_light_factor",
-        str(pbr_light_factor),
-        "--with_mtl",
-        "--device",
-        device,
-    ]
-    if gen_color_mp4:
-        command.append("--gen_color_mp4")
-    if gen_viewnormal_mp4:
-        command.append("--gen_viewnormal_mp4")
-    if gen_glonormal_mp4:
-        command.append("--gen_glonormal_mp4")
-    try:
-        subprocess.run(command, check=True)
-    except subprocess.CalledProcessError as e:
-        logger.error(f"Error occurred during rendering: {e}.")
-
-    dst_paths = glob(os.path.join(output_root, output_subdir, return_key))
-
-    return dst_paths
-
-
-def merge_images_video(color_images, normal_images, output_path) -> None:
-    width = color_images[0].shape[1]
-    combined_video = [
-        np.hstack([rgb_img[:, : width // 2], normal_img[:, width // 2 :]])
-        for rgb_img, normal_img in zip(color_images, normal_images)
-    ]
-    imageio.mimsave(output_path, combined_video, fps=50)
-
-    return
-
-
-def merge_video_video(
-    video_path1: str, video_path2: str, output_path: str
-) -> None:
-    """Merge two videos by the left half and the right half of the videos."""
-    clip1 = VideoFileClip(video_path1)
-    clip2 = VideoFileClip(video_path2)
-
-    if clip1.size != clip2.size:
-        raise ValueError("The resolutions of the two videos do not match.")
-
-    width, height = clip1.size
-    clip1_half = clip1.crop(x1=0, y1=0, x2=width // 2, y2=height)
-    clip2_half = clip2.crop(x1=width // 2, y1=0, x2=width, y2=height)
-    final_clip = clips_array([[clip1_half, clip2_half]])
-    final_clip.write_videofile(output_path, codec="libx264")
-
-
-def filter_small_connected_components(
-    mask: Union[Image.Image, np.ndarray],
-    area_ratio: float,
-    connectivity: int = 8,
-) -> np.ndarray:
-    if isinstance(mask, Image.Image):
-        mask = np.array(mask)
-    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
-        mask,
-        connectivity=connectivity,
-    )
-
-    small_components = np.zeros_like(mask, dtype=np.uint8)
-    mask_area = (mask != 0).sum()
-    min_area = mask_area // area_ratio
-    for label in range(1, num_labels):
-        area = stats[label, cv2.CC_STAT_AREA]
-        if area < min_area:
-            small_components[labels == label] = 255
-
-    mask = cv2.bitwise_and(mask, cv2.bitwise_not(small_components))
-
-    return mask
-
-
-def filter_image_small_connected_components(
-    image: Union[Image.Image, np.ndarray],
-    area_ratio: float = 10,
-    connectivity: int = 8,
-) -> np.ndarray:
-    if isinstance(image, Image.Image):
-        image = image.convert("RGBA")
-        image = np.array(image)
-
-    mask = image[..., 3]
-    mask = filter_small_connected_components(mask, area_ratio, connectivity)
-    image[..., 3] = mask
-
-    return image
-
-
-def combine_images_to_base64(
-    images: list[str | Image.Image],
-    cat_row_col: tuple[int, int] = None,
-    target_wh: tuple[int, int] = (512, 512),
-) -> str:
-    n_images = len(images)
-    if cat_row_col is None:
-        n_col = math.ceil(math.sqrt(n_images))
-        n_row = math.ceil(n_images / n_col)
-    else:
-        n_row, n_col = cat_row_col
-
-    images = [
-        Image.open(p).convert("RGB") if isinstance(p, str) else p
-        for p in images[: n_row * n_col]
-    ]
-    images = [img.resize(target_wh) for img in images]
-
-    grid_w, grid_h = n_col * target_wh[0], n_row * target_wh[1]
-    grid = Image.new("RGB", (grid_w, grid_h), (255, 255, 255))
-
-    for idx, img in enumerate(images):
-        row, col = divmod(idx, n_col)
-        grid.paste(img, (col * target_wh[0], row * target_wh[1]))
-
-    buffer = BytesIO()
-    grid.save(buffer, format="PNG")
-
-    return base64.b64encode(buffer.getvalue()).decode("utf-8")
-
-
-if __name__ == "__main__":
-    # Example usage:
-    merge_video_video(
-        "outputs/imageto3d/room_bottle7/room_bottle_007/URDF_room_bottle_007/mesh_glo_normal.mp4",  # noqa
-        "outputs/imageto3d/room_bottle7/room_bottle_007/URDF_room_bottle_007/mesh.mp4",  # noqa
-        "merge.mp4",
-    )
-
-    image_base64 = combine_images_to_base64(
-        [
-            "outputs/text2image/demo_objects/bed/sample_0.jpg",
-            "outputs/imageto3d/v2/cups/sample_69/URDF_sample_69/qa_renders/image_color/003.png",  # noqa
-            "outputs/text2image/demo_objects/cardboard/sample_1.jpg",
-        ]
-    )

asset3d_gen/utils/tags.py

@@ -1 +0,0 @@
-VERSION = "v0.0.2"

asset3d_gen/validators/aesthetic_predictor.py

@@ -1,134 +0,0 @@
-import os
-
-import clip
-import numpy as np
-import pytorch_lightning as pl
-import torch
-import torch.nn as nn
-from huggingface_hub import snapshot_download
-from PIL import Image
-
-
-class AestheticPredictor:
-    """Aesthetic Score Predictor.
-
-    Args:
-        clip_model_dir (str): Path to the directory of the CLIP model.
-        sac_model_path (str): Path to the pre-trained SAC model.
-        device (str): Device to use for computation ("cuda" or "cpu").
-    """
-
-    def __init__(self, clip_model_dir=None, sac_model_path=None, device=None):
-
-        self.device = device or (
-            "cuda" if torch.cuda.is_available() else "cpu"
-        )
-
-        if clip_model_dir is None:
-            model_path = snapshot_download(
-                repo_id="xinjjj/RoboAssetGen", allow_patterns="aesthetic/*"
-            )
-            suffix = "aesthetic"
-            model_path = snapshot_download(
-                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
-            )
-            clip_model_dir = os.path.join(model_path, suffix)
-
-        if sac_model_path is None:
-            model_path = snapshot_download(
-                repo_id="xinjjj/RoboAssetGen", allow_patterns="aesthetic/*"
-            )
-            suffix = "aesthetic"
-            model_path = snapshot_download(
-                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
-            )
-            sac_model_path = os.path.join(
-                model_path, suffix, "sac+logos+ava1-l14-linearMSE.pth"
-            )
-
-        self.clip_model, self.preprocess = self._load_clip_model(
-            clip_model_dir
-        )
-        self.sac_model = self._load_sac_model(sac_model_path, input_size=768)
-
-    class MLP(pl.LightningModule):  # noqa
-        def __init__(self, input_size):
-            super().__init__()
-            self.layers = nn.Sequential(
-                nn.Linear(input_size, 1024),
-                nn.Dropout(0.2),
-                nn.Linear(1024, 128),
-                nn.Dropout(0.2),
-                nn.Linear(128, 64),
-                nn.Dropout(0.1),
-                nn.Linear(64, 16),
-                nn.Linear(16, 1),
-            )
-
-        def forward(self, x):
-            return self.layers(x)
-
-    @staticmethod
-    def normalized(a, axis=-1, order=2):
-        """Normalize the array to unit norm."""
-        l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
-        l2[l2 == 0] = 1
-        return a / np.expand_dims(l2, axis)
-
-    def _load_clip_model(self, model_dir: str, model_name: str = "ViT-L/14"):
-        """Load the CLIP model."""
-        model, preprocess = clip.load(
-            model_name, download_root=model_dir, device=self.device
-        )
-        return model, preprocess
-
-    def _load_sac_model(self, model_path, input_size):
-        """Load the SAC model."""
-        model = self.MLP(input_size)
-        ckpt = torch.load(model_path)
-        model.load_state_dict(ckpt)
-        model.to(self.device)
-        model.eval()
-        return model
-
-    def predict(self, image_path):
-        """Predict the aesthetic score for a given image.
-
-        Args:
-            image_path (str): Path to the image file.
-
-        Returns:
-            float: Predicted aesthetic score.
-        """
-        pil_image = Image.open(image_path)
-        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
-
-        with torch.no_grad():
-            # Extract CLIP features
-            image_features = self.clip_model.encode_image(image)
-            # Normalize features
-            normalized_features = self.normalized(
-                image_features.cpu().detach().numpy()
-            )
-            # Predict score
-            prediction = self.sac_model(
-                torch.from_numpy(normalized_features)
-                .type(torch.FloatTensor)
-                .to(self.device)
-            )
-
-        return prediction.item()
-
-
-if __name__ == "__main__":
-    # Configuration
-    img_path = "/home/users/xinjie.wang/xinjie/asset3d-gen/outputs/imageto3d/demo_objects/bed/sample_0/sample_0_raw.png"  # noqa
-    # clip_model_dir = "/horizon-bucket/robot_lab/users/xinjie.wang/weights/clip"  # noqa
-    # sac_model_path = "/horizon-bucket/robot_lab/users/xinjie.wang/weights/sac/sac+logos+ava1-l14-linearMSE.pth"  # noqa
-
-    # Initialize the predictor
-    predictor = AestheticPredictor()
-
-    # Predict the aesthetic score
-    score = predictor.predict(img_path)
-    print("Aesthetic score predicted by the model:", score)

asset3d_gen/validators/quality_checkers.py

@@ -1,195 +0,0 @@
-import logging
-import os
-
-from tqdm import tqdm
-from asset3d_gen.utils.gpt_clients import GPT_CLIENT, GPTclient
-from asset3d_gen.utils.process_media import render_asset3d
-from asset3d_gen.validators.aesthetic_predictor import AestheticPredictor
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class BaseChecker:
-    def __init__(self, prompt: str = None, verbose: bool = False) -> None:
-        self.prompt = prompt
-        self.verbose = verbose
-
-    def query(self, *args, **kwargs):
-        raise NotImplementedError(
-            "Subclasses must implement the query method."
-        )
-
-    def __call__(self, *args, **kwargs) -> bool:
-        response = self.query(*args, **kwargs)
-        if response is None:
-            response = "Error when calling gpt api."
-
-        if self.verbose and response != "YES":
-            logger.info(response)
-
-        flag = "YES" in response
-        response = "YES" if flag else response
-
-        return flag, response
-
-    @staticmethod
-    def validate(
-        checkers: list["BaseChecker"], images_list: list[list[str]]
-    ) -> list:
-        assert len(checkers) == len(images_list)
-        results = []
-        overall_result = True
-        for checker, images in zip(checkers, images_list):
-            qa_flag, qa_info = checker(images)
-            if isinstance(qa_info, str):
-                qa_info = qa_info.replace("\n", ".")
-            results.append([checker.__class__.__name__, qa_info])
-            if qa_flag is False:
-                overall_result = False
-
-        results.append(["overall", "YES" if overall_result else "NO"])
-
-        return results
-
-
-class MeshGeoChecker(BaseChecker):
-    def __init__(
-        self,
-        gpt_client: GPTclient,
-        prompt: str = None,
-        verbose: bool = False,
-    ) -> None:
-        super().__init__(prompt, verbose)
-        self.gpt_client = gpt_client
-        if self.prompt is None:
-            self.prompt = """
-            Refer to the provided multi-view rendering images to evaluate
-            whether the geometry of the 3D object asset is complete and
-            whether the asset can be placed stably on the ground.
-            Return "YES" only if reach the requirments,
-            otherwise "NO" and explain the reason very briefly.
-            """
-
-    def query(self, image_paths: str) -> str:
-        # Hardcode tmp because of the openrouter can't input multi images.
-        if "openrouter" in self.gpt_client.endpoint:
-            from asset3d_gen.utils.process_media import (
-                combine_images_to_base64,
-            )
-
-            image_paths = combine_images_to_base64(image_paths)
-
-        return self.gpt_client.query(
-            text_prompt=self.prompt,
-            image_base64=image_paths,
-        )
-
-
-class ImageSegChecker(BaseChecker):
-    def __init__(
-        self,
-        gpt_client: GPTclient,
-        prompt: str = None,
-        verbose: bool = False,
-    ) -> None:
-        super().__init__(prompt, verbose)
-        self.gpt_client = gpt_client
-        if self.prompt is None:
-            self.prompt = """
-            The first image is the original, and the second image is the
-            result after segmenting the main object. Evaluate the segmentation
-            quality to ensure the main object is clearly segmented without
-            significant truncation. Note that the foreground of the object
-            needs to be extracted instead of the background.
-            Minor imperfections can be ignored. If segmentation is acceptable,
-            return "YES" only; otherwise, return "NO" with
-            very brief explanation.
-            """
-
-    def query(self, image_paths: list[str]) -> str:
-        if len(image_paths) != 2:
-            raise ValueError(
-                "ImageSegChecker requires exactly two images: [raw_image, seg_image]."  # noqa
-            )
-        # Hardcode tmp because of the openrouter can't input multi images.
-        if "openrouter" in self.gpt_client.endpoint:
-            from asset3d_gen.utils.process_media import (
-                combine_images_to_base64,
-            )
-
-            image_paths = combine_images_to_base64(image_paths)
-
-        return self.gpt_client.query(
-            text_prompt=self.prompt,
-            image_base64=image_paths,
-        )
-
-
-class ImageAestheticChecker(BaseChecker):
-    def __init__(
-        self,
-        clip_model_dir: str = None,
-        sac_model_path: str = None,
-        thresh: float = 4.50,
-        verbose: bool = False,
-    ) -> None:
-        super().__init__(verbose=verbose)
-        self.clip_model_dir = clip_model_dir
-        self.sac_model_path = sac_model_path
-        self.thresh = thresh
-        self.predictor = AestheticPredictor(clip_model_dir, sac_model_path)
-
-    def query(self, image_paths: list[str]) -> float:
-        scores = [self.predictor.predict(img_path) for img_path in image_paths]
-        return sum(scores) / len(scores)
-
-    def __call__(self, image_paths: list[str], **kwargs) -> bool:
-        avg_score = self.query(image_paths)
-        if self.verbose:
-            logger.info(f"Average aesthetic score: {avg_score}")
-        return avg_score > self.thresh, avg_score
-
-
-if __name__ == "__main__":
-    geo_checker = MeshGeoChecker(GPT_CLIENT)
-    seg_checker = ImageSegChecker(GPT_CLIENT)
-    aesthetic_checker = ImageAestheticChecker(
-        "/horizon-bucket/robot_lab/users/xinjie.wang/weights/clip",
-        "/horizon-bucket/robot_lab/users/xinjie.wang/weights/sac/sac+logos+ava1-l14-linearMSE.pth",  # noqa
-    )
-
-    checkers = [geo_checker, seg_checker, aesthetic_checker]
-
-    output_root = "outputs/test_gpt"
-
-    fails = []
-    for idx in tqdm(range(150)):
-        mesh_path = f"outputs/imageto3d/demo_objects/cups/sample_{idx}/sample_{idx}.obj"  # noqa
-        if not os.path.exists(mesh_path):
-            continue
-        image_paths = render_asset3d(
-            mesh_path,
-            f"{output_root}/{idx}",
-            num_images=8,
-            elevation=(30, -30),
-            distance=5.5,
-        )
-
-        for cid, checker in enumerate(checkers):
-            if isinstance(checker, ImageSegChecker):
-                images = [
-                    f"outputs/imageto3d/demo_objects/cups/sample_{idx}/sample_{idx}_raw.png",  # noqa
-                    f"outputs/imageto3d/demo_objects/cups/sample_{idx}/sample_{idx}_cond.png",  # noqa
-                ]
-            else:
-                images = image_paths
-            result, info = checker(images)
-            logger.info(
-                f"Checker {checker.__class__.__name__}: {result}, {info}, mesh {mesh_path}"  # noqa
-            )
-
-            if result is False:
-                fails.append((idx, cid, info))
-
-        break

asset3d_gen/validators/urdf_convertor.py

@@ -1,423 +0,0 @@
-import logging
-import os
-import shutil
-import xml.etree.ElementTree as ET
-import zipfile
-from datetime import datetime
-from xml.dom.minidom import parseString
-
-import numpy as np
-import trimesh
-from asset3d_gen.utils.gpt_clients import GPT_CLIENT, GPTclient
-from asset3d_gen.utils.process_media import render_asset3d
-from asset3d_gen.utils.tags import VERSION
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-__all__ = ["URDFGenerator"]
-
-
-URDF_TEMPLATE = """
-<robot name="template_robot">
-    <link name="template_link">
-        <visual>
-            <geometry>
-                <mesh filename="mesh.obj" scale="1.0 1.0 1.0"/>
-            </geometry>
-        </visual>
-        <collision>
-            <geometry>
-                <mesh filename="mesh.obj" scale="1.0 1.0 1.0"/>
-            </geometry>
-            <gazebo>
-                <mu1>0.8</mu1> <!-- 主摩擦系数 -->
-                <mu2>0.6</mu2> <!-- 次摩擦系数 -->
-            </gazebo>
-        </collision>
-        <inertial>
-            <mass value="1.0"/>
-            <origin xyz="0 0 0"/>
-            <inertia ixx="1.0" ixy="0.0" ixz="0.0" iyy="1.0" iyz="0.0" izz="1.0"/>
-        </inertial>
-        <extra_info>
-            <scale>1.0</scale>
-            <version>"0.0.0"</version>
-            <category>"unknown"</category>
-            <description>"unknown"</description>
-            <min_height>0.0</min_height>
-            <max_height>0.0</max_height>
-            <real_height>0.0</real_height>
-            <min_mass>0.0</min_mass>
-            <max_mass>0.0</max_mass>
-            <generate_time>"-1"</generate_time>
-            <gs_model>""</gs_model>
-        </extra_info>
-    </link>
-</robot>
-"""
-
-
-def zip_files(input_paths: list[str], output_zip: str) -> str:
-    with zipfile.ZipFile(output_zip, "w", zipfile.ZIP_DEFLATED) as zipf:
-        for input_path in input_paths:
-            if not os.path.exists(input_path):
-                raise FileNotFoundError(f"File not found: {input_path}")
-
-            if os.path.isdir(input_path):
-                for root, _, files in os.walk(input_path):
-                    for file in files:
-                        file_path = os.path.join(root, file)
-                        arcname = os.path.relpath(
-                            file_path, start=os.path.commonpath(input_paths)
-                        )
-                        zipf.write(file_path, arcname=arcname)
-            else:
-                arcname = os.path.relpath(
-                    input_path, start=os.path.commonpath(input_paths)
-                )
-                zipf.write(input_path, arcname=arcname)
-
-    return output_zip
-
-
-class URDFGenerator(object):
-    def __init__(
-        self,
-        gpt_client: GPTclient,
-        mesh_file_list: list[str] = ["material_0.png", "material.mtl"],
-        prompt_template: str = None,
-        attrs_name: list[str] = None,
-        render_dir: str = "urdf_renders",
-        render_view_num: int = 4,
-    ) -> None:
-        if mesh_file_list is None:
-            mesh_file_list = []
-        self.mesh_file_list = mesh_file_list
-        self.output_mesh_dir = "mesh"
-        self.output_render_dir = render_dir
-        self.gpt_client = gpt_client
-        self.render_view_num = render_view_num
-        if render_view_num == 4:
-            view_desc = "This is orthographic projection showing the front, left, right and back views "  # noqa
-        else:
-            view_desc = "This is the rendered views "
-
-        if prompt_template is None:
-            prompt_template = (
-                view_desc
-                + """of the 3D object asset,
-                category: {category}.
-                Give the category of this object asset (within 3 words),
-                (if category is already provided, use it directly),
-                accurately describe this 3D object asset (within 15 words),
-                and give the recommended geometric height range (unit: meter),
-                weight range (unit: kilogram), the average static friction
-                coefficient of the object relative to rubber and the average
-                dynamic friction coefficient of the object relative to rubber.
-                Return response format as shown in Example.
-
-                Example:
-                Category: cup
-                Description: shiny golden cup with floral design
-                Height: 0.1-0.15 m
-                Weight: 0.3-0.6 kg
-                Static friction coefficient: 1.1
-                Dynamic friction coefficient: 0.9
-            """
-            )
-
-        self.prompt_template = prompt_template
-        if attrs_name is None:
-            attrs_name = [
-                "category",
-                "description",
-                "min_height",
-                "max_height",
-                "real_height",
-                "min_mass",
-                "max_mass",
-                "version",
-                "generate_time",
-                "gs_model",
-            ]
-        self.attrs_name = attrs_name
-
-    def parse_response(self, response: str) -> dict[str, any]:
-        lines = response.split("\n")
-        lines = [line.strip() for line in lines if line]
-        category = lines[0].split(": ")[1]
-        description = lines[1].split(": ")[1]
-        min_height, max_height = map(
-            lambda x: float(x.strip().replace(",", "").split()[0]),
-            lines[2].split(": ")[1].split("-"),
-        )
-        min_mass, max_mass = map(
-            lambda x: float(x.strip().replace(",", "").split()[0]),
-            lines[3].split(": ")[1].split("-"),
-        )
-        mu1 = float(lines[4].split(": ")[1].replace(",", ""))
-        mu2 = float(lines[5].split(": ")[1].replace(",", ""))
-
-        return {
-            "category": category.lower(),
-            "description": description.lower(),
-            "min_height": round(min_height, 4),
-            "max_height": round(max_height, 4),
-            "real_height": round((min_height + max_height) / 2, 4),
-            "min_mass": round(min_mass, 4),
-            "max_mass": round(max_mass, 4),
-            "mu1": round(mu1, 2),
-            "mu2": round(mu2, 2),
-            "version": VERSION,
-            "generate_time": datetime.now().strftime("%Y%m%d%H%M%S"),
-        }
-
-    def generate_urdf(
-        self,
-        input_mesh: str,
-        output_dir: str,
-        attr_dict: dict,
-        output_name: str = None,
-    ) -> str:
-        """Generate a URDF file for a given mesh with specified attributes.
-
-        Args:
-            input_mesh (str): Path to the input mesh file.
-            output_dir (str): Directory to store the generated URDF
-                and processed mesh.
-            attr_dict (dict): Dictionary containing attributes like height,
-                mass, and friction coefficients.
-            output_name (str, optional): Name for the generated URDF and robot.
-
-        Returns:
-            str: Path to the generated URDF file.
-        """
-
-        # 1. Load and normalize the mesh
-        mesh = trimesh.load(input_mesh)
-        mesh_scale = np.ptp(mesh.vertices, axis=0).max()
-        mesh.vertices /= mesh_scale  # Normalize to [-0.5, 0.5]
-        raw_height = np.ptp(mesh.vertices, axis=0)[1]
-
-        # 2. Scale the mesh to real height
-        real_height = attr_dict["real_height"]
-        scale = round(real_height / raw_height, 6)
-        mesh = mesh.apply_scale(scale)
-
-        # 3. Prepare output directories and save scaled mesh
-        mesh_folder = os.path.join(output_dir, self.output_mesh_dir)
-        os.makedirs(mesh_folder, exist_ok=True)
-
-        obj_name = os.path.basename(input_mesh)
-        mesh_output_path = os.path.join(mesh_folder, obj_name)
-        mesh.export(mesh_output_path)
-
-        # 4. Copy additional mesh files, if any
-        input_dir = os.path.dirname(input_mesh)
-        for file in self.mesh_file_list:
-            src_file = os.path.join(input_dir, file)
-            dest_file = os.path.join(mesh_folder, file)
-            if os.path.isfile(src_file):
-                shutil.copy(src_file, dest_file)
-
-        # 5. Determine output name
-        if output_name is None:
-            output_name = os.path.splitext(obj_name)[0]
-
-        # 6. Load URDF template and update attributes
-        robot = ET.fromstring(URDF_TEMPLATE)
-        robot.set("name", output_name)
-
-        link = robot.find("link")
-        if link is None:
-            raise ValueError("URDF template is missing 'link' element.")
-        link.set("name", output_name)
-
-        # Update visual geometry
-        visual = link.find("visual/geometry/mesh")
-        if visual is not None:
-            visual.set(
-                "filename", os.path.join(self.output_mesh_dir, obj_name)
-            )
-            visual.set("scale", "1.0 1.0 1.0")
-
-        # Update collision geometry
-        collision = link.find("collision/geometry/mesh")
-        if collision is not None:
-            collision.set(
-                "filename", os.path.join(self.output_mesh_dir, obj_name)
-            )
-            collision.set("scale", "1.0 1.0 1.0")
-
-        # Update friction coefficients
-        gazebo = link.find("collision/gazebo")
-        if gazebo is not None:
-            for param, key in zip(["mu1", "mu2"], ["mu1", "mu2"]):
-                element = gazebo.find(param)
-                if element is not None:
-                    element.text = f"{attr_dict[key]:.2f}"
-
-        # Update mass
-        inertial = link.find("inertial/mass")
-        if inertial is not None:
-            mass_value = (attr_dict["min_mass"] + attr_dict["max_mass"]) / 2
-            inertial.set("value", f"{mass_value:.4f}")
-
-        # Add extra_info element to the link
-        extra_info = link.find("extra_info/scale")
-        if extra_info is not None:
-            extra_info.text = f"{scale:.6f}"
-
-        for key in self.attrs_name:
-            extra_info = link.find(f"extra_info/{key}")
-            if extra_info is not None and key in attr_dict:
-                extra_info.text = f"{attr_dict[key]}"
-
-        # 7. Write URDF to file
-        os.makedirs(output_dir, exist_ok=True)
-        urdf_path = os.path.join(output_dir, f"{output_name}.urdf")
-        tree = ET.ElementTree(robot)
-        tree.write(urdf_path, encoding="utf-8", xml_declaration=True)
-
-        logger.info(f"URDF file saved to {urdf_path}")
-
-        return urdf_path
-
-    @staticmethod
-    def get_attr_from_urdf(
-        urdf_path: str,
-        attr_root: str = ".//link/extra_info",
-        attr_name: str = "scale",
-    ) -> float:
-        if not os.path.exists(urdf_path):
-            raise FileNotFoundError(f"URDF file not found: {urdf_path}")
-
-        mesh_scale = 1.0
-        tree = ET.parse(urdf_path)
-        root = tree.getroot()
-        extra_info = root.find(attr_root)
-        if extra_info is not None:
-            scale_element = extra_info.find(attr_name)
-            if scale_element is not None:
-                mesh_scale = float(scale_element.text)
-
-        return mesh_scale
-
-    @staticmethod
-    def add_quality_tag(
-        urdf_path: str, results, output_path: str = None
-    ) -> None:
-        if output_path is None:
-            output_path = urdf_path
-
-        tree = ET.parse(urdf_path)
-        root = tree.getroot()
-        custom_data = ET.SubElement(root, "custom_data")
-        quality = ET.SubElement(custom_data, "quality")
-        for key, value in results:
-            checker_tag = ET.SubElement(quality, key)
-            checker_tag.text = str(value)
-
-        rough_string = ET.tostring(root, encoding="utf-8")
-        formatted_string = parseString(rough_string).toprettyxml(indent="   ")
-        cleaned_string = "\n".join(
-            [line for line in formatted_string.splitlines() if line.strip()]
-        )
-
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
-        with open(output_path, "w", encoding="utf-8") as f:
-            f.write(cleaned_string)
-
-        logger.info(f"URDF files saved to {output_path}")
-
-    def get_estimated_attributes(self, asset_attrs: dict):
-        estimated_attrs = {
-            "height": round(
-                (asset_attrs["min_height"] + asset_attrs["max_height"]) / 2, 4
-            ),
-            "mass": round(
-                (asset_attrs["min_mass"] + asset_attrs["max_mass"]) / 2, 4
-            ),
-            "mu": round((asset_attrs["mu1"] + asset_attrs["mu2"]) / 2, 4),
-            "category": asset_attrs["category"],
-        }
-
-        return estimated_attrs
-
-    def __call__(
-        self,
-        mesh_path: str,
-        output_root: str,
-        text_prompt: str = None,
-        category: str = "unknown",
-        **kwargs,
-    ):
-        if text_prompt is None or len(text_prompt) == 0:
-            text_prompt = self.prompt_template
-            text_prompt = text_prompt.format(category=category.lower())
-
-        image_path = render_asset3d(
-            mesh_path,
-            output_root,
-            num_images=self.render_view_num,
-            output_subdir=self.output_render_dir,
-        )
-
-        # Hardcode tmp because of the openrouter can't input multi images.
-        if "openrouter" in self.gpt_client.endpoint:
-            from asset3d_gen.utils.process_media import (
-                combine_images_to_base64,
-            )
-
-            image_path = combine_images_to_base64(image_path)
-
-        response = self.gpt_client.query(text_prompt, image_path)
-        if response is None:
-            asset_attrs = {
-                "category": "unknown",
-                "description": "unknown",
-                "min_height": 1,
-                "max_height": 1,
-                "real_height": 1,
-                "min_mass": 1,
-                "max_mass": 1,
-                "mu1": 0.8,
-                "mu2": 0.6,
-                "version": VERSION,
-                "generate_time": datetime.now().strftime("%Y%m%d%H%M%S"),
-            }
-        else:
-            asset_attrs = self.parse_response(response)
-        for key in self.attrs_name:
-            if key in kwargs:
-                asset_attrs[key] = kwargs[key]
-
-        self.estimated_attrs = self.get_estimated_attributes(asset_attrs)
-
-        urdf_path = self.generate_urdf(mesh_path, output_root, asset_attrs)
-
-        logger.info(f"response: {response}")
-
-        return urdf_path
-
-
-if __name__ == "__main__":
-    urdf_gen = URDFGenerator(GPT_CLIENT, render_view_num=4)
-    urdf_path = urdf_gen(
-        mesh_path="scripts/apps/assets/example_texture/meshes/robot.obj",
-        output_root="outputs/test_urdf",
-        # category="coffee machine",
-        # min_height=1.0,
-        # max_height=1.2,
-        version=VERSION,
-    )
-
-    # zip_files(
-    #     input_paths=[
-    #         "scripts/apps/tmp/2umpdum3e5n/URDF_sample/mesh",
-    #         "scripts/apps/tmp/2umpdum3e5n/URDF_sample/sample.urdf"
-    #     ],
-    #     output_zip="zip.zip"
-    # )