| from functools import partial | |
| import torch | |
| def get_plucker_embedding(intrinsics, cam_c2w, height, width, height_dit=None, width_dit=None, flip_flag=None): | |
| custom_meshgrid = partial(torch.meshgrid, indexing="ij") | |
| batch_size, num_frames = intrinsics.shape[:2] | |
| use_dit_hw = height_dit is not None and width_dit is not None | |
| if not use_dit_hw: | |
| height_dit = height | |
| width_dit = width | |
| else: | |
| patch_height = height / height_dit | |
| patch_width = width / width_dit | |
| j, i = custom_meshgrid( | |
| torch.linspace(0, height_dit - 1, height_dit, device=cam_c2w.device, dtype=cam_c2w.dtype), | |
| torch.linspace(0, width_dit - 1, width_dit, device=cam_c2w.device, dtype=cam_c2w.dtype), | |
| ) | |
| i = i.reshape(1, 1, height_dit * width_dit).expand(batch_size, num_frames, height_dit * width_dit) + 0.5 | |
| j = j.reshape(1, 1, height_dit * width_dit).expand(batch_size, num_frames, height_dit * width_dit) + 0.5 | |
| if use_dit_hw: | |
| i = i * patch_width + (patch_width / 2) | |
| j = j * patch_height + (patch_height / 2) | |
| if flip_flag is not None and torch.sum(flip_flag).item() > 0: | |
| j_flip, i_flip = custom_meshgrid( | |
| torch.linspace(0, height_dit - 1, height_dit, device=cam_c2w.device, dtype=cam_c2w.dtype), | |
| torch.linspace(width_dit - 1, 0, width_dit, device=cam_c2w.device, dtype=cam_c2w.dtype), | |
| ) | |
| i_flip = i_flip.reshape(1, 1, height_dit * width_dit).expand(batch_size, 1, height_dit * width_dit) + 0.5 | |
| j_flip = j_flip.reshape(1, 1, height_dit * width_dit).expand(batch_size, 1, height_dit * width_dit) + 0.5 | |
| if use_dit_hw: | |
| i_flip = i_flip * patch_width + (patch_width / 2) | |
| j_flip = j_flip * patch_height + (patch_height / 2) | |
| i[:, flip_flag, ...] = i_flip | |
| j[:, flip_flag, ...] = j_flip | |
| fx, fy, cx, cy = intrinsics.chunk(4, dim=-1) | |
| zs = torch.ones_like(i) | |
| xs = (i - cx) / fx * zs | |
| ys = (j - cy) / fy * zs | |
| zs = zs.expand_as(ys) | |
| directions = torch.stack((xs, ys, zs), dim=-1) | |
| directions = directions / directions.norm(dim=-1, keepdim=True) | |
| rays_d = directions @ cam_c2w[..., :3, :3].transpose(-1, -2) | |
| rays_o = cam_c2w[..., :3, 3] | |
| rays_o = rays_o[:, :, None].expand_as(rays_d) | |
| rays_dxo = torch.cross(rays_o, rays_d, dim=-1) | |
| plucker = torch.cat([rays_dxo, rays_d], dim=-1) | |
| return plucker.reshape(batch_size, num_frames, height_dit, width_dit, 6) | |