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,
)
| 29 | |
| 30 | |
| 31 | def bake_texture( |
| 32 | vertices: np.array, |
| 33 | faces: np.array, |
| 34 | uvs: np.array, |
| 35 | observations: List[np.array], |
| 36 | masks: List[np.array], |
| 37 | extrinsics: List[np.array], |
| 38 | intrinsics: List[np.array], |
| 39 | texture_size: int = 2048, |
| 40 | near: float = 0.1, |
| 41 | far: float = 10.0, |
| 42 | mode: Literal['fast', 'opt'] = 'opt', |
| 43 | lambda_tv: float = 1e-2, |
| 44 | verbose: bool = False, |
| 45 | ): |
| 46 | """ |
| 47 | Bake texture to a mesh from multiple observations. |
| 48 | Args: |
| 49 | vertices (np.array): Vertices of the mesh. Shape (V, 3). |
| 50 | faces (np.array): Faces of the mesh. Shape (F, 3). |
| 51 | uvs (np.array): UV coordinates of the mesh. Shape (V, 2). |
| 52 | observations (List[np.array]): List of observations. Each observation is a 2D image. Shape (H, W, 3). |
| 53 | masks (List[np.array]): List of masks. Each mask is a 2D image. Shape (H, W). |
| 54 | extrinsics (List[np.array]): List of extrinsics. Shape (4, 4). |
| 55 | intrinsics (List[np.array]): List of intrinsics. Shape (3, 3). |
| 56 | texture_size (int): Size of the texture. |
| 57 | near (float): Near plane of the camera. |
| 58 | far (float): Far plane of the camera. |
| 59 | mode (Literal['fast', 'opt']): Mode of texture baking. |
| 60 | lambda_tv (float): Weight of total variation loss in optimization. |
| 61 | verbose (bool): Whether to print progress. |
| 62 | """ |
| 63 | vertices = torch.tensor(vertices).float().cuda() |
| 64 | faces = torch.tensor(faces.astype(np.int32)).cuda() |
| 65 | uvs = torch.tensor(uvs).float().cuda() |
| 66 | observations = [torch.tensor(obs).float().cuda() for obs in observations] |
| 67 | masks = [torch.tensor(m>1e-2).bool().cuda() for m in masks] |
| 68 | views = [utils3d.torch.extrinsics_to_view(torch.tensor(extr).float().cuda()) for extr in extrinsics] |
| 69 | projections = [utils3d.torch.intrinsics_to_perspective(torch.tensor(intr).float().cuda(), near, far) for intr in intrinsics] |
| 70 | |
| 71 | if mode == 'fast': |
| 72 | texture = torch.zeros((texture_size * texture_size, 3), dtype=torch.float32).cuda() |
| 73 | texture_weights = torch.zeros((texture_size * texture_size), dtype=torch.float32).cuda() |
| 74 | rastctx = utils3d.torch.RastContext(backend='cuda') |
| 75 | for observation, view, projection in tqdm(zip(observations, views, projections), total=len(observations), disable=not verbose, desc='Texture baking (fast)'): |
| 76 | with torch.no_grad(): |
| 77 | rast = utils3d.torch.rasterize_triangle_faces( |
| 78 | rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection |
| 79 | ) |
| 80 | uv_map = rast['uv'][0].detach().flip(0) |
| 81 | mask = rast['mask'][0].detach().bool() & masks[0] |
| 82 | |
| 83 | # nearest neighbor interpolation |
| 84 | uv_map = (uv_map * texture_size).floor().long() |
| 85 | obs = observation[mask] |
| 86 | uv_map = uv_map[mask] |
| 87 | idx = uv_map[:, 0] + (texture_size - uv_map[:, 1] - 1) * texture_size |
| 88 | texture = texture.scatter_add(0, idx.view(-1, 1).expand(-1, 3), obs) |
no test coverage detected