HorizonRobotics/EmbodiedGen-Image-to-3D
47
1# Project EmbodiedGen2#3# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or14# implied. See the License for the specific language governing15# permissions and limitations under the License.16 17import argparse18import logging19import math20import os21from typing import Literal, Union22 23import cv224import numpy as np25import nvdiffrast.torch as dr26import spaces27import torch28import trimesh29import utils3d30import xatlas31from PIL import Image32from tqdm import tqdm33from embodied_gen.data.mesh_operator import MeshFixer34from embodied_gen.data.utils import (35 CameraSetting,36 init_kal_camera,37 kaolin_to_opencv_view,38 normalize_vertices_array,39 post_process_texture,40 save_mesh_with_mtl,41)42from embodied_gen.models.delight_model import DelightingModel43from embodied_gen.models.gs_model import load_gs_model44from embodied_gen.models.sr_model import ImageRealESRGAN45 46logging.basicConfig(47 format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO48)49logger = logging.getLogger(__name__)50 51 52__all__ = [53 "TextureBaker",54]55 56 57class TextureBaker(object):58 """Baking textures onto a mesh from multiple observations.59 60 This class take 3D mesh data, camera settings and texture baking parameters61 to generate texture map by projecting images to the mesh from diff views.62 It supports both a fast texture baking approach and a more optimized method63 with total variation regularization.64 65 Attributes:66 vertices (torch.Tensor): The vertices of the mesh.67 faces (torch.Tensor): The faces of the mesh, defined by vertex indices.68 uvs (torch.Tensor): The UV coordinates of the mesh.69 camera_params (CameraSetting): Camera setting (intrinsics, extrinsics).70 device (str): The device to run computations on ("cpu" or "cuda").71 w2cs (torch.Tensor): World-to-camera transformation matrices.72 projections (torch.Tensor): Camera projection matrices.73 74 Example:75 >>> vertices, faces, uvs = TextureBaker.parametrize_mesh(vertices, faces) # noqa76 >>> texture_backer = TextureBaker(vertices, faces, uvs, camera_params)77 >>> images = get_images_from_grid(args.color_path, image_size)78 >>> texture = texture_backer.bake_texture(79 ... images, texture_size=args.texture_size, mode=args.baker_mode80 ... )81 >>> texture = post_process_texture(texture)82 """83 84 def __init__(85 self,86 vertices: np.ndarray,87 faces: np.ndarray,88 uvs: np.ndarray,89 camera_params: CameraSetting,90 device: str = "cuda",91 ) -> None:92 self.vertices = (93 torch.tensor(vertices, device=device)94 if isinstance(vertices, np.ndarray)95 else vertices.to(device)96 )97 self.faces = (98 torch.tensor(faces.astype(np.int32), device=device)99 if isinstance(faces, np.ndarray)100 else faces.to(device)101 )102 self.uvs = (103 torch.tensor(uvs, device=device)104 if isinstance(uvs, np.ndarray)105 else uvs.to(device)106 )107 self.camera_params = camera_params108 self.device = device109 110 camera = init_kal_camera(camera_params)111 matrix_mv = camera.view_matrix() # (n_cam 4 4) world2cam112 matrix_mv = kaolin_to_opencv_view(matrix_mv)113 matrix_p = (114 camera.intrinsics.projection_matrix()115 ) # (n_cam 4 4) cam2pixel116 self.w2cs = matrix_mv.to(self.device)117 self.projections = matrix_p.to(self.device)118 119 @staticmethod120 def parametrize_mesh(121 vertices: np.array, faces: np.array122 ) -> Union[np.array, np.array, np.array]:123 vmapping, indices, uvs = xatlas.parametrize(vertices, faces)124 125 vertices = vertices[vmapping]126 faces = indices127 128 return vertices, faces, uvs129 130 def _bake_fast(self, observations, w2cs, projections, texture_size, masks):131 texture = torch.zeros(132 (texture_size * texture_size, 3), dtype=torch.float32133 ).cuda()134 texture_weights = torch.zeros(135 (texture_size * texture_size), dtype=torch.float32136 ).cuda()137 rastctx = utils3d.torch.RastContext(backend="cuda")138 for observation, w2c, projection in tqdm(139 zip(observations, w2cs, projections),140 total=len(observations),141 desc="Texture baking (fast)",142 ):143 with torch.no_grad():144 rast = utils3d.torch.rasterize_triangle_faces(145 rastctx,146 self.vertices[None],147 self.faces,148 observation.shape[1],149 observation.shape[0],150 uv=self.uvs[None],151 view=w2c,152 projection=projection,153 )154 uv_map = rast["uv"][0].detach().flip(0)155 mask = rast["mask"][0].detach().bool() & masks[0]156 157 # nearest neighbor interpolation158 uv_map = (uv_map * texture_size).floor().long()159 obs = observation[mask]160 uv_map = uv_map[mask]161 idx = (162 uv_map[:, 0] + (texture_size - uv_map[:, 1] - 1) * texture_size163 )164 texture = texture.scatter_add(165 0, idx.view(-1, 1).expand(-1, 3), obs166 )167 texture_weights = texture_weights.scatter_add(168 0,169 idx,170 torch.ones(171 (obs.shape[0]), dtype=torch.float32, device=texture.device172 ),173 )174 175 mask = texture_weights > 0176 texture[mask] /= texture_weights[mask][:, None]177 texture = np.clip(178 texture.reshape(texture_size, texture_size, 3).cpu().numpy() * 255,179 0,180 255,181 ).astype(np.uint8)182 183 # inpaint184 mask = (185 (texture_weights == 0)186 .cpu()187 .numpy()188 .astype(np.uint8)189 .reshape(texture_size, texture_size)190 )191 texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA)192 193 return texture194 195 def _bake_opt(196 self,197 observations,198 w2cs,199 projections,200 texture_size,201 lambda_tv,202 masks,203 total_steps,204 ):205 rastctx = utils3d.torch.RastContext(backend="cuda")206 observations = [observations.flip(0) for observations in observations]207 masks = [m.flip(0) for m in masks]208 _uv = []209 _uv_dr = []210 for observation, w2c, projection in tqdm(211 zip(observations, w2cs, projections),212 total=len(w2cs),213 ):214 with torch.no_grad():215 rast = utils3d.torch.rasterize_triangle_faces(216 rastctx,217 self.vertices[None],218 self.faces,219 observation.shape[1],220 observation.shape[0],221 uv=self.uvs[None],222 view=w2c,223 projection=projection,224 )225 _uv.append(rast["uv"].detach())226 _uv_dr.append(rast["uv_dr"].detach())227 228 texture = torch.nn.Parameter(229 torch.zeros(230 (1, texture_size, texture_size, 3), dtype=torch.float32231 ).cuda()232 )233 optimizer = torch.optim.Adam([texture], betas=(0.5, 0.9), lr=1e-2)234 235 def cosine_anealing(step, total_steps, start_lr, end_lr):236 return end_lr + 0.5 * (start_lr - end_lr) * (237 1 + np.cos(np.pi * step / total_steps)238 )239 240 def tv_loss(texture):241 return torch.nn.functional.l1_loss(242 texture[:, :-1, :, :], texture[:, 1:, :, :]243 ) + torch.nn.functional.l1_loss(244 texture[:, :, :-1, :], texture[:, :, 1:, :]245 )246 247 with tqdm(total=total_steps, desc="Texture baking") as pbar:248 for step in range(total_steps):249 optimizer.zero_grad()250 selected = np.random.randint(0, len(w2cs))251 uv, uv_dr, observation, mask = (252 _uv[selected],253 _uv_dr[selected],254 observations[selected],255 masks[selected],256 )257 render = dr.texture(texture, uv, uv_dr)[0]258 loss = torch.nn.functional.l1_loss(259 render[mask], observation[mask]260 )261 if lambda_tv > 0:262 loss += lambda_tv * tv_loss(texture)263 loss.backward()264 optimizer.step()265 266 optimizer.param_groups[0]["lr"] = cosine_anealing(267 step, total_steps, 1e-2, 1e-5268 )269 pbar.set_postfix({"loss": loss.item()})270 pbar.update()271 272 texture = np.clip(273 texture[0].flip(0).detach().cpu().numpy() * 255, 0, 255274 ).astype(np.uint8)275 mask = 1 - utils3d.torch.rasterize_triangle_faces(276 rastctx,277 (self.uvs * 2 - 1)[None],278 self.faces,279 texture_size,280 texture_size,281 )["mask"][0].detach().cpu().numpy().astype(np.uint8)282 texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA)283 284 return texture285 286 def bake_texture(287 self,288 images: list[np.array],289 texture_size: int = 1024,290 mode: Literal["fast", "opt"] = "opt",291 lambda_tv: float = 1e-2,292 opt_step: int = 2000,293 ):294 masks = [np.any(img > 0, axis=-1) for img in images]295 masks = [torch.tensor(m > 0).bool().to(self.device) for m in masks]296 images = [297 torch.tensor(obs / 255.0).float().to(self.device) for obs in images298 ]299 300 if mode == "fast":301 return self._bake_fast(302 images, self.w2cs, self.projections, texture_size, masks303 )304 elif mode == "opt":305 return self._bake_opt(306 images,307 self.w2cs,308 self.projections,309 texture_size,310 lambda_tv,311 masks,312 opt_step,313 )314 else:315 raise ValueError(f"Unknown mode: {mode}")316 317 318def parse_args():319 """Parses command-line arguments for texture backprojection.320 321 Returns:322 argparse.Namespace: Parsed arguments.323 """324 parser = argparse.ArgumentParser(description="Backproject texture")325 parser.add_argument(326 "--gs_path",327 type=str,328 help="Path to the GS.ply gaussian splatting model",329 )330 parser.add_argument(331 "--mesh_path",332 type=str,333 help="Mesh path, .obj, .glb or .ply",334 )335 parser.add_argument(336 "--output_path",337 type=str,338 help="Output mesh path with suffix",339 )340 parser.add_argument(341 "--num_images",342 type=int,343 default=180,344 help="Number of images to render.",345 )346 parser.add_argument(347 "--elevation",348 nargs="+",349 type=float,350 default=list(range(85, -90, -10)),351 help="Elevation angles for the camera",352 )353 parser.add_argument(354 "--distance",355 type=float,356 default=4.5,357 help="Camera distance (default: 4.5)",358 )359 parser.add_argument(360 "--resolution_hw",361 type=int,362 nargs=2,363 default=(512, 512),364 help="Resolution of the render images (default: (512, 512))",365 )366 parser.add_argument(367 "--fov",368 type=float,369 default=30,370 help="Field of view in degrees (default: 30)",371 )372 parser.add_argument(373 "--device",374 type=str,375 choices=["cpu", "cuda"],376 default="cuda",377 help="Device to run on (default: `cuda`)",378 )379 parser.add_argument(380 "--skip_fix_mesh", action="store_true", help="Fix mesh geometry."381 )382 parser.add_argument(383 "--texture_size",384 type=int,385 default=2048,386 help="Texture size for texture baking (default: 1024)",387 )388 parser.add_argument(389 "--baker_mode",390 type=str,391 default="opt",392 help="Texture baking mode, `fast` or `opt` (default: opt)",393 )394 parser.add_argument(395 "--opt_step",396 type=int,397 default=3000,398 help="Optimization steps for texture baking (default: 3000)",399 )400 parser.add_argument(401 "--mesh_sipmlify_ratio",402 type=float,403 default=0.85,404 help="Mesh simplification ratio (default: 0.85)",405 )406 parser.add_argument(407 "--delight", action="store_true", help="Use delighting model."408 )409 parser.add_argument(410 "--no_smooth_texture",411 action="store_true",412 help="Do not smooth the texture.",413 )414 parser.add_argument(415 "--no_coor_trans",416 action="store_true",417 help="Do not transform the asset coordinate system.",418 )419 parser.add_argument(420 "--save_glb_path", type=str, default=None, help="Save glb path."421 )422 parser.add_argument("--n_max_faces", type=int, default=50000)423 args, unknown = parser.parse_known_args()424 425 return args426 427 428@spaces.GPU429def entrypoint(430 delight_model: DelightingModel = None,431 imagesr_model: ImageRealESRGAN = None,432 **kwargs,433) -> trimesh.Trimesh:434 """Entrypoint for texture backprojection from multi-view images.435 436 Args:437 delight_model (DelightingModel, optional): Delighting model.438 imagesr_model (ImageRealESRGAN, optional): Super-resolution model.439 **kwargs: Additional arguments to override CLI.440 441 Returns:442 trimesh.Trimesh: Textured mesh.443 """444 args = parse_args()445 for k, v in kwargs.items():446 if hasattr(args, k) and v is not None:447 setattr(args, k, v)448 449 # Setup camera parameters.450 camera_params = CameraSetting(451 num_images=args.num_images,452 elevation=args.elevation,453 distance=args.distance,454 resolution_hw=args.resolution_hw,455 fov=math.radians(args.fov),456 device=args.device,457 )458 459 # GS render.460 camera = init_kal_camera(camera_params, flip_az=True)461 matrix_mv = camera.view_matrix() # (n_cam 4 4) world2cam462 matrix_mv[:, :3, 3] = -matrix_mv[:, :3, 3]463 w2cs = matrix_mv.to(camera_params.device)464 c2ws = [torch.linalg.inv(matrix) for matrix in w2cs]465 Ks = torch.tensor(camera_params.Ks).to(camera_params.device)466 gs_model = load_gs_model(args.gs_path, pre_quat=[0.0, 0.0, 1.0, 0.0])467 multiviews = []468 for idx in tqdm(range(len(c2ws)), desc="Rendering GS"):469 result = gs_model.render(470 c2ws[idx],471 Ks=Ks,472 image_width=camera_params.resolution_hw[1],473 image_height=camera_params.resolution_hw[0],474 )475 color = cv2.cvtColor(result.rgba, cv2.COLOR_BGRA2RGBA)476 multiviews.append(Image.fromarray(color))477 478 if args.delight and delight_model is None:479 delight_model = DelightingModel()480 481 if args.delight:482 for idx in range(len(multiviews)):483 multiviews[idx] = delight_model(multiviews[idx])484 485 multiviews = [img.convert("RGB") for img in multiviews]486 487 mesh = trimesh.load(args.mesh_path)488 if isinstance(mesh, trimesh.Scene):489 mesh = mesh.dump(concatenate=True)490 491 vertices, scale, center = normalize_vertices_array(mesh.vertices)492 493 # Transform mesh coordinate system by default.494 if not args.no_coor_trans:495 x_rot = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]])496 z_rot = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]])497 vertices = vertices @ x_rot498 vertices = vertices @ z_rot499 500 faces = mesh.faces.astype(np.int32)501 vertices = vertices.astype(np.float32)502 503 if not args.skip_fix_mesh:504 mesh_fixer = MeshFixer(vertices, faces, args.device)505 vertices, faces = mesh_fixer(506 filter_ratio=args.mesh_sipmlify_ratio,507 max_hole_size=0.04,508 resolution=1024,509 num_views=1000,510 norm_mesh_ratio=0.5,511 )512 if len(faces) > args.n_max_faces:513 mesh_fixer = MeshFixer(vertices, faces, args.device)514 vertices, faces = mesh_fixer(515 filter_ratio=max(0.1, args.mesh_sipmlify_ratio - 0.1),516 max_hole_size=0.04,517 resolution=1024,518 num_views=1000,519 norm_mesh_ratio=0.5,520 )521 522 vertices, faces, uvs = TextureBaker.parametrize_mesh(vertices, faces)523 texture_backer = TextureBaker(524 vertices,525 faces,526 uvs,527 camera_params,528 )529 530 multiviews = [np.array(img) for img in multiviews]531 texture = texture_backer.bake_texture(532 images=[img[..., :3] for img in multiviews],533 texture_size=args.texture_size,534 mode=args.baker_mode,535 opt_step=args.opt_step,536 )537 if not args.no_smooth_texture:538 texture = post_process_texture(texture)539 540 # Recover mesh original orientation, scale and center.541 if not args.no_coor_trans:542 vertices = vertices @ np.linalg.inv(z_rot)543 vertices = vertices @ np.linalg.inv(x_rot)544 vertices = vertices / scale545 vertices = vertices + center546 547 textured_mesh = save_mesh_with_mtl(548 vertices, faces, uvs, texture, args.output_path549 )550 if args.save_glb_path is not None:551 os.makedirs(os.path.dirname(args.save_glb_path), exist_ok=True)552 textured_mesh.export(args.save_glb_path)553 554 return textured_mesh555 556 557if __name__ == "__main__":558 entrypoint()559