# LICENSE HEADER MANAGED BY add-license-header # # Copyright 2018 Kornia Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import math from typing import List, Tuple, Union import torch from kornia.core import Device, Tensor, cos, sin, stack from kornia.geometry.camera import PinholeCamera from kornia.geometry.conversions import quaternion_to_rotation_matrix def parse_colmap_output( cameras_path: str, images_path: str, device: Device, dtype: torch.dtype ) -> Tuple[List[str], PinholeCamera]: r"""Parse colmap output to create an PinholeCamera for aligned scene cameras. Args: cameras_path: Path to camera.txt Colmap file with camera intrinsics: str images_path: Path to images.txt Colmap file with camera extrinsics for each image: str device: device for created camera object: Union[str, torch.device] dtype: Intrinsics and extrinsics dtype. Returns: image names: List[str] scene camera object: PinholeCamera """ # Parse camera intrinsics with open(cameras_path) as f: lines = f.readlines() class CameraParams: def __init__(self, line: str) -> None: split_line = line.split(" ") model = split_line[1] if model == "SIMPLE_PINHOLE": self._width = int(split_line[2]) self._height = int(split_line[3]) self._fx = float(split_line[4]) self._fy = self._fx self._cx = int(split_line[5]) self._cy = int(split_line[6]) elif model == "PINHOLE": self._width = int(split_line[2]) self._height = int(split_line[3]) self._fx = float(split_line[4]) self._fy = float(split_line[5]) self._cx = int(split_line[6]) self._cy = int(split_line[7]) cameras_params: List[CameraParams] = [] for line in lines: if line.startswith("#"): continue camera_params = CameraParams(line) cameras_params.append(camera_params) # Parse camera quaternions and translation vectors with open(images_path) as f: lines = f.readlines() intrinsics: List[Tensor] = [] extrinsics: List[Tensor] = [] heights: List[int] = [] widths: List[int] = [] img_names: List[str] = [] for line in lines: if line.startswith("#"): continue # Read line with camera quaternion line = line.strip() if line.endswith(("jpg", "png")): split_line = line.split(" ") qw = float(split_line[1]) qx = float(split_line[2]) qy = float(split_line[3]) qz = float(split_line[4]) tx = float(split_line[5]) ty = float(split_line[6]) tz = float(split_line[7]) camera_ind = int(split_line[8]) - 1 img_name = split_line[9] img_names.append(img_name) # Intrinsic camera_params = cameras_params[camera_ind] intrinsic = torch.eye(4, device=device, dtype=dtype) intrinsic[0, 0] = camera_params._fx intrinsic[1, 1] = camera_params._fy intrinsic[0, 2] = camera_params._cx intrinsic[1, 2] = camera_params._cy intrinsics.append(intrinsic) heights.append(camera_params._height) widths.append(camera_params._width) # Extrinsic q = torch.tensor([qw, qx, qy, qz], device=device) R = quaternion_to_rotation_matrix(q) t = torch.tensor([tx, ty, tz], device=device) extrinsic = torch.eye(4, device=device, dtype=dtype) extrinsic[:3, :3] = R extrinsic[:3, 3] = t extrinsics.append(extrinsic) cameras = PinholeCamera( torch.stack(intrinsics), torch.stack(extrinsics), torch.tensor(heights, device=device), torch.tensor(widths, device=device), ) return img_names, cameras def cameras_for_ids(cameras: PinholeCamera, camera_ids: Union[List[int], Tensor]) -> PinholeCamera: r"""Take a PinholeCamera camera and camera indices to create a new PinholeCamera for requested cameras. Args: cameras: Scene camera object: PinholeCamera camera_ids: List of camera indices to copy: List[int] Return: A new PinholeCamera object with a sub-set of cameras: PinholeCamera """ intrinsics = cameras.intrinsics[camera_ids] extrinsics = cameras.extrinsics[camera_ids] height = cameras.height[camera_ids] width = cameras.width[camera_ids] return PinholeCamera(intrinsics, extrinsics, height, width) def create_spiral_path(cameras: PinholeCamera, rad: float, num_views: int, num_circles: int) -> PinholeCamera: r"""Create a PinholeCamera object with cameras that follow a spiral path. Used for novel view synthesis for face facing models. Args: cameras: Scene cameras used to train the NeRF model: PinholeCamera rad: Spiral radius: float num_views: Number of created cameras: int num_circles: Number of spiral circles: int """ # Average locations over all cameras mean_center = cameras.translation_vector.mean(0, False).squeeze(-1) device = cameras.intrinsics.device t = torch.linspace(0, 2 * math.pi * num_circles, num_views, device=device) cos_t = cos(t) * rad sin_t = -sin(t) * rad sin_05t = -sin(0.5 * t) * rad translation_vector = torch.unsqueeze(mean_center, dim=0) + stack((cos_t, sin_t, sin_05t)).permute((1, 0)) mean_intrinsics = cameras.intrinsics.mean(0, True).repeat(num_views, 1, 1) mean_extrinsics = cameras.extrinsics.mean(0, True).repeat(num_views, 1, 1) extrinsics = mean_extrinsics extrinsics[:, :3, 3] = translation_vector height = torch.tensor([cameras.height[0]] * num_views, device=device) width = torch.tensor([cameras.width[0]] * num_views, device=device) return PinholeCamera(mean_intrinsics, extrinsics, height, width)