Spaces:
Running on Zero
Running on Zero
| import random | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from easydict import EasyDict as edict | |
| from kornia.geometry.transform import pyrdown | |
| from rich import print | |
| import itertools | |
| import torch.nn.functional as F | |
| from pdb import set_trace as stx | |
| # 将输入的 data_batch 拆分成 input 和 target | |
| class SplitData(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| def forward(self, data_batch, target_has_input=True): | |
| """ | |
| image (torch.tensor): [b, v, c, h, w] | |
| fxfycxcy (torch.tensor): [b, v, 4] | |
| c2w (torch.tensor): [b, v, 4, 4] | |
| label (torch.tensor): [b, v, c, h, w] | |
| input: | |
| data_batch: {"image": (b, v, c, h, w)} - a dict | |
| 举个例子: [8, 8, 3, 256, 256] | |
| return: | |
| input: {"image": (b, v, c, h, w)} - a dict | |
| target: {"image": (b, v, c, h, w)} - a dict | |
| """ | |
| input, target = {}, {} | |
| b,v,c,h,w = data_batch['image'].size() | |
| index = None | |
| for key, value in data_batch.items(): | |
| try: | |
| input[key] = value[:, : self.config.training.num_input_views, ...] # 取前 num_input_views 个视图, 在应用中 num_input_views = 4 | |
| except Exception as exc: | |
| raise ValueError(f"Batch field {key!r} must have batch and view dimensions") from exc | |
| if self.config.training.num_target_views >= value.size(1): # 如果 target 的数量大于等于输入的数量, 4 < 8, 这一步不执行 | |
| target[key] = value | |
| else: | |
| if index is None: # 只在 index 尚未被分配时执行 | |
| b, v = value.shape[:2] | |
| ''' | |
| 分两种情况： | |
| [1] target 里面包含 input | |
| [2] target 里面不包含 input | |
| ''' | |
| if target_has_input: | |
| index = np.array( # 生成一个随机的 index，用于选择 target, input 和 target 可能有重叠 | |
| [ | |
| random.sample( | |
| range(v), self.config.training.num_target_views | |
| ) | |
| for _ in range(b) | |
| ] | |
| ) | |
| else: | |
| assert ( | |
| self.config.training.num_input_views | |
| + self.config.training.num_target_views | |
| <= self.config.training.num_views | |
| ), "num_input_views + num_target_views must <= num_views to avoid duplicate views" | |
| index = np.array( | |
| [ | |
| [ | |
| self.config.training.num_views - 1 - j | |
| for j in range( | |
| self.config.training.num_target_views | |
| ) | |
| ] | |
| for _ in range(b) | |
| ] | |
| ) | |
| index = torch.from_numpy(index).long().to(value.device) # 把随机取好的序号放到 GPU 上 | |
| # match shape of index to v, insert dummy dimensions to match v | |
| value_index = index | |
| if value.dim() > 2: # 此时包含了维度 v | |
| dummy_dims = [ | |
| 1, | |
| ] * (value.dim() - 2) | |
| value_index = index.reshape( | |
| index.size(0), index.size(1), *dummy_dims | |
| ) | |
| try: | |
| ''' | |
| torch.gather - 按照指定的张量在特定的维度按照索引来取元素 | |
| 此处是根据 index 来取 target 的元素 | |
| 在 label wild 测试中, not eval, value: torch.Size([16, 8, 1, 256, 256]), index: torch.Size([16, 6, 1, 1, 1]) | |
| 在 label wild 测试中, eval, value: torch.Size([16, 8, 1, 256, 256]), index: torch.Size([16, 6, 1, 1, 1]) | |
| ''' | |
| # if key == 'label': | |
| # stx() | |
| target[key] = torch.gather( | |
| value, | |
| dim=1, | |
| index=value_index.expand(-1, -1, *value.size()[2:]), | |
| ) | |
| except Exception as e: | |
| print(f"key: {key}") | |
| print(f"value: {value.size()}") | |
| print(f"index: {value_index.size()}") | |
| raise e | |
| return edict(input), edict(target) | |
| ''' | |
| [1] 从输入的数据批次中提取图像、内参、外参和索引。 | |
| [2] 对图像、内参和外参的维度进行检查，确保它们的维度符合预期。 | |
| [3] 对图像、内参和外参进行重塑，以便于后续的处理。 | |
| [4] 生成图像的网格坐标，并将这些坐标转移到图像的设备上。 | |
| [5] 对网格坐标进行归一化处理，使其范围在 [-1, 1] 之间。 | |
| [6] 对归一化后的坐标进行扩展，以便于后续的处理。 | |
| [7] 计算射线的方向(ray_d)和原点(ray_o)。 | |
| [8] 如果提供了 patch_size,则会对图像进行分块处理,并计算每个块的颜色、射线的方向和原点。 | |
| [9] 最后，将处理后的数据打包并返回。 | |
| ''' | |
| class TransformInput(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| def forward(self, data_batch, patch_size=None): | |
| """transform input image before feeding into transformer | |
| Args: | |
| image (torch.tensor): [b, v, c, h, w] # 图像 | |
| fxfycxcy (torch.tensor): [b, v, 4] # 内参 | |
| c2w (torch.tensor): [b, v, 4, 4] # 外参 | |
| Returns: | |
| image: [b, v, c, h, w] | |
| ray_o: [b, v, 3, h, w] | |
| ray_d: [b, v, 3, h, w] | |
| Input 这个字典所包含的 key 有: (['image', 'ray_o', 'ray_d', 'ray_d_cam', 'fxfycxcy', 'c2w', 'index', 'xy_norm', 'ray_color_patch', 'ray_o_patch', 'ray_d_patch', 'ray_xy_norm_patch', 'proj_mat']) | |
| """ | |
| image, fxfycxcy, c2w, image_noisy = ( | |
| data_batch.image, | |
| data_batch.fxfycxcy, | |
| data_batch.c2w, | |
| # data_batch.index, | |
| data_batch.image_noisy, | |
| ) | |
| # stx() | |
| #image[:,1:4,...] = image_noisy # [8, 4, 3, 256, 256] | |
| # stx() | |
| assert image.dim() == 5, f"image dim should be 5, but got {image.dim()}" # b, v, c, h, w | |
| assert ( | |
| fxfycxcy.dim() == 3 | |
| ), f"fxfycxcy dim should be 3, but got {fxfycxcy.dim()}" # b, v, 4 | |
| assert c2w.dim() == 4, f"c2w dim should be 4, but got {c2w.dim()}" # b, v, 4, 4 | |
| b, v, c, h, w = image.size() | |
| image = image.reshape(b * v, c, h * w) | |
| fxfycxcy = fxfycxcy.reshape(b * v, 4) | |
| c2w = c2w.reshape(b * v, 4, 4) # 进来就把 b 和 v 合并 | |
| # y - i - 行 - h, x - j - 列 - w | |
| # y 和 x 都是形状为 [h, w] 的矩阵, y 的每一行都是 0 到 w-1, x 的每一列都是 0 到 h-1 | |
| y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") | |
| # print(f"y device: {y.device}, x device: {x.device}, image device: {image.device}") | |
| # print(f"y dtype: {y.dtype}, x dtype: {x.dtype}, image dtype: {image.dtype}") | |
| # print(f"y shape: {y.shape}, x shape: {x.shape}, image shape: {image.shape}") | |
| # stx() | |
| ''' | |
| 1. not eval: | |
| y device: cpu, x device: cpu, image device: cuda:0 | |
| y dtype: torch.int64, x dtype: torch.int64, image dtype: torch.float32 | |
| y shape: torch.Size([256, 256]), x shape: torch.Size([256, 256]), image shape: torch.Size([64, 3, 65536]) | |
| 2. eval: | |
| y device: cpu, x device: cpu, image device: cuda:0 | |
| y dtype: torch.int64, x dtype: torch.int64, image dtype: torch.float32 | |
| y shape: torch.Size([256, 256]), x shape: torch.Size([256, 256]), image shape: torch.Size([64, 3, 65536]) | |
| ''' | |
| y, x = y.to(image.device), x.to(image.device) | |
| # [h, w], normalize 到 [-1, 1] | |
| y_norm, x_norm = (y + 0.5) / h * 2 - 1, (x + 0.5) / w * 2 - 1 | |
| # [b, v, 2, h, w] | |
| xy_norm = torch.stack([x_norm, y_norm], dim=0)[None, None, :, :, :].expand( | |
| b, v, -1, -1, -1 | |
| ) | |
| x = x[None, :, :].expand(b * v, -1, -1).reshape(b * v, -1) | |
| y = y[None, :, :].expand(b * v, -1, -1).reshape(b * v, -1) | |
| # 到此处, x, y 是图像坐标 (u, v), 它与相机坐标 (x_c, y_c, z_c) 的转换关系是: | |
| # u = f_x * x_c / z_c + cx | |
| # v = f_y * y_c / z_c + cy | |
| # 所以反算 x_c, y_c 是 | |
| # x_c / z_c = (u - cx) / f_x, y_c / z_c = (v - cy) / f_y | |
| # 那么 per-pixel 的射线方向就是 (x_c, y_c, z_c) = (x_c / z_c, y_c / z_c, 1) | |
| x = (x + 0.5 - fxfycxcy[:, 2:3]) / fxfycxcy[:, 0:1] | |
| y = (y + 0.5 - fxfycxcy[:, 3:4]) / fxfycxcy[:, 1:2] | |
| z = torch.ones_like(x) | |
| # ray_d 是相机坐标系下的射线方向 | |
| ray_d = torch.stack([x, y, z], dim=2) # [b*v, h*w, 3] | |
| ray_d_cam = ray_d.clone() | |
| # bmm = batch matrix multiplication, 这里转成世界坐标系 | |
| ray_d = torch.bmm(ray_d, c2w[:, :3, :3].transpose(1, 2)) # [b*v, h*w, 3] | |
| # norm | |
| ray_d = ray_d / torch.norm(ray_d, dim=2, keepdim=True) # [b*v, h*w, 3] | |
| ray_d_cam = ray_d_cam / torch.norm(ray_d_cam, dim=2, keepdim=True) | |
| ray_o = c2w[:, :3, 3][:, None, :].expand_as(ray_d) # [b*v, h*w, 3] | |
| # fetch color, ray_o, ray_d for all patch centers | |
| ray_color_patch, ray_o_patch, ray_d_patch = None, None, None | |
| ray_xy_norm_patch, proj_mat = None, None | |
| if patch_size is not None: | |
| start_patch_center = patch_size / 2.0 | |
| y, x = torch.meshgrid( | |
| torch.arange(h // patch_size) * patch_size + start_patch_center, | |
| torch.arange(w // patch_size) * patch_size + start_patch_center, | |
| indexing="ij", | |
| ) | |
| y, x = y.to(image.device), x.to(image.device) | |
| x = x[None, :, :].expand(b * v, -1, -1).reshape(b * v, -1) | |
| y = y[None, :, :].expand(b * v, -1, -1).reshape(b * v, -1) | |
| ray_xy_norm_patch = torch.stack( | |
| [x / w, y / h], dim=2 | |
| ) # use [0,1] for patch center | |
| K_norm = ( | |
| torch.eye(3, device=image.device).unsqueeze(0).repeat(b * v, 1, 1) | |
| ) # [b*v, 3, 3] | |
| K_norm[:, 0, 0] = fxfycxcy[:, 0] / w | |
| K_norm[:, 1, 1] = fxfycxcy[:, 1] / h | |
| K_norm[:, 0, 2] = fxfycxcy[:, 2] / w | |
| K_norm[:, 1, 2] = fxfycxcy[:, 3] / h | |
| w2c = torch.inverse(c2w) # [b*v, 4, 4] | |
| proj_mat = torch.bmm(K_norm, w2c[:, :3, :4]) # [b*v, 3, 4] | |
| proj_mat = proj_mat.reshape(b * v, 12) | |
| proj_mat = proj_mat / (proj_mat.norm(dim=1, keepdim=True) + 1e-6) | |
| proj_mat = proj_mat.reshape(b * v, 3, 4) | |
| proj_mat = proj_mat * proj_mat[:, 0:1, 0:1].sign() | |
| # fetch colors for all patch centers | |
| ray_color_patch = ( | |
| F.grid_sample( | |
| image.reshape(b * v, c, h, w), | |
| torch.stack([x / w * 2.0 - 1.0, y / h * 2.0 - 1.0], dim=2).reshape( | |
| b * v, -1, 1, 2 | |
| ), | |
| align_corners=False, | |
| ) | |
| .squeeze(-1) | |
| .permute(0, 2, 1) | |
| ).contiguous() # [b*v, h'*w', c] | |
| x = (x - fxfycxcy[:, 2:3]) / fxfycxcy[:, 0:1] | |
| y = (y - fxfycxcy[:, 3:4]) / fxfycxcy[:, 1:2] | |
| z = torch.ones_like(x) | |
| ray_d_patch = torch.stack([x, y, z], dim=2) # [b*v, h'*w', 3] | |
| ray_d_patch = torch.bmm( | |
| ray_d_patch, c2w[:, :3, :3].transpose(1, 2) | |
| ) # [b*v, h'*w', 3] | |
| ray_d_patch = ray_d_patch / torch.norm( | |
| ray_d_patch, dim=2, keepdim=True | |
| ) # [b*v, h'*w', 3] | |
| ray_o_patch = c2w[:, :3, 3][:, None, :].expand_as( | |
| ray_d_patch | |
| ) # [b*v, h'*w', 3] | |
| n_patch = ray_color_patch.size(1) | |
| ray_color_patch = ray_color_patch.reshape(b, v, n_patch, c) | |
| ray_o_patch = ray_o_patch.reshape(b, v, n_patch, 3) | |
| ray_d_patch = ray_d_patch.reshape(b, v, n_patch, 3) | |
| ray_xy_norm_patch = ray_xy_norm_patch.reshape(b, v, n_patch, 2) | |
| proj_mat = proj_mat.reshape(b, v, 3, 4) | |
| ray_o = ray_o.reshape(b, v, h, w, 3).permute(0, 1, 4, 2, 3) | |
| ray_d = ray_d.reshape(b, v, h, w, 3).permute(0, 1, 4, 2, 3) | |
| ray_d_cam = ray_d_cam.reshape(b, v, h, w, 3).permute(0, 1, 4, 2, 3) | |
| image = image.reshape(b, v, c, h, w) | |
| fxfycxcy = fxfycxcy.reshape(b, v, 4) | |
| c2w = c2w.reshape(b, v, 4, 4) | |
| prepared_input = edict( | |
| image=image, | |
| ray_o=ray_o, | |
| ray_d=ray_d, | |
| ray_d_cam=ray_d_cam, | |
| fxfycxcy=fxfycxcy, | |
| c2w=c2w, | |
| ) | |
| return prepared_input | |
| ''' | |
| 为每个 batch 的输入数据生成网格 | |
| x 坐标和 y 坐标分别存在 x1_list 和 x2_list 中 | |
| ''' | |
| def batched_meshgrid(x1, x2, indexing="ij"): | |
| """_summary_ | |
| Args: | |
| x1 (torch.tensor): [b, n1] | |
| x2 (torch.tensor): [b, n2] | |
| Returns: | |
| x1_grid (torch.tensor): [b, n1, n2] | |
| x2_grid (torch.tensor): [b, n1, n2] | |
| """ | |
| x1_list, x2_list = [], [] | |
| for i in range(x1.size(0)): | |
| x1_i, x2_i = torch.meshgrid(x1[i], x2[i], indexing=indexing) | |
| x1_list.append(x1_i) | |
| x2_list.append(x2_i) | |
| return torch.stack(x1_list, dim=0), torch.stack(x2_list, dim=0) | |
| class TransformTarget(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| def forward(self, data_batch): | |
| """Transform target before computing loss | |
| Args: | |
| image (torch.tensor): [b, v, c, h, w] | |
| fxfycxcy (torch.tensor): [b, v, 4] | |
| c2w (torch.tensor): [b, v, 4, 4] | |
| Returns: | |
| target: {"image": (b, v*(crop_size*crop_size+n_rand_rays), 3), | |
| "ray_o": (b, v*(crop_size*crop_size+n_rand_rays), 3), | |
| "ray_d": (b, v*(crop_size*crop_size+n_rand_rays), 3), | |
| "crop_size": crop_size, | |
| "n_rand_rays": n_rand_rays} | |
| """ | |
| image, depth, normal, fxfycxcy, c2w = ( | |
| data_batch["image"], | |
| data_batch["depth"], | |
| data_batch["normal"], | |
| data_batch["fxfycxcy"], | |
| data_batch["c2w"], | |
| ) | |
| crop_size, n_rand_rays, downsample = ( | |
| self.config.training.crop_size, | |
| self.config.training.n_rand_rays, | |
| self.config.training.downsample, | |
| ) | |
| # downsample 表示降采样的 scale factor | |
| if downsample > 1: | |
| with torch.cuda.amp.autocast(enabled=False): # autocast 决定是否启用混合精度 | |
| b, v, c, h, w = image.size() | |
| dtype = image.dtype | |
| image = pyrdown( | |
| image.reshape(-1, c, h, w).to(dtype=torch.float32), | |
| factor=downsample, | |
| ).to(dtype=dtype) | |
| _, c, h, w = image.size() | |
| image = image.reshape(b, v, c, h, w) | |
| depth = F.interpolate( | |
| depth.reshape(-1, 1, h, w).to(dtype=torch.float32), | |
| size=(h, w), | |
| interpolation="nearest", | |
| ) | |
| depth = depth.reshape(b, v, 1, h, w) | |
| normal = F.interpolate( | |
| normal.reshape(-1, 3, h, w).to(dtype=torch.float32), | |
| size=(h, w), | |
| interpolation="nearest", | |
| ) | |
| normal = normal.reshape(b, v, 3, h, w) | |
| fxfycxcy = fxfycxcy / downsample | |
| b, v, c, h, w = image.size() | |
| image = image.reshape(b * v, c, h * w) | |
| depth = depth.reshape(b * v, 1, h * w) | |
| normal = normal.reshape(b * v, 3, h * w) | |
| # sample crop supervision | |
| # [b*v, crop_size] | |
| is_crop_performed = (w > crop_size) or (h > crop_size) | |
| if w > crop_size: | |
| crop_idx_x = ( | |
| torch.randint(low=0, high=w - crop_size, size=(b * v, 1)) | |
| + torch.arange(crop_size)[None, :] | |
| ).long() | |
| else: | |
| crop_idx_x = torch.arange(w)[None, :].expand(b * v, -1).long() | |
| if h > crop_size: | |
| crop_idx_y = ( | |
| torch.randint(low=0, high=h - crop_size, size=(b * v, 1)) | |
| + torch.arange(crop_size)[None, :] | |
| ).long() | |
| else: | |
| crop_idx_y = torch.arange(h)[None, :].expand(b * v, -1).long() | |
| # modify fxfycxcy | |
| # 并不影响 focal，只是影响 cx 和 cy | |
| crop_fxfycxcy = fxfycxcy.clone() # [b, v, 4] | |
| if is_crop_performed: | |
| device = fxfycxcy.device | |
| crop_fxfycxcy[:, :, 2] = crop_fxfycxcy[:, :, 2] - crop_idx_x[:, 0].reshape( | |
| b, v | |
| ).to( | |
| device | |
| ) # cx | |
| crop_fxfycxcy[:, :, 3] = crop_fxfycxcy[:, :, 3] - crop_idx_y[:, 0].reshape( | |
| b, v | |
| ).to( | |
| device | |
| ) # cy | |
| # [b*v, crop_size, crop_size] | |
| # not supported yet | |
| # crop_idx_y, crop_idx_x = torch.meshgrid(crop_idx_y, crop_idx_x, indexing="ij") | |
| # batched meshgrid | |
| crop_idx_y, crop_idx_x = batched_meshgrid(crop_idx_y, crop_idx_x, indexing="ij") | |
| # [b*v, crop_size*crop_size] | |
| crop_idx_x, crop_idx_y = crop_idx_x.reshape(b * v, -1), crop_idx_y.reshape( | |
| b * v, -1 | |
| ) | |
| sample_idx_x, sample_idx_y = crop_idx_x, crop_idx_y | |
| if n_rand_rays > 0: | |
| # sample random supervision | |
| # [b*v, n_rand_rays] | |
| rand_idx = torch.randint( | |
| low=0, high=h * w, size=(b * v, n_rand_rays) | |
| ).long() | |
| rand_idx_y, rand_idx_x = rand_idx // w, rand_idx % w | |
| # combine crop and random supervision | |
| # [b*v, crop_size*crop_size+n_rand_rays] | |
| sample_idx_x, sample_idx_y = torch.cat( | |
| [sample_idx_x, rand_idx_x], dim=1 | |
| ), torch.cat([sample_idx_y, rand_idx_y], dim=1) | |
| sample_idx_x, sample_idx_y = sample_idx_x.to(image.device), sample_idx_y.to( | |
| image.device | |
| ) | |
| sample_idx = sample_idx_y * w + sample_idx_x | |
| # shortcut if no crop and random supervision | |
| if (not is_crop_performed) and (n_rand_rays == 0): | |
| ray_color = image # [b*v, c, h*w] | |
| ray_depth = depth # [b*v, 1, h*w] | |
| ray_normal = normal # [b*v, 3, h*w] | |
| else: | |
| ray_color = torch.gather( | |
| image, dim=2, index=sample_idx[:, None, :].expand(-1, c, -1) | |
| ) # [b*v, c, crop_size*crop_size+n_rand_rays] | |
| ray_depth = torch.gather( | |
| depth, dim=2, index=sample_idx[:, None, :].expand(-1, 1, -1) | |
| ) # [b*v, 1, crop_size*crop_size+n_rand_rays] | |
| ray_normal = torch.gather( | |
| normal, dim=2, index=sample_idx[:, None, :].expand(-1, 3, -1) | |
| ) # [b*v, 3, crop_size*crop_size+n_rand_rays] | |
| fxfycxcy = fxfycxcy.reshape(b * v, 4) | |
| c2w = c2w.reshape(b * v, 4, 4) | |
| x = (sample_idx_x + 0.5 - fxfycxcy[:, 2:3]) / fxfycxcy[:, 0:1] | |
| y = (sample_idx_y + 0.5 - fxfycxcy[:, 3:4]) / fxfycxcy[:, 1:2] | |
| z = torch.ones_like(x) | |
| ray_d = torch.stack( | |
| [x, y, z], dim=2 | |
| ) # [b*v, crop_size*crop_size+n_rand_rays, 3] | |
| # 一个 batch 的矩阵乘法，bmm = batch matrix multiplication, 全部转成世界坐标系 | |
| ray_d = torch.bmm(ray_d, c2w[:, :3, :3].transpose(1, 2)) | |
| ray_d = ray_d / torch.norm(ray_d, dim=2, keepdim=True) | |
| ray_o = c2w[:, :3, 3][:, None, :].expand_as(ray_d) | |
| ray_color = ( | |
| ray_color.reshape(b, v, c, -1).permute(0, 1, 3, 2).reshape(b, -1, c) | |
| ) # [b, v*(crop_size*crop_size+n_rand_rays), c] | |
| ray_depth = ( | |
| ray_depth.reshape(b, v, 1, -1).permute(0, 1, 3, 2).reshape(b, -1, 1) | |
| ) # [b, v*(crop_size*crop_size+n_rand_rays), 1] | |
| ray_normal = ( | |
| ray_normal.reshape(b, v, 3, -1).permute(0, 1, 3, 2).reshape(b, -1, 3) | |
| ) # [b, v*(crop_size*crop_size+n_rand_rays), 3] | |
| ray_o = ray_o.reshape(b, -1, 3) # [b, v*(crop_size*crop_size+n_rand_rays), 3] | |
| ray_d = ray_d.reshape(b, -1, 3) # [b, v*(crop_size*crop_size+n_rand_rays), 3] | |
| crop_fxfycxcy = crop_fxfycxcy.reshape(b, -1, 4) # [b, v, 4] | |
| c2w = c2w.reshape(b, -1, 4, 4) # [b, v, 4, 4] | |
| target = edict( | |
| ray_color=ray_color.contiguous(), | |
| ray_o=ray_o.contiguous(), | |
| ray_d=ray_d.contiguous(), | |
| batch=b, | |
| view=v, | |
| channel=c, | |
| crop_size=crop_size, | |
| n_rand_rays=n_rand_rays, | |
| crop_fxfycxcy=crop_fxfycxcy, | |
| c2w=c2w, | |
| ray_depth=ray_depth.contiguous(), | |
| ray_normal=ray_normal.contiguous(), | |
| ) | |
| # # debug | |
| # view_ray_o = ray_o.reshape(b, v, -1, 3)[0, :, :, :] | |
| # view_ray_d = ray_d.reshape(b, v, -1, 3)[0, :, :, :] | |
| # from .utils_nerf import sample_pts_uniform | |
| # is_intersect, xyz, t_vals, dists = sample_pts_uniform( | |
| # view_ray_o, view_ray_d, num_samples_per_ray=16, randomize_samples=False | |
| # ) | |
| # import trimesh | |
| # import os | |
| # print(f"debug b,v: {b},{v}") | |
| # print(f"{view_ray_o[:, 0, :]}") | |
| # os.makedirs("debug", exist_ok=True) | |
| # trimesh.PointCloud(view_ray_o.reshape(-1, 3).cpu().numpy()).export( | |
| # f"debug/ray_o.ply" | |
| # ) | |
| # trimesh.PointCloud( | |
| # (view_ray_o + 0.2 * view_ray_d).reshape(-1, 3).cpu().numpy() | |
| # ).export(f"debug/ray_d.ply") | |
| # for i in range(xyz.shape[0]): | |
| # trimesh.PointCloud(xyz[i].reshape(-1, 3).cpu().numpy()).export( | |
| # f"debug/xyz_{i}.ply" | |
| # ) | |
| # trimesh.PointCloud( | |
| # (view_ray_o + 0.2 * view_ray_d)[i].reshape(-1, 3).cpu().numpy() | |
| # ).export(f"debug/ray_d_{i}.ply") | |
| # exit(0) | |
| return target | |
| def generate_drop_mask(b, v): | |
| drop_mask_list = [] | |
| for _ in range(b): | |
| num_to_drop = random.choice(list(range(v + 1))) | |
| combinations = list(itertools.combinations(list(range(v)), num_to_drop)) | |
| items_to_drop = random.choice(combinations) | |
| drop_mask = torch.ones(v, dtype=torch.bool) | |
| for item in items_to_drop: | |
| drop_mask[item] = False | |
| drop_mask_list.append(drop_mask) | |
| drop_mask_list = torch.stack(drop_mask_list, dim=0) | |
| return drop_mask_list | |
| class TransformMeshTarget(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| def forward(self, data_batch, generate_ray=False): | |
| """transform input image before feeding into transformer | |
| Args: | |
| image (torch.tensor): [b, v, c, h, w] | |
| fxfycxcy (torch.tensor): [b, v, 4] | |
| c2w (torch.tensor): [b, v, 4, 4] | |
| mv (torch.tensor): [b, v, 4, 4] | |
| mvp (torch.tensor): [b, v, 4, 4] | |
| Returns: | |
| image: [b, v, c, h, w] | |
| """ | |
| image, fxfycxcy, c2w, mv, mvp, depth, normal, index = ( | |
| data_batch.image, | |
| data_batch.fxfycxcy, | |
| data_batch.c2w, | |
| data_batch.mv, | |
| data_batch.mvp, | |
| data_batch.depth, | |
| data_batch.normal, | |
| data_batch.index, | |
| ) | |
| assert image.dim() == 5, f"image dim should be 5, but got {image.dim()}" | |
| assert ( | |
| fxfycxcy.dim() == 3 | |
| ), f"fxfycxcy dim should be 3, but got {fxfycxcy.dim()}" | |
| assert c2w.dim() == 4, f"c2w dim should be 4, but got {c2w.dim()}" | |
| assert mv.dim() == 4, f"mv dim should be 4, but got {mv.dim()}" | |
| assert mvp.dim() == 4, f"mvp dim should be 4, but got {mvp.dim()}" | |
| assert depth.dim() == 5, f"depth dim should be 5, but got {depth.dim()}" | |
| assert normal.dim() == 5, f"normal dim should be 5, but got {normal.dim()}" | |
| crop_size, n_rand_rays, downsample = image.size(3), 0, 1 | |
| b, v, c, h, w = image.size() | |
| image = image.reshape(b * v, c, h * w) | |
| # sample crop supervision | |
| # [b*v, crop_size] | |
| if w > crop_size: | |
| crop_idx_x = ( | |
| torch.randint(low=0, high=w - crop_size, size=(b * v, 1)) | |
| + torch.arange(crop_size)[None, :] | |
| ).long() | |
| else: | |
| crop_idx_x = torch.arange(w)[None, :].expand(b * v, -1).long() | |
| if h > crop_size: | |
| crop_idx_y = ( | |
| torch.randint(low=0, high=h - crop_size, size=(b * v, 1)) | |
| + torch.arange(crop_size)[None, :] | |
| ).long() | |
| else: | |
| crop_idx_y = torch.arange(h)[None, :].expand(b * v, -1).long() | |
| # [b*v, crop_size, crop_size] | |
| # not supported yet | |
| # crop_idx_y, crop_idx_x = torch.meshgrid(crop_idx_y, crop_idx_x, indexing="ij") | |
| # batched meshgrid | |
| crop_idx_y, crop_idx_x = batched_meshgrid(crop_idx_y, crop_idx_x, indexing="ij") | |
| # [b*v, crop_size*crop_size] | |
| crop_idx_x, crop_idx_y = crop_idx_x.reshape(b * v, -1), crop_idx_y.reshape( | |
| b * v, -1 | |
| ) | |
| # sample random supervision | |
| # [b*v, n_rand_rays] | |
| rand_idx = torch.randint(low=0, high=h * w, size=(b * v, n_rand_rays)).long() | |
| rand_idx_y, rand_idx_x = rand_idx // w, rand_idx % w | |
| # combine crop and random supervision | |
| # [b*v, crop_size*crop_size+n_rand_rays] | |
| sample_idx_x, sample_idx_y = torch.cat( | |
| [crop_idx_x, rand_idx_x], dim=1 | |
| ), torch.cat([crop_idx_y, rand_idx_y], dim=1) | |
| sample_idx_x, sample_idx_y = sample_idx_x.to(image.device), sample_idx_y.to( | |
| image.device | |
| ) | |
| sample_idx = sample_idx_y * w + sample_idx_x | |
| ray_color = torch.gather( | |
| image, dim=2, index=sample_idx[:, None, :].expand(-1, c, -1) | |
| ) # [b*v, c, crop_size*crop_size+n_rand_rays] | |
| fxfycxcy = fxfycxcy.reshape(b * v, 4) | |
| c2w = c2w.reshape(b * v, 4, 4) | |
| x = (sample_idx_x + 0.5 - fxfycxcy[:, 2:3]) / fxfycxcy[:, 0:1] | |
| y = (sample_idx_y + 0.5 - fxfycxcy[:, 3:4]) / fxfycxcy[:, 1:2] | |
| z = torch.ones_like(x) | |
| ray_d = torch.stack( | |
| [x, y, z], dim=2 | |
| ) # [b*v, crop_size*crop_size+n_rand_rays, 3] | |
| ray_d = torch.bmm(ray_d, c2w[:, :3, :3].transpose(1, 2)) | |
| ray_d = ray_d / torch.norm(ray_d, dim=2, keepdim=True) | |
| ray_o = c2w[:, :3, 3][:, None, :].expand_as(ray_d) | |
| ray_color = ( | |
| ray_color.reshape(b, v, c, -1).permute(0, 1, 3, 2).reshape(b, -1, c) | |
| ) # [b, v*(crop_size*crop_size+n_rand_rays), c] | |
| ray_o = ray_o.reshape(b, -1, 3) # [b, v*(crop_size*crop_size+n_rand_rays), 3] | |
| ray_d = ray_d.reshape(b, -1, 3) # [b, v*(crop_size*crop_size+n_rand_rays), 3] | |
| image = image.reshape(b, v, c, h, w) | |
| fxfycxcy = fxfycxcy.reshape(b, v, 4) | |
| c2w = c2w.reshape(b, v, 4, 4) | |
| input = edict( | |
| image=image, | |
| fxfycxcy=fxfycxcy, | |
| c2w=c2w, | |
| mv=mv, | |
| mvp=mvp, | |
| depth=depth, | |
| normal=normal, | |
| index=index, | |
| ray_o=ray_o, | |
| ray_d=ray_d, | |
| ) | |
| return input | |