tohid4n/PartCrafter
0
1from src.utils.typing_utils import *2 3import os4import numpy as np5from PIL import Image6import trimesh7from trimesh.transformations import rotation_matrix8import pyrender9from diffusers.utils import export_to_video10from diffusers.utils.loading_utils import load_video11import torch12from torchvision.utils import make_grid13import math14 15os.environ['PYOPENGL_PLATFORM'] = 'egl'16 17def explode_mesh(mesh, explosion_scale=0.4):18 # ensure we have a Scene19 if isinstance(mesh, trimesh.Trimesh):20 scene = trimesh.Scene(mesh)21 elif isinstance(mesh, trimesh.Scene):22 scene = mesh23 else:24 raise ValueError(f"Expected Trimesh or Scene, got {type(mesh)}")25 26 if len(scene.geometry) <= 1:27 print("Nothing to explode")28 return scene29 30 # 1) collect (name, geom, world_center)31 parts = []32 for name, geom in scene.geometry.items():33 # ← get(name) returns (4×4 world‐space matrix, parent_frame)34 world_tf, _ = scene.graph.get(name)35 pts = trimesh.transformations.transform_points(geom.vertices, world_tf)36 center = pts.mean(axis=0)37 parts.append((name, geom, center))38 39 # compute global center40 all_centers = np.stack([c for _,_,c in parts], axis=0)41 global_center = all_centers.mean(axis=0)42 43 exploded = trimesh.Scene()44 for name, geom, center in parts:45 dir_vec = center - global_center46 norm = np.linalg.norm(dir_vec)47 if norm < 1e-6:48 dir_vec = np.random.randn(3)49 dir_vec /= np.linalg.norm(dir_vec)50 else:51 dir_vec /= norm52 53 offset = dir_vec * explosion_scale54 55 # fetch the same 4×4, then bump just the translation56 world_tf, _ = scene.graph.get(name)57 world_tf = world_tf.copy()58 world_tf[:3, 3] += offset59 60 exploded.add_geometry(geom, transform=world_tf, geom_name=name)61 print(f"[explode] {name} moved by {np.linalg.norm(offset):.4f}")62 63 return exploded64 65 66 67def render(68 scene: pyrender.Scene,69 renderer: pyrender.Renderer,70 camera: pyrender.Camera,71 pose: np.ndarray,72 light: Optional[pyrender.Light] = None,73 normalize_depth: bool = False,74 flags: int = pyrender.constants.RenderFlags.NONE,75 return_type: Literal['pil', 'ndarray'] = 'pil'76) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[Image.Image, Image.Image]]:77 camera_node = scene.add(camera, pose=pose)78 if light is not None:79 light_node = scene.add(light, pose=pose)80 image, depth = renderer.render(81 scene, 82 flags=flags83 )84 scene.remove_node(camera_node)85 if light is not None:86 scene.remove_node(light_node)87 if normalize_depth or return_type == 'pil':88 depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.089 if return_type == 'pil':90 image = Image.fromarray(image)91 depth = Image.fromarray(depth.astype(np.uint8))92 return image, depth93 94def rotation_matrix_from_vectors(vec1, vec2):95 a, b = vec1 / np.linalg.norm(vec1), vec2 / np.linalg.norm(vec2)96 v = np.cross(a, b)97 c = np.dot(a, b)98 s = np.linalg.norm(v)99 if s == 0:100 return np.eye(3) if c > 0 else -np.eye(3)101 kmat = np.array([102 [0, -v[2], v[1]],103 [v[2], 0, -v[0]],104 [-v[1], v[0], 0]105 ])106 return np.eye(3) + kmat + kmat @ kmat * ((1 - c) / (s ** 2))107 108def create_circular_camera_positions(109 num_views: int,110 radius: float,111 axis: np.ndarray = np.array([0.0, 1.0, 0.0])112) -> List[np.ndarray]:113 # Create a list of positions for a circular camera trajectory114 # around the given axis with the given radius.115 positions = []116 axis = axis / np.linalg.norm(axis)117 for i in range(num_views):118 theta = 2 * np.pi * i / num_views119 position = np.array([120 np.sin(theta) * radius,121 0.0,122 np.cos(theta) * radius123 ])124 if not np.allclose(axis, np.array([0.0, 1.0, 0.0])):125 R = rotation_matrix_from_vectors(np.array([0.0, 1.0, 0.0]), axis)126 position = R @ position127 positions.append(position)128 return positions129 130def create_circular_camera_poses(131 num_views: int,132 radius: float,133 axis: np.ndarray = np.array([0.0, 1.0, 0.0])134) -> List[np.ndarray]:135 # Create a list of poses for a circular camera trajectory136 # around the given axis with the given radius.137 # The camera always looks at the origin.138 # The up vector is always [0, 1, 0].139 canonical_pose = np.array([140 [1.0, 0.0, 0.0, 0.0],141 [0.0, 1.0, 0.0, 0.0],142 [0.0, 0.0, 1.0, radius],143 [0.0, 0.0, 0.0, 1.0]144 ])145 poses = []146 for i in range(num_views):147 theta = 2 * np.pi * i / num_views148 R = rotation_matrix(149 angle=theta,150 direction=axis,151 point=[0, 0, 0]152 )153 pose = R @ canonical_pose154 poses.append(pose)155 return poses156 157def render_views_around_mesh(158 mesh: Union[trimesh.Trimesh, trimesh.Scene],159 num_views: int = 36,160 radius: float = 3.5,161 axis: np.ndarray = np.array([0.0, 1.0, 0.0]),162 image_size: tuple = (512, 512),163 fov: float = 40.0,164 light_intensity: Optional[float] = 5.0,165 znear: float = 0.1,166 zfar: float = 10.0, 167 normalize_depth: bool = False,168 flags: int = pyrender.constants.RenderFlags.NONE,169 return_depth: bool = False, 170 return_type: Literal['pil', 'ndarray'] = 'pil'171) -> Union[172 List[Image.Image], 173 List[np.ndarray], 174 Tuple[List[Image.Image], List[Image.Image]], 175 Tuple[List[np.ndarray], List[np.ndarray]]176 ]:177 178 meshes = []179 scenes = []180 181 if not isinstance(mesh, (trimesh.Trimesh, trimesh.Scene)):182 raise ValueError("mesh must be a trimesh.Trimesh or trimesh.Scene object")183 if isinstance(mesh, trimesh.Trimesh):184 for i in range(num_views):185 scenes.append(pyrender.Scene.from_trimesh_scene(trimesh.Scene(mesh)))186 else:187 for i in range(num_views):188 value = math.sin(math.pi * (i - 1) / num_views)189 scenes.append(pyrender.Scene.from_trimesh_scene(explode_mesh(mesh, 0.2 * value), 190 ambient_light=[0.02, 0.02, 0.02],191 bg_color=[0.0, 0.0, 0.0, 1.0]))192 193 light = pyrender.DirectionalLight(194 color=np.ones(3), 195 intensity=light_intensity196 ) if light_intensity is not None else None197 camera = pyrender.PerspectiveCamera(198 yfov=np.deg2rad(fov),199 aspectRatio=image_size[0]/image_size[1],200 znear=znear,201 zfar=zfar202 )203 renderer = pyrender.OffscreenRenderer(*image_size)204 205 camera_poses = create_circular_camera_poses(206 num_views, 207 radius, 208 axis = axis209 )210 211 images, depths = [], []212 for i, pose in enumerate(camera_poses):213 image, depth = render(214 scenes[i], renderer, camera, pose, light, 215 normalize_depth=normalize_depth,216 flags=flags,217 return_type=return_type218 )219 images.append(image)220 depths.append(depth)221 222 renderer.delete()223 224 if return_depth:225 return images, depths226 return images227 228def render_normal_views_around_mesh(229 mesh: Union[trimesh.Trimesh, trimesh.Scene],230 num_views: int = 36,231 radius: float = 3.5,232 axis: np.ndarray = np.array([0.0, 1.0, 0.0]),233 image_size: tuple = (512, 512),234 fov: float = 40.0,235 light_intensity: Optional[float] = 5.0,236 znear: float = 0.1,237 zfar: float = 10.0,238 normalize_depth: bool = False,239 flags: int = pyrender.constants.RenderFlags.NONE,240 return_depth: bool = False, 241 return_type: Literal['pil', 'ndarray'] = 'pil'242) -> Union[243 List[Image.Image], 244 List[np.ndarray], 245 Tuple[List[Image.Image], List[Image.Image]], 246 Tuple[List[np.ndarray], List[np.ndarray]]247 ]:248 249 if not isinstance(mesh, (trimesh.Trimesh, trimesh.Scene)):250 raise ValueError("mesh must be a trimesh.Trimesh or trimesh.Scene object")251 if isinstance(mesh, trimesh.Scene):252 mesh = mesh.to_geometry()253 normals = mesh.vertex_normals254 colors = ((normals + 1.0) / 2.0 * 255).astype(np.uint8)255 mesh.visual = trimesh.visual.ColorVisuals(256 mesh=mesh,257 vertex_colors=colors258 )259 mesh = trimesh.Scene(mesh)260 return render_views_around_mesh(261 mesh, num_views, radius, axis, 262 image_size, fov, light_intensity, znear, zfar, 263 normalize_depth, flags,264 return_depth, return_type265 )266 267def create_camera_pose_on_sphere(268 azimuth: float = 0.0, # in degrees269 elevation: float = 0.0, # in degrees270 radius: float = 3.5,271) -> np.ndarray:272 # Create a camera pose for a given azimuth and elevation273 # with the given radius.274 # The camera always looks at the origin.275 # The up vector is always [0, 1, 0].276 canonical_pose = np.array([277 [1.0, 0.0, 0.0, 0.0],278 [0.0, 1.0, 0.0, 0.0],279 [0.0, 0.0, 1.0, radius],280 [0.0, 0.0, 0.0, 1.0]281 ])282 azimuth = np.deg2rad(azimuth)283 elevation = np.deg2rad(elevation)284 position = np.array([285 np.cos(elevation) * np.sin(azimuth),286 np.sin(elevation),287 np.cos(elevation) * np.cos(azimuth),288 ])289 R = np.eye(4)290 R[:3, :3] = rotation_matrix_from_vectors(291 np.array([0.0, 0.0, 1.0]), 292 position293 )294 pose = R @ canonical_pose295 return pose296 297def render_single_view(298 mesh: Union[trimesh.Trimesh, trimesh.Scene],299 azimuth: float = 0.0, # in degrees300 elevation: float = 0.0, # in degrees301 radius: float = 3.5,302 image_size: tuple = (512, 512),303 fov: float = 40.0,304 light_intensity: Optional[float] = 5.0,305 num_env_lights: int = 0, 306 znear: float = 0.1,307 zfar: float = 10.0,308 normalize_depth: bool = False,309 flags: int = pyrender.constants.RenderFlags.NONE,310 return_depth: bool = False, 311 return_type: Literal['pil', 'ndarray'] = 'pil'312) -> Union[313 Image.Image, 314 np.ndarray, 315 Tuple[Image.Image, Image.Image], 316 Tuple[np.ndarray, np.ndarray]317 ]:318 319 if not isinstance(mesh, (trimesh.Trimesh, trimesh.Scene)):320 raise ValueError("mesh must be a trimesh.Trimesh or trimesh.Scene object")321 if isinstance(mesh, trimesh.Trimesh):322 mesh = trimesh.Scene(mesh)323 324 scene = pyrender.Scene.from_trimesh_scene(mesh)325 light = pyrender.DirectionalLight(326 color=np.ones(3), 327 intensity=light_intensity328 ) if light_intensity is not None else None329 camera = pyrender.PerspectiveCamera(330 yfov=np.deg2rad(fov),331 aspectRatio=image_size[0]/image_size[1],332 znear=znear,333 zfar=zfar334 )335 renderer = pyrender.OffscreenRenderer(*image_size)336 337 camera_pose = create_camera_pose_on_sphere(338 azimuth,339 elevation,340 radius341 )342 343 if num_env_lights > 0:344 env_light_poses = create_circular_camera_poses(345 num_env_lights,346 radius,347 axis = np.array([0.0, 1.0, 0.0])348 )349 for pose in env_light_poses:350 scene.add(pyrender.DirectionalLight(351 color=np.ones(3),352 intensity=light_intensity353 ), pose=pose)354 # set light to None355 light = None356 357 image, depth = render(358 scene, renderer, camera, camera_pose, light,359 normalize_depth=normalize_depth,360 flags=flags,361 return_type=return_type362 )363 renderer.delete()364 365 if return_depth:366 return image, depth367 return image368 369def render_normal_single_view(370 mesh: Union[trimesh.Trimesh, trimesh.Scene],371 azimuth: float = 0.0, # in degrees372 elevation: float = 0.0, # in degrees373 radius: float = 3.5,374 image_size: tuple = (512, 512),375 fov: float = 40.0,376 light_intensity: Optional[float] = 5.0,377 znear: float = 0.1,378 zfar: float = 10.0,379 normalize_depth: bool = False,380 flags: int = pyrender.constants.RenderFlags.NONE,381 return_depth: bool = False,382 return_type: Literal['pil', 'ndarray'] = 'pil'383) -> Union[384 Image.Image,385 np.ndarray,386 Tuple[Image.Image, Image.Image],387 Tuple[np.ndarray, np.ndarray]388 ]:389 390 if not isinstance(mesh, (trimesh.Trimesh, trimesh.Scene)):391 raise ValueError("mesh must be a trimesh.Trimesh or trimesh.Scene object")392 if isinstance(mesh, trimesh.Scene):393 mesh = mesh.to_geometry()394 normals = mesh.vertex_normals395 colors = ((normals + 1.0) / 2.0 * 255).astype(np.uint8)396 mesh.visual = trimesh.visual.ColorVisuals(397 mesh=mesh,398 vertex_colors=colors399 )400 mesh = trimesh.Scene(mesh)401 return render_single_view(402 mesh, azimuth, elevation, radius, 403 image_size, fov, light_intensity, znear, zfar,404 normalize_depth, flags, 405 return_depth, return_type406 )407 408def export_renderings(409 images: List[Image.Image],410 export_path: str,411 fps: int = 36, 412 loop: int = 0413): 414 export_type = export_path.split('.')[-1]415 if export_type == 'mp4':416 export_to_video(417 images,418 export_path,419 fps=fps,420 )421 elif export_type == 'gif':422 duration = 1000 / fps423 images[0].save(424 export_path,425 save_all=True,426 append_images=images[1:],427 duration=duration,428 loop=loop429 )430 else:431 raise ValueError(f'Unknown export type: {export_type}')432 433def make_grid_for_images_or_videos(434 images_or_videos: Union[List[Image.Image], List[List[Image.Image]]],435 nrow: int = 4, 436 padding: int = 0, 437 pad_value: int = 0, 438 image_size: tuple = (512, 512),439 return_type: Literal['pil', 'ndarray'] = 'pil'440) -> Union[Image.Image, List[Image.Image], np.ndarray]:441 if isinstance(images_or_videos[0], Image.Image):442 images = [np.array(image.resize(image_size).convert('RGB')) for image in images_or_videos]443 images = np.stack(images, axis=0).transpose(0, 3, 1, 2) # [N, C, H, W]444 images = torch.from_numpy(images)445 image_grid = make_grid(446 images,447 nrow=nrow,448 padding=padding,449 pad_value=pad_value,450 normalize=False451 ) # [C, H', W']452 image_grid = image_grid.cpu().numpy()453 if return_type == 'pil':454 image_grid = Image.fromarray(image_grid.transpose(1, 2, 0))455 return image_grid456 elif isinstance(images_or_videos[0], list) and isinstance(images_or_videos[0][0], Image.Image):457 image_grids = []458 for i in range(len(images_or_videos[0])):459 images = [video[i] for video in images_or_videos]460 image_grid = make_grid_for_images_or_videos(461 images,462 nrow=nrow,463 padding=padding,464 return_type=return_type465 )466 image_grids.append(image_grid)467 if return_type == 'ndarray':468 image_grids = np.stack(image_grids, axis=0)469 return image_grids470 else:471 raise ValueError(f'Unknown input type: {type(images_or_videos[0])}')