Bake texture to a mesh from multiple observations. Args: vertices (np.array): Vertices of the mesh. Shape (V, 3). faces (np.array): Faces of the mesh. Shape (F, 3). uvs (np.array): UV coordinates of the mesh. Shape (V, 2). observations (List[np.array]): List
(
vertices: np.array,
faces: np.array,
uvs: np.array,
observations: List[np.array],
masks: List[np.array],
extrinsics: List[np.array],
intrinsics: List[np.array],
texture_size: int = 2048,
near: float = 0.1,
far: float = 10.0,
mode: Literal['fast', 'opt'] = 'opt',
lambda_tv: float = 1e-2,
verbose: bool = False,
)
| 274 | |
| 275 | |
| 276 | def bake_texture( |
| 277 | vertices: np.array, |
| 278 | faces: np.array, |
| 279 | uvs: np.array, |
| 280 | observations: List[np.array], |
| 281 | masks: List[np.array], |
| 282 | extrinsics: List[np.array], |
| 283 | intrinsics: List[np.array], |
| 284 | texture_size: int = 2048, |
| 285 | near: float = 0.1, |
| 286 | far: float = 10.0, |
| 287 | mode: Literal['fast', 'opt'] = 'opt', |
| 288 | lambda_tv: float = 1e-2, |
| 289 | verbose: bool = False, |
| 290 | ): |
| 291 | """ |
| 292 | Bake texture to a mesh from multiple observations. |
| 293 | |
| 294 | Args: |
| 295 | vertices (np.array): Vertices of the mesh. Shape (V, 3). |
| 296 | faces (np.array): Faces of the mesh. Shape (F, 3). |
| 297 | uvs (np.array): UV coordinates of the mesh. Shape (V, 2). |
| 298 | observations (List[np.array]): List of observations. Each observation is a 2D image. Shape (H, W, 3). |
| 299 | masks (List[np.array]): List of masks. Each mask is a 2D image. Shape (H, W). |
| 300 | extrinsics (List[np.array]): List of extrinsics. Shape (4, 4). |
| 301 | intrinsics (List[np.array]): List of intrinsics. Shape (3, 3). |
| 302 | texture_size (int): Size of the texture. |
| 303 | near (float): Near plane of the camera. |
| 304 | far (float): Far plane of the camera. |
| 305 | mode (Literal['fast', 'opt']): Mode of texture baking. |
| 306 | lambda_tv (float): Weight of total variation loss in optimization. |
| 307 | verbose (bool): Whether to print progress. |
| 308 | """ |
| 309 | vertices = torch.tensor(vertices).cuda() |
| 310 | faces = torch.tensor(faces.astype(np.int32)).cuda() |
| 311 | uvs = torch.tensor(uvs).cuda() |
| 312 | observations = [torch.tensor(obs / 255.0).float().cuda() for obs in observations] |
| 313 | masks = [torch.tensor(m>0).bool().cuda() for m in masks] |
| 314 | views = [utils3d.torch.extrinsics_to_view(torch.tensor(extr).cuda()) for extr in extrinsics] |
| 315 | projections = [utils3d.torch.intrinsics_to_perspective(torch.tensor(intr).cuda(), near, far) for intr in intrinsics] |
| 316 | |
| 317 | if mode == 'fast': |
| 318 | texture = torch.zeros((texture_size * texture_size, 3), dtype=torch.float32).cuda() |
| 319 | texture_weights = torch.zeros((texture_size * texture_size), dtype=torch.float32).cuda() |
| 320 | rastctx = utils3d.torch.RastContext(backend='cuda') |
| 321 | for observation, view, projection in tqdm(zip(observations, views, projections), total=len(observations), disable=not verbose, desc='Texture baking (fast)'): |
| 322 | with torch.no_grad(): |
| 323 | rast = utils3d.torch.rasterize_triangle_faces( |
| 324 | rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection |
| 325 | ) |
| 326 | uv_map = rast['uv'][0].detach().flip(0) |
| 327 | mask = rast['mask'][0].detach().bool() & masks[0] |
| 328 | |
| 329 | # nearest neighbor interpolation |
| 330 | uv_map = (uv_map * texture_size).floor().long() |
| 331 | obs = observation[mask] |
| 332 | uv_map = uv_map[mask] |
| 333 | idx = uv_map[:, 0] + (texture_size - uv_map[:, 1] - 1) * texture_size |