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 logging19import math20import os21import time22import zipfile23from contextlib import contextmanager24from copy import deepcopy25from dataclasses import dataclass, field26from shutil import rmtree27from typing import List, Tuple, Union28 29import cv230import kaolin as kal31import numpy as np32import nvdiffrast.torch as dr33import torch34import torch.nn.functional as F35import trimesh36from kaolin.render.camera import Camera37from PIL import Image, ImageEnhance38 39logger = logging.getLogger(__name__)40 41 42__all__ = [43 "DiffrastRender",44 "save_images",45 "render_pbr",46 "calc_vertex_normals",47 "normalize_vertices_array",48 "as_list",49 "CameraSetting",50 "import_kaolin_mesh",51 "save_mesh_with_mtl",52 "get_images_from_grid",53 "post_process_texture",54 "quat_mult",55 "quat_to_rotmat",56 "gamma_shs",57 "resize_pil",58 "trellis_preprocess",59 "delete_dir",60 "kaolin_to_opencv_view",61 "model_device_ctx",62]63 64 65class DiffrastRender(object):66 """A class to handle differentiable rendering using nvdiffrast.67 68 This class provides methods to render position, depth, and normal maps69 with optional anti-aliasing and gradient disabling for rasterization.70 71 Attributes:72 p_mtx (torch.Tensor): Projection matrix.73 mv_mtx (torch.Tensor): Model-view matrix.74 mvp_mtx (torch.Tensor): Model-view-projection matrix, calculated as75 p_mtx @ mv_mtx if not provided.76 resolution_hw (Tuple[int, int]): Height and width of the rendering resolution. # noqa77 _ctx (Union[dr.RasterizeCudaContext, dr.RasterizeGLContext]): Rasterization context. # noqa78 mask_thresh (float): Threshold for mask creation.79 grad_db (bool): Whether to disable gradients during rasterization.80 antialias_mask (bool): Whether to apply anti-aliasing to the mask.81 device (str): Device used for rendering ('cuda' or 'cpu').82 """83 84 def __init__(85 self,86 p_matrix: torch.Tensor,87 mv_matrix: torch.Tensor,88 resolution_hw: Tuple[int, int],89 context: Union[dr.RasterizeCudaContext, dr.RasterizeGLContext] = None,90 mvp_matrix: torch.Tensor = None,91 mask_thresh: float = 0.5,92 grad_db: bool = False,93 antialias_mask: bool = True,94 align_coordinate: bool = True,95 device: str = "cuda",96 ) -> None:97 self.p_mtx = p_matrix98 self.mv_mtx = mv_matrix99 if mvp_matrix is None:100 self.mvp_mtx = torch.bmm(p_matrix, mv_matrix)101 102 self.resolution_hw = resolution_hw103 if context is None:104 context = dr.RasterizeCudaContext(device=device)105 self._ctx = context106 self.mask_thresh = mask_thresh107 self.grad_db = grad_db108 self.antialias_mask = antialias_mask109 self.align_coordinate = align_coordinate110 self.device = device111 112 def compute_dr_raster(113 self,114 vertices: torch.Tensor,115 faces: torch.Tensor,116 ) -> Tuple[torch.Tensor, torch.Tensor]:117 vertices_clip = self.transform_vertices(vertices, matrix=self.mvp_mtx)118 rast, _ = dr.rasterize(119 self._ctx,120 vertices_clip,121 faces.int(),122 resolution=self.resolution_hw,123 grad_db=self.grad_db,124 )125 126 return rast, vertices_clip127 128 def transform_vertices(129 self,130 vertices: torch.Tensor,131 matrix: torch.Tensor,132 ) -> torch.Tensor:133 verts_ones = torch.ones(134 (len(vertices), 1), device=vertices.device, dtype=vertices.dtype135 )136 verts_homo = torch.cat([vertices, verts_ones], dim=-1)137 trans_vertices = torch.matmul(verts_homo, matrix.permute(0, 2, 1))138 139 return trans_vertices140 141 def normalize_map_by_mask_separately(142 self, map: torch.Tensor, mask: torch.Tensor143 ) -> torch.Tensor:144 # Normalize each map separately by mask, normalized map in [0, 1].145 normalized_maps = []146 for map_item, mask_item in zip(map, mask):147 normalized_map = self.normalize_map_by_mask(map_item, mask_item)148 normalized_maps.append(normalized_map)149 150 normalized_maps = torch.stack(normalized_maps, dim=0)151 152 return normalized_maps153 154 @staticmethod155 def normalize_map_by_mask(156 map: torch.Tensor, mask: torch.Tensor157 ) -> torch.Tensor:158 # Normalize all maps in total by mask, normalized map in [0, 1].159 foreground = (mask == 1).squeeze(dim=-1)160 foreground_elements = map[foreground]161 if len(foreground_elements) == 0:162 return map163 164 min_val, _ = foreground_elements.min(dim=0)165 max_val, _ = foreground_elements.max(dim=0)166 val_range = (max_val - min_val).clip(min=1e-6)167 168 normalized_map = (map - min_val) / val_range169 normalized_map = torch.lerp(170 torch.zeros_like(normalized_map), normalized_map, mask171 )172 normalized_map[normalized_map < 0] = 0173 174 return normalized_map175 176 def _compute_mask(177 self,178 rast: torch.Tensor,179 vertices_clip: torch.Tensor,180 faces: torch.Tensor,181 ) -> torch.Tensor:182 mask = (rast[..., 3:] > 0).float()183 mask = mask.clip(min=0, max=1)184 185 if self.antialias_mask is True:186 mask = dr.antialias(mask, rast, vertices_clip, faces)187 else:188 foreground = mask > self.mask_thresh189 mask[foreground] = 1190 mask[~foreground] = 0191 192 return mask193 194 def render_rast_alpha(195 self,196 vertices: torch.Tensor,197 faces: torch.Tensor,198 ):199 faces = faces.to(torch.int32)200 rast, vertices_clip = self.compute_dr_raster(vertices, faces)201 mask = self._compute_mask(rast, vertices_clip, faces)202 203 return mask, rast204 205 def render_position(206 self,207 vertices: torch.Tensor,208 faces: torch.Tensor,209 ) -> Union[torch.Tensor, torch.Tensor]:210 # Vertices in model coordinate system, real position coordinate number.211 faces = faces.to(torch.int32)212 mask, rast = self.render_rast_alpha(vertices, faces)213 214 vertices_model = vertices[None, ...].contiguous().float()215 position_map, _ = dr.interpolate(vertices_model, rast, faces)216 # Align with blender.217 if self.align_coordinate:218 position_map = position_map[..., [0, 2, 1]]219 position_map[..., 1] = -position_map[..., 1]220 221 position_map = torch.lerp(222 torch.zeros_like(position_map), position_map, mask223 )224 225 return position_map, mask226 227 def render_uv(228 self,229 vertices: torch.Tensor,230 faces: torch.Tensor,231 vtx_uv: torch.Tensor,232 ) -> Union[torch.Tensor, torch.Tensor]:233 faces = faces.to(torch.int32)234 mask, rast = self.render_rast_alpha(vertices, faces)235 uv_map, _ = dr.interpolate(vtx_uv, rast, faces)236 uv_map = torch.lerp(torch.zeros_like(uv_map), uv_map, mask)237 238 return uv_map, mask239 240 def render_depth(241 self,242 vertices: torch.Tensor,243 faces: torch.Tensor,244 ) -> Union[torch.Tensor, torch.Tensor]:245 # Vertices in model coordinate system, real depth coordinate number.246 faces = faces.to(torch.int32)247 mask, rast = self.render_rast_alpha(vertices, faces)248 249 vertices_camera = self.transform_vertices(vertices, matrix=self.mv_mtx)250 vertices_camera = vertices_camera[..., 2:3].contiguous().float()251 depth_map, _ = dr.interpolate(vertices_camera, rast, faces)252 # Change camera depth minus to positive.253 if self.align_coordinate:254 depth_map = -depth_map255 depth_map = torch.lerp(torch.zeros_like(depth_map), depth_map, mask)256 257 return depth_map, mask258 259 def render_global_normal(260 self,261 vertices: torch.Tensor,262 faces: torch.Tensor,263 vertice_normals: torch.Tensor,264 ) -> Union[torch.Tensor, torch.Tensor]:265 # NOTE: vertice_normals in [-1, 1], return normal in [0, 1].266 # vertices / vertice_normals in model coordinate system.267 faces = faces.to(torch.int32)268 mask, rast = self.render_rast_alpha(vertices, faces)269 im_base_normals, _ = dr.interpolate(270 vertice_normals[None, ...].float(), rast, faces271 )272 273 if im_base_normals is not None:274 faces = faces.to(torch.int64)275 vertices_cam = self.transform_vertices(276 vertices, matrix=self.mv_mtx277 )278 face_vertices_ndc = kal.ops.mesh.index_vertices_by_faces(279 vertices_cam[..., :3], faces280 )281 face_normal_sign = kal.ops.mesh.face_normals(face_vertices_ndc)[282 ..., 2283 ]284 for idx in range(len(im_base_normals)):285 face_idx = (rast[idx, ..., -1].long() - 1).contiguous()286 im_normal_sign = torch.sign(face_normal_sign[idx, face_idx])287 im_normal_sign[face_idx == -1] = 0288 im_base_normals[idx] *= im_normal_sign.unsqueeze(-1)289 290 normal = (im_base_normals + 1) / 2291 normal = normal.clip(min=0, max=1)292 normal = torch.lerp(torch.zeros_like(normal), normal, mask)293 294 return normal, mask295 296 def transform_normal(297 self,298 normals: torch.Tensor,299 trans_matrix: torch.Tensor,300 masks: torch.Tensor,301 to_view: bool,302 ) -> torch.Tensor:303 # NOTE: input normals in [0, 1], output normals in [0, 1].304 normals = normals.clone()305 assert len(normals) == len(trans_matrix)306 307 if not to_view:308 # Flip the sign on the x-axis to match inv bae system for global transformation. # noqa309 normals[..., 0] = 1 - normals[..., 0]310 311 normals = 2 * normals - 1312 b, h, w, c = normals.shape313 314 transformed_normals = []315 for normal, matrix in zip(normals, trans_matrix):316 # Transform normals using the transformation matrix (4x4).317 reshaped_normals = normal.view(-1, c) # (h w 3) -> (hw 3)318 padded_vectors = torch.nn.functional.pad(319 reshaped_normals, pad=(0, 1), mode="constant", value=0.0320 )321 transformed_normal = torch.matmul(322 padded_vectors, matrix.transpose(0, 1)323 )[..., :3]324 325 # Normalize and clip the normals to [0, 1] range.326 transformed_normal = F.normalize(transformed_normal, p=2, dim=-1)327 transformed_normal = (transformed_normal + 1) / 2328 329 if to_view:330 # Flip the sign on the x-axis to match bae system for view transformation. # noqa331 transformed_normal[..., 0] = 1 - transformed_normal[..., 0]332 333 transformed_normals.append(transformed_normal.view(h, w, c))334 335 transformed_normals = torch.stack(transformed_normals, dim=0)336 337 if masks is not None:338 transformed_normals = torch.lerp(339 torch.zeros_like(transformed_normals),340 transformed_normals,341 masks,342 )343 344 return transformed_normals345 346 347def _az_el_to_points(348 azimuths: np.ndarray, elevations: np.ndarray349) -> np.ndarray:350 x = np.cos(azimuths) * np.cos(elevations)351 y = np.sin(azimuths) * np.cos(elevations)352 z = np.sin(elevations)353 354 return np.stack([x, y, z], axis=-1)355 356 357def _compute_az_el_by_views(358 num_view: int, el: float359) -> Tuple[np.ndarray, np.ndarray]:360 azimuths = np.arange(num_view) / num_view * np.pi * 2361 elevations = np.deg2rad(np.array([el] * num_view))362 363 return azimuths, elevations364 365 366def _compute_cam_pts_by_az_el(367 azs: np.ndarray,368 els: np.ndarray,369 distance: float | list[float] | np.ndarray,370 extra_pts: np.ndarray = None,371) -> np.ndarray:372 if np.isscalar(distance) or isinstance(distance, (float, int)):373 distances = np.full(len(azs), distance)374 else:375 distances = np.array(distance)376 if len(distances) != len(azs):377 raise ValueError(378 f"Length of distances ({len(distances)}) must match length of azs ({len(azs)})"379 )380 381 cam_pts = _az_el_to_points(azs, els) * distances[:, None]382 383 if extra_pts is not None:384 cam_pts = np.concatenate([cam_pts, extra_pts], axis=0)385 386 # Align coordinate system.387 cam_pts = cam_pts[:, [0, 2, 1]] # xyz -> xzy388 cam_pts[..., 2] = -cam_pts[..., 2]389 390 return cam_pts391 392 393def compute_cam_pts_by_views(394 num_view: int, el: float, distance: float, extra_pts: np.ndarray = None395) -> torch.Tensor:396 """Computes object-center camera points for a given number of views.397 398 Args:399 num_view (int): The number of views (camera positions) to compute.400 el (float): The elevation angle in degrees.401 distance (float): The distance from the origin to the camera.402 extra_pts (np.ndarray): Extra camera points postion.403 404 Returns:405 torch.Tensor: A tensor containing the camera points for each view, with shape `(num_view, 3)`. # noqa406 """407 azimuths, elevations = _compute_az_el_by_views(num_view, el)408 cam_pts = _compute_cam_pts_by_az_el(409 azimuths, elevations, distance, extra_pts410 )411 412 return cam_pts413 414 415def save_images(416 images: Union[list[np.ndarray], list[torch.Tensor]],417 output_dir: str,418 cvt_color: str = None,419 format: str = ".png",420 to_uint8: bool = True,421 verbose: bool = False,422) -> List[str]:423 # NOTE: images in [0, 1]424 os.makedirs(output_dir, exist_ok=True)425 save_paths = []426 for idx, image in enumerate(images):427 if isinstance(image, torch.Tensor):428 image = image.detach().cpu().numpy()429 if to_uint8:430 image = image.clip(min=0, max=1)431 image = (255.0 * image).astype(np.uint8)432 if cvt_color is not None:433 image = cv2.cvtColor(image, cvt_color)434 save_path = os.path.join(output_dir, f"{idx:04d}{format}")435 save_paths.append(save_path)436 437 cv2.imwrite(save_path, image)438 439 if verbose:440 logger.info(f"Images saved in {output_dir}")441 442 return save_paths443 444 445def _disable_metallic_for_render(materials):446 if materials is None:447 return448 449 for material in materials:450 if hasattr(material, "metallic_texture"):451 material.metallic_texture = None452 if (453 hasattr(material, "metallic_value")454 and material.metallic_value is not None455 ):456 if torch.is_tensor(material.metallic_value):457 material.metallic_value = torch.zeros_like(458 material.metallic_value459 )460 else:461 material.metallic_value = 0.0462 463 464def _build_render_materials(mesh, metallic: bool = False):465 if metallic:466 return None467 468 if mesh.materials is None:469 return None470 471 render_materials = deepcopy(mesh.materials)472 _disable_metallic_for_render(render_materials)473 return render_materials474 475 476def _current_lighting(477 azimuths: List[float],478 elevations: List[float],479 light_factor: float = 1.0,480 device: str = "cuda",481):482 # azimuths, elevations in degress.483 directions = []484 for az, el in zip(azimuths, elevations):485 az, el = math.radians(az), math.radians(el)486 direction = kal.render.lighting.sg_direction_from_azimuth_elevation(487 az, el488 )489 directions.append(direction)490 directions = torch.cat(directions, dim=0)491 492 amplitude = torch.ones_like(directions) * light_factor493 light_condition = kal.render.lighting.SgLightingParameters(494 amplitude=amplitude,495 direction=directions,496 sharpness=3,497 ).to(device)498 499 # light_condition = kal.render.lighting.SgLightingParameters.from_sun(500 # directions, strength=1, angle=90, color=None501 # ).to(device)502 503 return light_condition504 505 506def _uniform_lighting(507 light_factor: float = 1.0,508 device: str = "cuda",509 sharpness: float = 0.5,510 num_lights: int = 1024,511):512 indices = torch.arange(num_lights, dtype=torch.float32, device=device)513 golden_angle = math.pi * (3.0 - math.sqrt(5.0))514 z = 1.0 - 2.0 * (indices + 0.5) / num_lights515 radius = torch.sqrt(torch.clamp(1.0 - z * z, min=0.0))516 theta = golden_angle * indices517 directions = torch.stack(518 [519 radius * torch.cos(theta),520 radius * torch.sin(theta),521 z,522 ],523 dim=1,524 )525 directions = F.normalize(directions, dim=1)526 527 amplitude = torch.ones_like(directions) * (light_factor / len(directions))528 light_condition = kal.render.lighting.SgLightingParameters(529 amplitude=amplitude,530 direction=directions,531 sharpness=sharpness,532 ).to(device)533 534 return light_condition535 536 537def render_pbr(538 mesh,539 camera,540 device="cuda",541 cxt=None,542 light_factor=1.0,543 metallic: bool = False,544):545 if cxt is None:546 cxt = dr.RasterizeCudaContext()547 548 light_condition = _uniform_lighting(549 light_factor=light_factor,550 device=device,551 )552 render_materials = _build_render_materials(mesh, metallic)553 render_res = kal.render.easy_render.render_mesh(554 camera,555 mesh,556 lighting=light_condition,557 nvdiffrast_context=cxt,558 custom_materials=render_materials,559 )560 561 image = render_res[kal.render.easy_render.RenderPass.render]562 image = image.clip(0, 1)563 564 albedo = render_res[kal.render.easy_render.RenderPass.albedo]565 albedo = albedo.clip(0, 1)566 567 diffuse = render_res[kal.render.easy_render.RenderPass.diffuse]568 diffuse = diffuse.clip(0, 1)569 570 normal = render_res[kal.render.easy_render.RenderPass.normals]571 normal = normal.clip(-1, 1)572 573 return image, albedo, diffuse, normal574 575 576def _calc_face_normals(577 vertices: torch.Tensor, # V,3 first vertex may be unreferenced578 faces: torch.Tensor, # F,3 long, first face may be all zero579 normalize: bool = False,580) -> torch.Tensor: # F,3581 full_vertices = vertices[faces] # F,C=3,3582 v0, v1, v2 = full_vertices.unbind(dim=1) # F,3583 face_normals = torch.cross(v1 - v0, v2 - v0, dim=1) # F,3584 if normalize:585 face_normals = F.normalize(586 face_normals, eps=1e-6, dim=1587 ) # TODO inplace?588 return face_normals # F,3589 590 591def calc_vertex_normals(592 vertices: torch.Tensor, # V,3 first vertex may be unreferenced593 faces: torch.Tensor, # F,3 long, first face may be all zero594 face_normals: torch.Tensor = None, # F,3, not normalized595) -> torch.Tensor: # F,3596 _F = faces.shape[0]597 598 if face_normals is None:599 face_normals = _calc_face_normals(vertices, faces)600 601 vertex_normals = torch.zeros(602 (vertices.shape[0], 3, 3), dtype=vertices.dtype, device=vertices.device603 ) # V,C=3,3604 vertex_normals.scatter_add_(605 dim=0,606 index=faces[:, :, None].expand(_F, 3, 3),607 src=face_normals[:, None, :].expand(_F, 3, 3),608 )609 vertex_normals = vertex_normals.sum(dim=1) # V,3610 return F.normalize(vertex_normals, eps=1e-6, dim=1)611 612 613def normalize_vertices_array(614 vertices: Union[torch.Tensor, np.ndarray],615 mesh_scale: float = 1.0,616 exec_norm: bool = True,617):618 if isinstance(vertices, torch.Tensor):619 bbmin, bbmax = vertices.min(0)[0], vertices.max(0)[0]620 else:621 bbmin, bbmax = vertices.min(0), vertices.max(0) # (3,)622 center = (bbmin + bbmax) * 0.5623 bbsize = bbmax - bbmin624 scale = 2 * mesh_scale / bbsize.max()625 if exec_norm:626 vertices = (vertices - center) * scale627 628 return vertices, scale, center629 630 631def as_list(obj):632 if isinstance(obj, (list, tuple)):633 return obj634 elif isinstance(obj, set):635 return list(obj)636 elif obj is None:637 return obj638 else:639 return [obj]640 641 642@dataclass643class CameraSetting:644 """Camera settings for images rendering."""645 646 num_images: int647 elevation: list[float]648 distance: float | list[float]649 resolution_hw: tuple[int, int]650 fov: float651 at: tuple[float, float, float] = field(652 default_factory=lambda: (0.0, 0.0, 0.0)653 )654 up: tuple[float, float, float] = field(655 default_factory=lambda: (0.0, 1.0, 0.0)656 )657 device: str = "cuda"658 near: float = 1e-2659 far: float = 1e2660 661 def __post_init__(662 self,663 ):664 h = self.resolution_hw[0]665 f = (h / 2) / math.tan(self.fov / 2)666 cx = self.resolution_hw[1] / 2667 cy = self.resolution_hw[0] / 2668 Ks = [669 [f, 0, cx],670 [0, f, cy],671 [0, 0, 1],672 ]673 674 self.Ks = Ks675 676 677def _compute_az_el_by_camera_params(678 camera_params: CameraSetting, flip_az: bool = False679):680 num_view = camera_params.num_images // len(camera_params.elevation)681 view_interval = 2 * np.pi / num_view / 2682 if num_view == 1:683 view_interval = np.pi / 2684 azimuths = []685 elevations = []686 for idx, el in enumerate(camera_params.elevation):687 azs = np.arange(num_view) / num_view * np.pi * 2 + idx * view_interval688 if flip_az:689 azs *= -1690 els = np.deg2rad(np.array([el] * num_view))691 azimuths.append(azs)692 elevations.append(els)693 694 azimuths = np.concatenate(azimuths, axis=0)695 elevations = np.concatenate(elevations, axis=0)696 697 return azimuths, elevations698 699 700def init_kal_camera(701 camera_params: CameraSetting,702 flip_az: bool = False,703) -> Camera:704 azimuths, elevations = _compute_az_el_by_camera_params(705 camera_params, flip_az706 )707 cam_pts = _compute_cam_pts_by_az_el(708 azimuths, elevations, camera_params.distance709 )710 711 up = torch.cat(712 [713 torch.tensor(camera_params.up).repeat(camera_params.num_images, 1),714 ],715 dim=0,716 )717 718 camera = Camera.from_args(719 eye=torch.tensor(cam_pts),720 at=torch.tensor(camera_params.at),721 up=up,722 fov=camera_params.fov,723 height=camera_params.resolution_hw[0],724 width=camera_params.resolution_hw[1],725 near=camera_params.near,726 far=camera_params.far,727 device=camera_params.device,728 )729 730 return camera731 732 733def import_kaolin_mesh(mesh_path: str, with_mtl: bool = False):734 if mesh_path.endswith(".glb"):735 mesh = kal.io.gltf.import_mesh(mesh_path)736 elif mesh_path.endswith(".obj"):737 with_material = True if with_mtl else False738 mesh = kal.io.obj.import_mesh(mesh_path, with_materials=with_material)739 if with_mtl and mesh.materials and len(mesh.materials) > 0:740 material = kal.render.materials.PBRMaterial()741 assert "map_Kd" in mesh.materials[0], (742 "'map_Kd' not found in materials."743 )744 material.diffuse_texture = mesh.materials[0]["map_Kd"] / 255.0745 mesh.materials = [material]746 elif mesh_path.endswith(".ply"):747 mesh = trimesh.load(mesh_path)748 mesh_path = mesh_path.replace(".ply", ".obj")749 mesh.export(mesh_path)750 mesh = kal.io.obj.import_mesh(mesh_path)751 elif mesh_path.endswith(".off"):752 mesh = kal.io.off.import_mesh(mesh_path)753 else:754 raise RuntimeError(755 f"{mesh_path} mesh type not supported, "756 "supported mesh type `.glb`, `.obj`, `.ply`, `.off`."757 )758 759 return mesh760 761 762def kaolin_to_opencv_view(raw_matrix):763 R_orig = raw_matrix[:, :3, :3]764 t_orig = raw_matrix[:, :3, 3]765 766 R_target = torch.zeros_like(R_orig)767 R_target[:, :, 0] = R_orig[:, :, 2]768 R_target[:, :, 1] = R_orig[:, :, 0]769 R_target[:, :, 2] = R_orig[:, :, 1]770 771 t_target = t_orig772 773 target_matrix = (774 torch.eye(4, device=raw_matrix.device)775 .unsqueeze(0)776 .repeat(raw_matrix.size(0), 1, 1)777 )778 target_matrix[:, :3, :3] = R_target779 target_matrix[:, :3, 3] = t_target780 781 return target_matrix782 783 784def save_mesh_with_mtl(785 vertices: np.ndarray,786 faces: np.ndarray,787 uvs: np.ndarray,788 texture: Union[Image.Image, np.ndarray],789 output_path: str,790 material_base=(250, 250, 250, 255),791 mesh_process: bool = True,792 glossiness: float = 250.0,793) -> trimesh.Trimesh:794 if isinstance(texture, np.ndarray):795 texture = Image.fromarray(texture)796 797 mesh = trimesh.Trimesh(798 vertices,799 faces,800 visual=trimesh.visual.TextureVisuals(uv=uvs, image=texture),801 process=mesh_process, # True for preventing modification of vertices802 )803 mesh.visual.material = trimesh.visual.material.SimpleMaterial(804 image=texture,805 diffuse=material_base,806 ambient=material_base,807 specular=material_base,808 # 250 gives a tight visible highlight similar to glossy plastic.809 glossiness=glossiness,810 )811 812 dir_name = os.path.dirname(output_path)813 os.makedirs(dir_name, exist_ok=True)814 815 _ = mesh.export(output_path)816 # texture.save(os.path.join(dir_name, f"{file_name}_texture.png"))817 818 logger.info(f"Saved mesh with texture to {output_path}")819 820 return mesh821 822 823def get_images_from_grid(824 image: Union[str, Image.Image], img_size: int825) -> list[Image.Image]:826 if isinstance(image, str):827 image = Image.open(image)828 829 view_images = np.array(image)830 height, width, _ = view_images.shape831 rows = height // img_size832 cols = width // img_size833 blocks = []834 for i in range(rows):835 for j in range(cols):836 block = view_images[837 i * img_size : (i + 1) * img_size,838 j * img_size : (j + 1) * img_size,839 :,840 ]841 blocks.append(Image.fromarray(block))842 843 return blocks844 845 846def enhance_image(847 image: Image.Image,848 contrast_factor: float = 1.3,849 color_factor: float = 1.2,850 brightness_factor: float = 0.95,851) -> Image.Image:852 enhancer_contrast = ImageEnhance.Contrast(image)853 img_contrasted = enhancer_contrast.enhance(contrast_factor)854 855 enhancer_color = ImageEnhance.Color(img_contrasted)856 img_colored = enhancer_color.enhance(color_factor)857 858 enhancer_brightness = ImageEnhance.Brightness(img_colored)859 enhanced_image = enhancer_brightness.enhance(brightness_factor)860 861 return enhanced_image862 863 864def post_process_texture(texture: np.ndarray, iter: int = 1) -> np.ndarray:865 for _ in range(iter):866 texture = cv2.fastNlMeansDenoisingColored(texture, None, 2, 2, 7, 15)867 texture = cv2.bilateralFilter(868 texture, d=5, sigmaColor=20, sigmaSpace=20869 )870 871 texture = enhance_image(872 image=Image.fromarray(texture),873 contrast_factor=1.3,874 color_factor=1.2,875 brightness_factor=0.95,876 )877 878 return np.array(texture)879 880 881def quat_mult(q1, q2):882 # NOTE:883 # Q1 is the quaternion that rotates the vector from the original position to the final position # noqa884 # Q2 is the quaternion that been rotated885 w1, x1, y1, z1 = q1.T886 w2, x2, y2, z2 = q2.T887 w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2888 x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2889 y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2890 z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2891 return torch.stack([w, x, y, z]).T892 893 894def quat_to_rotmat(quats: torch.Tensor, mode="wxyz") -> torch.Tensor:895 """Convert quaternion to rotation matrix."""896 quats = F.normalize(quats, p=2, dim=-1)897 898 if mode == "xyzw":899 x, y, z, w = torch.unbind(quats, dim=-1)900 elif mode == "wxyz":901 w, x, y, z = torch.unbind(quats, dim=-1)902 else:903 raise ValueError(f"Invalid mode: {mode}.")904 905 R = torch.stack(906 [907 1 - 2 * (y**2 + z**2),908 2 * (x * y - w * z),909 2 * (x * z + w * y),910 2 * (x * y + w * z),911 1 - 2 * (x**2 + z**2),912 2 * (y * z - w * x),913 2 * (x * z - w * y),914 2 * (y * z + w * x),915 1 - 2 * (x**2 + y**2),916 ],917 dim=-1,918 )919 920 return R.reshape(quats.shape[:-1] + (3, 3))921 922 923def gamma_shs(shs: torch.Tensor, gamma: float) -> torch.Tensor:924 C0 = 0.28209479177387814 # Constant for normalization in spherical harmonics # noqa925 # Clip to the range [0.0, 1.0], apply gamma correction, and then un-clip back # noqa926 new_shs = torch.clip(shs * C0 + 0.5, 0.0, 1.0)927 new_shs = (torch.pow(new_shs, gamma) - 0.5) / C0928 return new_shs929 930 931def resize_pil(image: Image.Image, max_size: int = 1024) -> Image.Image:932 current_max_dim = max(image.size)933 scale = min(1, max_size / current_max_dim)934 935 if scale < 1:936 new_size = (int(image.width * scale), int(image.height * scale))937 image = image.resize(new_size, Image.Resampling.LANCZOS)938 939 return image940 941 942def trellis_preprocess(image: Image.Image) -> Image.Image:943 """Process the input image as trellis done."""944 image_np = np.array(image)945 alpha = image_np[:, :, 3]946 bbox = np.argwhere(alpha > 0.8 * 255)947 bbox = (948 np.min(bbox[:, 1]),949 np.min(bbox[:, 0]),950 np.max(bbox[:, 1]),951 np.max(bbox[:, 0]),952 )953 center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2954 size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])955 size = int(size * 1.2)956 bbox = (957 center[0] - size // 2,958 center[1] - size // 2,959 center[0] + size // 2,960 center[1] + size // 2,961 )962 image = image.crop(bbox)963 image = image.resize((518, 518), Image.Resampling.LANCZOS)964 image = np.array(image).astype(np.float32) / 255965 image = image[:, :, :3] * image[:, :, 3:4]966 image = Image.fromarray((image * 255).astype(np.uint8))967 968 return image969 970 971def zip_files(input_paths: list[str], output_zip: str) -> str:972 with zipfile.ZipFile(output_zip, "w", zipfile.ZIP_DEFLATED) as zipf:973 for input_path in input_paths:974 if not os.path.exists(input_path):975 raise FileNotFoundError(f"File not found: {input_path}")976 977 if os.path.isdir(input_path):978 for root, _, files in os.walk(input_path):979 for file in files:980 file_path = os.path.join(root, file)981 arcname = os.path.relpath(982 file_path, start=os.path.commonpath(input_paths)983 )984 zipf.write(file_path, arcname=arcname)985 else:986 arcname = os.path.relpath(987 input_path, start=os.path.commonpath(input_paths)988 )989 zipf.write(input_path, arcname=arcname)990 991 return output_zip992 993 994def delete_dir(folder_path: str, keep_subs: list[str] = None) -> None:995 for item in os.listdir(folder_path):996 if keep_subs is not None and item in keep_subs:997 continue998 item_path = os.path.join(folder_path, item)999 if os.path.isdir(item_path):1000 rmtree(item_path)1001 else:1002 os.remove(item_path)1003 1004 1005@contextmanager1006def model_device_ctx(1007 *models,1008 src_device: str = "cpu",1009 dst_device: str = "cuda",1010 verbose: bool = False,1011):1012 start = time.perf_counter()1013 for m in models:1014 if m is None:1015 continue1016 m.to(dst_device)1017 to_cuda_time = time.perf_counter() - start1018 1019 try:1020 yield1021 finally:1022 start = time.perf_counter()1023 for m in models:1024 if m is None:1025 continue1026 m.to(src_device)1027 to_cpu_time = time.perf_counter() - start1028 1029 if verbose:1030 model_names = [m.__class__.__name__ for m in models]1031 logger.info(1032 f"[model_device_ctx] {model_names} to cuda: {to_cuda_time:.1f}s, to cpu: {to_cpu_time:.1f}s"1033 )1034