Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as nnf | |
| import addict | |
| class SafeDict(addict.Dict): | |
| def __missing__(self, name): | |
| if object.__getattribute__(self, '__frozen'): | |
| raise AttributeError(name) | |
| return self.__class__(__parent=self, __key=name) | |
| def __setitem__(self, name, value): | |
| is_frozen = (hasattr(self, '__frozen') and | |
| object.__getattribute__(self, '__frozen')) | |
| assert not is_frozen, f"Cannot set field '{name}': frozen dict." | |
| super().__setitem__(name, value) | |
| try: | |
| p = object.__getattribute__(self, '__parent') | |
| key = object.__getattribute__(self, '__key') | |
| except AttributeError: | |
| p = None | |
| key = None | |
| if p is not None: | |
| p[key] = self | |
| object.__delattr__(self, '__parent') | |
| object.__delattr__(self, '__key') | |
| class TransformInput(nn.Module): | |
| def __init__(self, wo_grad=True): | |
| super().__init__() | |
| self.wo_grad = wo_grad | |
| def forward(self, image, intrinsics, c2w, patch_size=None): | |
| # pose stage model requires grad | |
| grad_context = torch.no_grad if self.wo_grad else torch.enable_grad | |
| with grad_context(): | |
| return self._forward(image, intrinsics, c2w, patch_size) | |
| def _forward(self, image, intrinsics, c2w, patch_size=None): | |
| """transform input image before feeding into transformer | |
| Args: | |
| data_batch: SafeDict | |
| patch_size: int, optional | |
| Returns: | |
| image: [b, v, c, h, w] | |
| ray_o: [b, v, 3, h, w] | |
| ray_d: [b, v, 3, h, w] | |
| """ | |
| # image, fxfycxcy, c2w, index = ( | |
| # data_batch.image, | |
| # data_batch.fxfycxcy, | |
| # data_batch.c2w, | |
| # 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()}" | |
| fxfycxcy = torch.stack( | |
| [intrinsics[..., 0,0], intrinsics[..., 1,1], | |
| intrinsics[..., 0,2], intrinsics[..., 1,2]], dim=-1 | |
| ) | |
| 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) | |
| y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") | |
| y, x = y.to(image.device), x.to(image.device) | |
| # [h, w] | |
| 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 = (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 = torch.stack([x, y, z], dim=2) # [b*v, h*w, 3] | |
| ray_d_cam = ray_d.clone() | |
| ray_d = torch.bmm(ray_d, c2w[:, :3, :3].transpose(1, 2)) # [b*v, h*w, 3] | |
| 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 = ( | |
| nnf.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) | |
| ret = SafeDict( | |
| image=image, | |
| ray_o=ray_o, | |
| ray_d=ray_d, | |
| ray_d_cam=ray_d_cam, | |
| fxfycxcy=fxfycxcy, | |
| c2w=c2w, | |
| # index=index, | |
| xy_norm=xy_norm, | |
| ray_color_patch=ray_color_patch, | |
| ray_o_patch=ray_o_patch, | |
| ray_d_patch=ray_d_patch, | |
| ray_xy_norm_patch=ray_xy_norm_patch, | |
| proj_mat=proj_mat, | |
| ) | |
| return ret | |