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 17 18import argparse19import logging20import math21import os22 23import cv224import numpy as np25import nvdiffrast.torch as dr26import spaces27import torch28import torch.nn.functional as F29import trimesh30import xatlas31from PIL import Image32from embodied_gen.data.mesh_operator import MeshFixer33from embodied_gen.data.utils import (34 CameraSetting,35 DiffrastRender,36 as_list,37 get_images_from_grid,38 init_kal_camera,39 normalize_vertices_array,40 post_process_texture,41 save_mesh_with_mtl,42)43from embodied_gen.models.delight_model import DelightingModel44from embodied_gen.models.sr_model import ImageRealESRGAN45from embodied_gen.utils.process_media import vcat_pil_images46 47logging.basicConfig(48 format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO49)50logger = logging.getLogger(__name__)51 52 53__all__ = [54 "TextureBacker",55]56 57 58def _transform_vertices(59 mtx: torch.Tensor, pos: torch.Tensor, keepdim: bool = False60) -> torch.Tensor:61 """Transforms 3D vertices using a projection matrix.62 63 Args:64 mtx (torch.Tensor): Projection matrix.65 pos (torch.Tensor): Vertex positions.66 keepdim (bool, optional): If True, keeps the batch dimension.67 68 Returns:69 torch.Tensor: Transformed vertices.70 """71 t_mtx = torch.as_tensor(mtx, device=pos.device, dtype=pos.dtype)72 if pos.size(-1) == 3:73 pos = torch.cat([pos, torch.ones_like(pos[..., :1])], dim=-1)74 75 result = pos @ t_mtx.T76 77 return result if keepdim else result.unsqueeze(0)78 79 80def _bilinear_interpolation_scattering(81 image_h: int, image_w: int, coords: torch.Tensor, values: torch.Tensor82) -> torch.Tensor:83 """Performs bilinear interpolation scattering for grid-based value accumulation.84 85 Args:86 image_h (int): Image height.87 image_w (int): Image width.88 coords (torch.Tensor): Normalized coordinates.89 values (torch.Tensor): Values to scatter.90 91 Returns:92 torch.Tensor: Interpolated grid.93 """94 device = values.device95 dtype = values.dtype96 C = values.shape[-1]97 98 indices = coords * torch.tensor(99 [image_h - 1, image_w - 1], dtype=dtype, device=device100 )101 i, j = indices.unbind(-1)102 103 i0, j0 = (104 indices.floor()105 .long()106 .clamp(0, image_h - 2)107 .clamp(0, image_w - 2)108 .unbind(-1)109 )110 i1, j1 = i0 + 1, j0 + 1111 112 w_i = i - i0.float()113 w_j = j - j0.float()114 weights = torch.stack(115 [(1 - w_i) * (1 - w_j), (1 - w_i) * w_j, w_i * (1 - w_j), w_i * w_j],116 dim=1,117 )118 119 indices_comb = torch.stack(120 [121 torch.stack([i0, j0], dim=1),122 torch.stack([i0, j1], dim=1),123 torch.stack([i1, j0], dim=1),124 torch.stack([i1, j1], dim=1),125 ],126 dim=1,127 )128 129 grid = torch.zeros(image_h, image_w, C, device=device, dtype=dtype)130 cnt = torch.zeros(image_h, image_w, 1, device=device, dtype=dtype)131 132 for k in range(4):133 idx = indices_comb[:, k]134 w = weights[:, k].unsqueeze(-1)135 136 stride = torch.tensor([image_w, 1], device=device, dtype=torch.long)137 flat_idx = (idx * stride).sum(-1)138 139 grid.view(-1, C).scatter_add_(140 0, flat_idx.unsqueeze(-1).expand(-1, C), values * w141 )142 cnt.view(-1, 1).scatter_add_(0, flat_idx.unsqueeze(-1), w)143 144 mask = cnt.squeeze(-1) > 0145 grid[mask] = grid[mask] / cnt[mask].repeat(1, C)146 147 return grid148 149 150def _texture_inpaint_smooth(151 texture: np.ndarray,152 mask: np.ndarray,153 vertices: np.ndarray,154 faces: np.ndarray,155 uv_map: np.ndarray,156) -> tuple[np.ndarray, np.ndarray]:157 """Performs texture inpainting using vertex-based color propagation.158 159 Args:160 texture (np.ndarray): Texture image.161 mask (np.ndarray): Mask image.162 vertices (np.ndarray): Mesh vertices.163 faces (np.ndarray): Mesh faces.164 uv_map (np.ndarray): UV coordinates.165 166 Returns:167 tuple[np.ndarray, np.ndarray]: Inpainted texture and updated mask.168 """169 image_h, image_w, C = texture.shape170 N = vertices.shape[0]171 172 # Initialize vertex data structures173 vtx_mask = np.zeros(N, dtype=np.float32)174 vtx_colors = np.zeros((N, C), dtype=np.float32)175 unprocessed = []176 adjacency = [[] for _ in range(N)]177 178 # Build adjacency graph and initial color assignment179 for face_idx in range(faces.shape[0]):180 for k in range(3):181 uv_idx_k = faces[face_idx, k]182 v_idx = faces[face_idx, k]183 184 # Convert UV to pixel coordinates with boundary clamping185 u = np.clip(186 int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1187 )188 v = np.clip(189 int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),190 0,191 image_h - 1,192 )193 194 if mask[v, u]:195 vtx_mask[v_idx] = 1.0196 vtx_colors[v_idx] = texture[v, u]197 elif v_idx not in unprocessed:198 unprocessed.append(v_idx)199 200 # Build undirected adjacency graph201 neighbor = faces[face_idx, (k + 1) % 3]202 if neighbor not in adjacency[v_idx]:203 adjacency[v_idx].append(neighbor)204 if v_idx not in adjacency[neighbor]:205 adjacency[neighbor].append(v_idx)206 207 # Color propagation with dynamic stopping208 remaining_iters, prev_count = 2, 0209 while remaining_iters > 0:210 current_unprocessed = []211 212 for v_idx in unprocessed:213 valid_neighbors = [n for n in adjacency[v_idx] if vtx_mask[n] > 0]214 if not valid_neighbors:215 current_unprocessed.append(v_idx)216 continue217 218 # Calculate inverse square distance weights219 neighbors_pos = vertices[valid_neighbors]220 dist_sq = np.sum((vertices[v_idx] - neighbors_pos) ** 2, axis=1)221 weights = 1 / np.maximum(dist_sq, 1e-8)222 223 vtx_colors[v_idx] = np.average(224 vtx_colors[valid_neighbors], weights=weights, axis=0225 )226 vtx_mask[v_idx] = 1.0227 228 # Update iteration control229 if len(current_unprocessed) == prev_count:230 remaining_iters -= 1231 else:232 remaining_iters = min(remaining_iters + 1, 2)233 prev_count = len(current_unprocessed)234 unprocessed = current_unprocessed235 236 # Generate output texture237 inpainted_texture, updated_mask = texture.copy(), mask.copy()238 for face_idx in range(faces.shape[0]):239 for k in range(3):240 v_idx = faces[face_idx, k]241 if not vtx_mask[v_idx]:242 continue243 244 # UV coordinate conversion245 uv_idx_k = faces[face_idx, k]246 u = np.clip(247 int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1248 )249 v = np.clip(250 int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),251 0,252 image_h - 1,253 )254 255 inpainted_texture[v, u] = vtx_colors[v_idx]256 updated_mask[v, u] = 255257 258 return inpainted_texture, updated_mask259 260 261class TextureBacker:262 """Texture baking pipeline for multi-view projection and fusion.263 264 This class generates UV-based textures for a 3D mesh using multi-view images,265 depth, and normal information. It includes mesh normalization, UV unwrapping,266 visibility-aware back-projection, confidence-weighted fusion, and inpainting.267 268 Args:269 camera_params (CameraSetting): Camera intrinsics and extrinsics.270 view_weights (list[float]): Weights for each view in texture fusion.271 render_wh (tuple[int, int], optional): Intermediate rendering resolution.272 texture_wh (tuple[int, int], optional): Output texture resolution.273 bake_angle_thresh (int, optional): Max angle for valid projection.274 mask_thresh (float, optional): Threshold for visibility masks.275 smooth_texture (bool, optional): Apply post-processing to texture.276 inpaint_smooth (bool, optional): Apply inpainting smoothing.277 mesh_post_process (bool, optional): False for preventing modification of vertices.278 279 Example:280 ```py281 from embodied_gen.data.backproject_v2 import TextureBacker282 from embodied_gen.data.utils import CameraSetting283 import trimesh284 from PIL import Image285 286 camera_params = CameraSetting(287 num_images=6,288 elevation=[20, -10],289 distance=5,290 resolution_hw=(2048,2048),291 fov=math.radians(30),292 device='cuda',293 )294 view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]295 mesh = trimesh.load('mesh.obj')296 images = [Image.open(f'view_{i}.png') for i in range(6)]297 texture_backer = TextureBacker(camera_params, view_weights)298 textured_mesh = texture_backer(images, mesh, 'output.obj')299 ```300 """301 302 def __init__(303 self,304 camera_params: CameraSetting,305 view_weights: list[float],306 render_wh: tuple[int, int] = (2048, 2048),307 texture_wh: tuple[int, int] = (2048, 2048),308 bake_angle_thresh: int = 75,309 mask_thresh: float = 0.5,310 smooth_texture: bool = True,311 inpaint_smooth: bool = False,312 mesh_post_process: bool = True,313 ) -> None:314 self.camera_params = camera_params315 self.renderer = None316 self.view_weights = view_weights317 self.device = camera_params.device318 self.render_wh = render_wh319 self.texture_wh = texture_wh320 self.mask_thresh = mask_thresh321 self.smooth_texture = smooth_texture322 self.inpaint_smooth = inpaint_smooth323 self.mesh_post_process = mesh_post_process324 325 self.bake_angle_thresh = bake_angle_thresh326 self.bake_unreliable_kernel_size = int(327 (2 / 512) * max(self.render_wh[0], self.render_wh[1])328 )329 330 def _lazy_init_render(self, camera_params, mask_thresh):331 """Lazily initializes the renderer.332 333 Args:334 camera_params (CameraSetting): Camera settings.335 mask_thresh (float): Mask threshold.336 """337 if self.renderer is None:338 camera = init_kal_camera(camera_params)339 mv = camera.view_matrix() # (n 4 4) world2cam340 p = camera.intrinsics.projection_matrix()341 # NOTE: add a negative sign at P[0, 2] as the y axis is flipped in `nvdiffrast` output. # noqa342 p[:, 1, 1] = -p[:, 1, 1]343 self.renderer = DiffrastRender(344 p_matrix=p,345 mv_matrix=mv,346 resolution_hw=camera_params.resolution_hw,347 context=dr.RasterizeCudaContext(),348 mask_thresh=mask_thresh,349 grad_db=False,350 device=self.device,351 antialias_mask=True,352 )353 354 def load_mesh(self, mesh: trimesh.Trimesh) -> trimesh.Trimesh:355 """Normalizes mesh and unwraps UVs.356 357 Args:358 mesh (trimesh.Trimesh): Input mesh.359 360 Returns:361 trimesh.Trimesh: Mesh with normalized vertices and UVs.362 """363 mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)364 self.scale, self.center = scale, center365 366 vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)367 uvs[:, 1] = 1 - uvs[:, 1]368 mesh.vertices = mesh.vertices[vmapping]369 mesh.faces = indices370 mesh.visual.uv = uvs371 372 return mesh373 374 def get_mesh_np_attrs(375 self,376 mesh: trimesh.Trimesh,377 scale: float = None,378 center: np.ndarray = None,379 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:380 """Gets mesh attributes as numpy arrays.381 382 Args:383 mesh (trimesh.Trimesh): Input mesh.384 scale (float, optional): Scale factor.385 center (np.ndarray, optional): Center offset.386 387 Returns:388 tuple: (vertices, faces, uv_map)389 """390 vertices = mesh.vertices.copy()391 faces = mesh.faces.copy()392 uv_map = mesh.visual.uv.copy()393 uv_map[:, 1] = 1.0 - uv_map[:, 1]394 395 if scale is not None:396 vertices = vertices / scale397 if center is not None:398 vertices = vertices + center399 400 return vertices, faces, uv_map401 402 def _render_depth_edges(self, depth_image: torch.Tensor) -> torch.Tensor:403 """Computes edge image from depth map.404 405 Args:406 depth_image (torch.Tensor): Depth map.407 408 Returns:409 torch.Tensor: Edge image.410 """411 depth_image_np = depth_image.cpu().numpy()412 depth_image_np = (depth_image_np * 255).astype(np.uint8)413 depth_edges = cv2.Canny(depth_image_np, 30, 80)414 sketch_image = (415 torch.from_numpy(depth_edges).to(depth_image.device).float() / 255416 )417 sketch_image = sketch_image.unsqueeze(-1)418 419 return sketch_image420 421 def compute_enhanced_viewnormal(422 self, mv_mtx: torch.Tensor, vertices: torch.Tensor, faces: torch.Tensor423 ) -> torch.Tensor:424 """Computes enhanced view normals for mesh faces.425 426 Args:427 mv_mtx (torch.Tensor): View matrices.428 vertices (torch.Tensor): Mesh vertices.429 faces (torch.Tensor): Mesh faces.430 431 Returns:432 torch.Tensor: View normals.433 """434 rast, _ = self.renderer.compute_dr_raster(vertices, faces)435 rendered_view_normals = []436 for idx in range(len(mv_mtx)):437 pos_cam = _transform_vertices(mv_mtx[idx], vertices, keepdim=True)438 pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]439 v0, v1, v2 = (pos_cam[faces[:, i]] for i in range(3))440 face_norm = F.normalize(441 torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1442 )443 vertex_norm = (444 torch.from_numpy(445 trimesh.geometry.mean_vertex_normals(446 len(pos_cam), faces.cpu(), face_norm.cpu()447 )448 )449 .to(vertices.device)450 .contiguous()451 )452 im_base_normals, _ = dr.interpolate(453 vertex_norm[None, ...].float(),454 rast[idx : idx + 1],455 faces.to(torch.int32),456 )457 rendered_view_normals.append(im_base_normals)458 459 rendered_view_normals = torch.cat(rendered_view_normals, dim=0)460 461 return rendered_view_normals462 463 def back_project(464 self, image, vis_mask, depth, normal, uv465 ) -> tuple[torch.Tensor, torch.Tensor]:466 """Back-projects image and confidence to UV texture space.467 468 Args:469 image (PIL.Image or np.ndarray): Input image.470 vis_mask (torch.Tensor): Visibility mask.471 depth (torch.Tensor): Depth map.472 normal (torch.Tensor): Normal map.473 uv (torch.Tensor): UV coordinates.474 475 Returns:476 tuple[torch.Tensor, torch.Tensor]: Texture and confidence map.477 """478 image = np.array(image)479 image = torch.as_tensor(image, device=self.device, dtype=torch.float32)480 if image.ndim == 2:481 image = image.unsqueeze(-1)482 image = image / 255483 484 depth_inv = (1.0 - depth) * vis_mask485 sketch_image = self._render_depth_edges(depth_inv)486 487 cos = F.cosine_similarity(488 torch.tensor([[0, 0, 1]], device=self.device),489 normal.view(-1, 3),490 ).view_as(normal[..., :1])491 cos[cos < np.cos(np.radians(self.bake_angle_thresh))] = 0492 493 k = self.bake_unreliable_kernel_size * 2 + 1494 kernel = torch.ones((1, 1, k, k), device=self.device)495 496 vis_mask = vis_mask.permute(2, 0, 1).unsqueeze(0).float()497 vis_mask = F.conv2d(498 1.0 - vis_mask,499 kernel,500 padding=k // 2,501 )502 vis_mask = 1.0 - (vis_mask > 0).float()503 vis_mask = vis_mask.squeeze(0).permute(1, 2, 0)504 505 sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)506 sketch_image = F.conv2d(sketch_image, kernel, padding=k // 2)507 sketch_image = (sketch_image > 0).float()508 sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)509 vis_mask = vis_mask * (sketch_image < 0.5)510 511 cos[vis_mask == 0] = 0512 valid_pixels = (vis_mask != 0).view(-1)513 514 return (515 self._scatter_texture(uv, image, valid_pixels),516 self._scatter_texture(uv, cos, valid_pixels),517 )518 519 def _scatter_texture(self, uv, data, mask):520 """Scatters data to texture using UV coordinates and mask.521 522 Args:523 uv (torch.Tensor): UV coordinates.524 data (torch.Tensor): Data to scatter.525 mask (torch.Tensor): Mask for valid pixels.526 527 Returns:528 torch.Tensor: Scattered texture.529 """530 531 def __filter_data(data, mask):532 return data.view(-1, data.shape[-1])[mask]533 534 return _bilinear_interpolation_scattering(535 self.texture_wh[1],536 self.texture_wh[0],537 __filter_data(uv, mask)[..., [1, 0]],538 __filter_data(data, mask),539 )540 541 @torch.no_grad()542 def fast_bake_texture(543 self, textures: list[torch.Tensor], confidence_maps: list[torch.Tensor]544 ) -> tuple[torch.Tensor, torch.Tensor]:545 """Fuses multiple textures and confidence maps.546 547 Args:548 textures (list[torch.Tensor]): List of textures.549 confidence_maps (list[torch.Tensor]): List of confidence maps.550 551 Returns:552 tuple[torch.Tensor, torch.Tensor]: Fused texture and mask.553 """554 channel = textures[0].shape[-1]555 texture_merge = torch.zeros(self.texture_wh + [channel]).to(556 self.device557 )558 trust_map_merge = torch.zeros(self.texture_wh + [1]).to(self.device)559 for texture, cos_map in zip(textures, confidence_maps):560 view_sum = (cos_map > 0).sum()561 painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()562 if painted_sum / view_sum > 0.99:563 continue564 texture_merge += texture * cos_map565 trust_map_merge += cos_map566 texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)567 568 return texture_merge, trust_map_merge > 1e-8569 570 def uv_inpaint(571 self, mesh: trimesh.Trimesh, texture: np.ndarray, mask: np.ndarray572 ) -> np.ndarray:573 """Inpaints missing regions in the UV texture.574 575 Args:576 mesh (trimesh.Trimesh): Mesh.577 texture (np.ndarray): Texture image.578 mask (np.ndarray): Mask image.579 580 Returns:581 np.ndarray: Inpainted texture.582 """583 if self.inpaint_smooth:584 vertices, faces, uv_map = self.get_mesh_np_attrs(mesh)585 texture, mask = _texture_inpaint_smooth(586 texture, mask, vertices, faces, uv_map587 )588 589 texture = texture.clip(0, 1)590 texture = cv2.inpaint(591 (texture * 255).astype(np.uint8),592 255 - mask,593 3,594 cv2.INPAINT_NS,595 )596 597 return texture598 599 @spaces.GPU600 def compute_texture(601 self,602 colors: list[Image.Image],603 mesh: trimesh.Trimesh,604 ) -> trimesh.Trimesh:605 """Computes the fused texture for the mesh from multi-view images.606 607 Args:608 colors (list[Image.Image]): List of view images.609 mesh (trimesh.Trimesh): Mesh to texture.610 611 Returns:612 tuple[np.ndarray, np.ndarray]: Texture and mask.613 """614 self._lazy_init_render(self.camera_params, self.mask_thresh)615 616 vertices = torch.from_numpy(mesh.vertices).to(self.device).float()617 faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)618 uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()619 620 rendered_depth, masks = self.renderer.render_depth(vertices, faces)621 norm_deps = self.renderer.normalize_map_by_mask(rendered_depth, masks)622 render_uvs, _ = self.renderer.render_uv(vertices, faces, uv_map)623 view_normals = self.compute_enhanced_viewnormal(624 self.renderer.mv_mtx, vertices, faces625 )626 627 textures, weighted_cos_maps = [], []628 for color, mask, dep, normal, uv, weight in zip(629 colors,630 masks,631 norm_deps,632 view_normals,633 render_uvs,634 self.view_weights,635 ):636 texture, cos_map = self.back_project(color, mask, dep, normal, uv)637 textures.append(texture)638 weighted_cos_maps.append(weight * (cos_map**4))639 640 texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)641 642 texture_np = texture.cpu().numpy()643 mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)644 645 return texture_np, mask_np646 647 def __call__(648 self,649 colors: list[Image.Image],650 mesh: trimesh.Trimesh,651 output_path: str,652 ) -> trimesh.Trimesh:653 """Runs the texture baking and exports the textured mesh.654 655 Args:656 colors (list[Image.Image]): List of input view images.657 mesh (trimesh.Trimesh): Input mesh to be textured.658 output_path (str): Path to save the output textured mesh.659 660 Returns:661 trimesh.Trimesh: The textured mesh with UV and texture image.662 """663 mesh = self.load_mesh(mesh)664 texture_np, mask_np = self.compute_texture(colors, mesh)665 666 texture_np = self.uv_inpaint(mesh, texture_np, mask_np)667 if self.smooth_texture:668 texture_np = post_process_texture(texture_np)669 670 vertices, faces, uv_map = self.get_mesh_np_attrs(671 mesh, self.scale, self.center672 )673 textured_mesh = save_mesh_with_mtl(674 vertices,675 faces,676 uv_map,677 texture_np,678 output_path,679 mesh_process=self.mesh_post_process,680 )681 682 return textured_mesh683 684 685def parse_args():686 """Parses command-line arguments for texture backprojection.687 688 Returns:689 argparse.Namespace: Parsed arguments.690 """691 parser = argparse.ArgumentParser(description="Backproject texture")692 parser.add_argument(693 "--color_path",694 nargs="+",695 type=str,696 help="Multiview color image in grid file paths",697 )698 parser.add_argument(699 "--mesh_path",700 type=str,701 help="Mesh path, .obj, .glb or .ply",702 )703 parser.add_argument(704 "--output_path",705 type=str,706 help="Output mesh path with suffix",707 )708 parser.add_argument(709 "--num_images", type=int, default=6, help="Number of images to render."710 )711 parser.add_argument(712 "--elevation",713 nargs="+",714 type=float,715 default=[20.0, -10.0],716 help="Elevation angles for the camera (default: [20.0, -10.0])",717 )718 parser.add_argument(719 "--distance",720 type=float,721 default=5,722 help="Camera distance (default: 5)",723 )724 parser.add_argument(725 "--resolution_hw",726 type=int,727 nargs=2,728 default=(2048, 2048),729 help="Resolution of the output images (default: (2048, 2048))",730 )731 parser.add_argument(732 "--fov",733 type=float,734 default=30,735 help="Field of view in degrees (default: 30)",736 )737 parser.add_argument(738 "--device",739 type=str,740 choices=["cpu", "cuda"],741 default="cuda",742 help="Device to run on (default: `cuda`)",743 )744 parser.add_argument(745 "--skip_fix_mesh", action="store_true", help="Fix mesh geometry."746 )747 parser.add_argument(748 "--texture_wh",749 nargs=2,750 type=int,751 default=[2048, 2048],752 help="Texture resolution width and height",753 )754 parser.add_argument(755 "--mesh_sipmlify_ratio",756 type=float,757 default=0.9,758 help="Mesh simplification ratio (default: 0.9)",759 )760 parser.add_argument(761 "--delight", action="store_true", help="Use delighting model."762 )763 parser.add_argument(764 "--no_smooth_texture",765 action="store_true",766 help="Do not smooth the texture.",767 )768 parser.add_argument(769 "--save_glb_path", type=str, default=None, help="Save glb path."770 )771 parser.add_argument(772 "--no_save_delight_img",773 action="store_true",774 help="Disable saving delight image",775 )776 parser.add_argument("--n_max_faces", type=int, default=30000)777 parser.add_argument("--no_mesh_post_process", action="store_true")778 args, unknown = parser.parse_known_args()779 780 return args781 782 783def entrypoint(784 delight_model: DelightingModel = None,785 imagesr_model: ImageRealESRGAN = None,786 **kwargs,787) -> trimesh.Trimesh:788 """Entrypoint for texture backprojection from multi-view images.789 790 Args:791 delight_model (DelightingModel, optional): Delighting model.792 imagesr_model (ImageRealESRGAN, optional): Super-resolution model.793 **kwargs: Additional arguments to override CLI.794 795 Returns:796 trimesh.Trimesh: Textured mesh.797 """798 args = parse_args()799 for k, v in kwargs.items():800 if hasattr(args, k) and v is not None:801 setattr(args, k, v)802 803 # Setup camera parameters.804 camera_params = CameraSetting(805 num_images=args.num_images,806 elevation=args.elevation,807 distance=args.distance,808 resolution_hw=args.resolution_hw,809 fov=math.radians(args.fov),810 device=args.device,811 )812 813 args.color_path = as_list(args.color_path)814 if args.delight and delight_model is None:815 delight_model = DelightingModel()816 817 color_grid = [Image.open(color_path) for color_path in args.color_path]818 color_grid = vcat_pil_images(color_grid, image_mode="RGBA")819 if args.delight:820 color_grid = delight_model(color_grid)821 if not args.no_save_delight_img:822 save_dir = os.path.dirname(args.output_path)823 os.makedirs(save_dir, exist_ok=True)824 color_grid.save(f"{save_dir}/color_delight.png")825 826 multiviews = get_images_from_grid(color_grid, img_size=512)827 view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]828 view_weights += [0.01] * (len(multiviews) - len(view_weights))829 830 # Use RealESRGAN_x4plus for x4 (512->2048) image super resolution.831 if imagesr_model is None:832 imagesr_model = ImageRealESRGAN(outscale=4)833 multiviews = [imagesr_model(img) for img in multiviews]834 multiviews = [img.convert("RGB") for img in multiviews]835 mesh = trimesh.load(args.mesh_path)836 if isinstance(mesh, trimesh.Scene):837 mesh = mesh.dump(concatenate=True)838 839 if not args.skip_fix_mesh:840 mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)841 mesh_fixer = MeshFixer(mesh.vertices, mesh.faces, args.device)842 mesh.vertices, mesh.faces = mesh_fixer(843 filter_ratio=args.mesh_sipmlify_ratio,844 max_hole_size=0.04,845 resolution=1024,846 num_views=1000,847 norm_mesh_ratio=0.5,848 )849 if len(mesh.faces) > args.n_max_faces:850 mesh.vertices, mesh.faces = mesh_fixer(851 filter_ratio=0.8,852 max_hole_size=0.04,853 resolution=1024,854 num_views=1000,855 norm_mesh_ratio=0.5,856 )857 # Restore scale.858 mesh.vertices = mesh.vertices / scale859 mesh.vertices = mesh.vertices + center860 861 # Baking texture to mesh.862 texture_backer = TextureBacker(863 camera_params=camera_params,864 view_weights=view_weights,865 render_wh=args.resolution_hw,866 texture_wh=args.texture_wh,867 smooth_texture=not args.no_smooth_texture,868 mesh_post_process=not args.no_mesh_post_process,869 )870 871 textured_mesh = texture_backer(multiviews, mesh, args.output_path)872 873 if args.save_glb_path is not None:874 os.makedirs(os.path.dirname(args.save_glb_path), exist_ok=True)875 textured_mesh.export(args.save_glb_path)876 877 return textured_mesh878 879 880if __name__ == "__main__":881 entrypoint()882