From 741e39e17088682cfa21a06af234756312e5ed8f Mon Sep 17 00:00:00 2001 From: mokad6 Date: Wed, 27 Sep 2023 00:31:04 -0400 Subject: [PATCH] colmap with multiple exposure --- .../colmap_converter_to_nerfstudio_dataset.py | 40 ++++ nerfstudio/process_data/colmap_utils.py | 101 +++++++++ nerfstudio/process_data/equirect_utils.py | 191 ++++++++++++++++++ .../lantern_images_to_nerfstudio_dataset.py | 181 +++++++++++++++++ nerfstudio/scripts/process_data.py | 4 + 5 files changed, 517 insertions(+) create mode 100644 nerfstudio/process_data/lantern_images_to_nerfstudio_dataset.py diff --git a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py index 87287757d..fefd54479 100644 --- a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py @@ -109,6 +109,46 @@ class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset): def absolute_colmap_path(self) -> Path: return self.output_dir / "colmap" + + def _save_transforms_multiple_exposures( + self, + num_frames: int, + mask1_path: Path, + mask2_path: Path, + e1_path: Path, + e2_path: Path, + image_rename_map: Optional[Dict[str, str]] = None, + ) -> List[str]: + """Save colmap transforms into the output folder + + Args: + image_id_to_depth_path: When including sfm-based depth, embed these depth file paths in the exported json + image_rename_map: Use these image names instead of the names embedded in the COLMAP db + """ + summary_log = [] + if (self.absolute_colmap_model_path / "cameras.bin").exists(): + with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): + num_matched_frames = colmap_utils.colmap_to_json_for_different_Exposure( + recon_dir=self.absolute_colmap_model_path, + output_dir=self.output_dir, + mask1_path=mask1_path, + mask2_path=mask2_path, + e1_path=e1_path, + e2_path=e2_path, + image_rename_map=image_rename_map, + ) + summary_log.append(f"Colmap matched {num_matched_frames} images") + summary_log.append(colmap_utils.get_matching_summary(num_frames, num_matched_frames)) + + else: + CONSOLE.log( + "[bold yellow]Warning: Could not find existing COLMAP results. " "Not generating transforms.json" + ) + return summary_log + + + + def _save_transforms( self, num_frames: int, diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 449da2208..2231d7e72 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -379,6 +379,107 @@ def parse_colmap_camera_params(camera) -> Dict[str, Any]: return out +def colmap_to_json_for_different_Exposure( + recon_dir: Path, + output_dir: Path, + mask1_path: Path = None, + mask2_path: Path = None, + e1_path: Path = None, + e2_path: Path = None, + image_rename_map: Optional[Dict[str, str]] = None, +) -> int: + """Converts COLMAP's cameras.bin and images.bin to a JSON file. + + Args: + recon_dir: Path to the reconstruction directory, e.g. "sparse/0" + output_dir: Path to the output directory. + camera_model: Camera model used. + camera_mask_path: Path to the camera mask. + image_id_to_depth_path: When including sfm-based depth, embed these depth file paths in the exported json + image_rename_map: Use these image names instead of the names embedded in the COLMAP db + + Returns: + The number of registered images. + """ + + # TODO(1480) use pycolmap + # recon = pycolmap.Reconstruction(recon_dir) + # cam_id_to_camera = recon.cameras + # im_id_to_image = recon.images + cam_id_to_camera = read_cameras_binary(recon_dir / "cameras.bin") + im_id_to_image = read_images_binary(recon_dir / "images.bin") + + frames = [] + for im_id, im_data in im_id_to_image.items(): + # NB: COLMAP uses Eigen / scalar-first quaternions + # * https://colmap.github.io/format.html + # * https://github.com/colmap/colmap/blob/bf3e19140f491c3042bfd85b7192ef7d249808ec/src/base/pose.cc#L75 + # the `rotation_matrix()` handles that format for us. + + # TODO(1480) BEGIN use pycolmap API + # rotation = im_data.rotation_matrix() + rotation = qvec2rotmat(im_data.qvec) + + translation = im_data.tvec.reshape(3, 1) + w2c = np.concatenate([rotation, translation], 1) + w2c = np.concatenate([w2c, np.array([[0, 0, 0, 1]])], 0) + c2w = np.linalg.inv(w2c) + # Convert from COLMAP's camera coordinate system (OpenCV) to ours (OpenGL) + c2w[0:3, 1:3] *= -1 + c2w = c2w[np.array([1, 0, 2, 3]), :] + c2w[2, :] *= -1 + + name = im_data.name + if image_rename_map is not None: + name = image_rename_map[name] + e1_name = name.replace('.png', '.exr') + name = Path(e1_path.absolute() / e1_name) + mask_name = im_data.name + mask1_name = mask_name #.replace('.png', '.png') + mask_path = Path(mask1_path.absolute() / mask1_name).as_posix() + + frame = { + "file_path": name.as_posix(), + "transform_matrix": c2w.tolist(), + # "colmap_im_id": im_id, + "mask_path": mask_path, + } + frames.append(frame) + + name = im_data.name + if image_rename_map is not None: + name = image_rename_map[name] + e2_name = name.replace('.png', '.exr') + name = Path(e2_path.absolute() / e2_name) + mask_name = im_data.name + mask2_name = mask_name #.replace('.png', '.png') + mask_path = Path(mask2_path.absolute() / mask2_name).as_posix() + + frame = { + "file_path": name.as_posix(), + "transform_matrix": c2w.tolist(), + # "colmap_im_id": im_id, + "mask_path": mask_path, + } + frames.append(frame) + + if set(cam_id_to_camera.keys()) != {1}: + raise RuntimeError("Only single camera shared for all images is supported.") + out = parse_colmap_camera_params(cam_id_to_camera[1]) + out["frames"] = frames + + applied_transform = np.eye(4)[:3, :] + applied_transform = applied_transform[np.array([1, 0, 2]), :] + applied_transform[2, :] *= -1 + out["applied_transform"] = applied_transform.tolist() + + with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: + json.dump(out, f, indent=4) + + return len(frames) + + + def colmap_to_json( recon_dir: Path, output_dir: Path, diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index 722f8a2c5..34e55e989 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -272,6 +272,197 @@ def generate_planar_projections_from_equirectangular( return output_dir, output_mask_dir, output_hdr_dir + +def generate_planar_projections_from_equirectangular_with_two_exposures( + image_dir: Path, + planar_image_size: Tuple[int, int], + samples_per_im: int, + mask_e1_dir: Path, + mask_e2_dir: Path, + e1_dir: Path, + e2_dir: Path, + e1_factor: float, + e2_factor: float, + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), +) -> Path: + """Generate planar projections from an equirectangular image. + + Args: + image_dir: The directory containing the equirectangular image. + planar_image_size: The size of the planar projections [width, height]. + samples_per_im: The number of samples to take per image. + crop_factor: The portion of the image to crop from the (top, bottom, left, and right). + Values should be in [0, 1]. + returns: + The path to the planar projections directory. + """ + + for i in crop_factor: + if i < 0 or i > 1: + CONSOLE.print("[bold red] Invalid crop factor. All values must be in [0,1].") + sys.exit(1) + + device = torch.device("cuda") + + fov = 120 + yaw_pitch_pairs = [] + left_bound, right_bound = -180, 180 + if crop_factor[3] > 0: + left_bound = -180 + 360 * crop_factor[3] + if crop_factor[2] > 0: + right_bound = 180 - 360 * crop_factor[2] + + if samples_per_im == 8: + fov = 120 + bound_arr = [-45, 0, 45] + bound_arr = _crop_bound_arr_vertical(bound_arr, fov, crop_factor) + if bound_arr[1] is not None: + for i in np.arange(left_bound, right_bound, 90): + yaw_pitch_pairs.append((i, bound_arr[1])) + if bound_arr[2] is not None: + for i in np.arange(left_bound, right_bound, 180): + yaw_pitch_pairs.append((i, bound_arr[2])) + if bound_arr[0] is not None: + for i in np.arange(left_bound, right_bound, 180): + yaw_pitch_pairs.append((i, bound_arr[0])) + elif samples_per_im == 14: + fov = 110 + bound_arr = [-45, 0, 45] + bound_arr = _crop_bound_arr_vertical(bound_arr, fov, crop_factor) + if bound_arr[1] is not None: + for i in np.arange(left_bound, right_bound, 60): + yaw_pitch_pairs.append((i, bound_arr[1])) + if bound_arr[2] is not None: + for i in np.arange(left_bound, right_bound, 90): + yaw_pitch_pairs.append((i, bound_arr[2])) + if bound_arr[0] is not None: + for i in np.arange(left_bound, right_bound, 90): + yaw_pitch_pairs.append((i, bound_arr[0])) + + equi2pers = Equi2Pers(height=planar_image_size[1], width=planar_image_size[0], fov_x=fov, mode="bilinear") + frame_dir = image_dir + output_dir = image_dir / "planar_projections" + output_dir.mkdir(exist_ok=True) + + output_mask_e1_dir = image_dir / "mask_e1_planar_projections" + output_mask_e1_dir.mkdir(exist_ok=True) + + output_mask_e2_dir = image_dir / "mask_e2_planar_projections" + output_mask_e2_dir.mkdir(exist_ok=True) + + output_e1_dir = image_dir / "e1_planar_projections" + output_e1_dir.mkdir(exist_ok=True) + + output_e2_dir = image_dir / "e2_planar_projections" + output_e2_dir.mkdir(exist_ok=True) + + num_ims = len(os.listdir(frame_dir)) + progress = Progress( + TextColumn("[bold blue]Generating Planar Images", justify="right"), + BarColumn(), + TaskProgressColumn(show_speed=True), + ItersPerSecColumn(suffix="equirect frames/s"), + TimeRemainingColumn(elapsed_when_finished=True, compact=True), + ) + + with progress: + for i in progress.track(os.listdir(frame_dir), description="", total=num_ims): + if i.lower().endswith((".jpg", ".png", ".jpeg")): + im = np.array(cv2.imread(os.path.join(frame_dir, i))) + im = torch.tensor(im, dtype=torch.float32, device=device) + im = torch.permute(im, (2, 0, 1)) + im /=255.0 + + mask_e1_addr = os.path.join(mask_e1_dir, i) + mask_e1_addr = mask_e1_addr.replace('.png', '_e1.png') + if not os.path.isfile(mask_e1_addr): + raise Exception("The corresponding mask E1 for " + i + " does not exist!!!") + mask_e1 = np.array(cv2.imread(mask_e1_addr)) + mask_e1 = torch.tensor(mask_e1, dtype=torch.uint8, device=device) + mask_e1 = torch.permute(mask_e1, (2, 0, 1)) + + mask_e2_addr = os.path.join(mask_e2_dir, i) + mask_e2_addr = mask_e2_addr.replace('.png', '_e2.png') + if not os.path.isfile(mask_e2_addr): + raise Exception("The corresponding mask E2 for " + i + " does not exist!!!") + mask_e2 = np.array(cv2.imread(mask_e2_addr)) + mask_e2 = torch.tensor(mask_e2, dtype=torch.uint8, device=device) + mask_e2 = torch.permute(mask_e2, (2, 0, 1)) + + e1_addr = os.path.join(e1_dir, i) + if e1_addr.endswith('.png'): + e1_addr = e1_addr.replace('.png', '.exr') + elif e1_addr.endswith('.jpg'): + e1_addr = e1_addr.replace('.jpg', '.exr') + else: + raise Exception("Image format for " + i + " is not supported!!!") + e1_addr = e1_addr.replace('.exr', '_e1.exr') + if not os.path.isfile(e1_addr): + raise Exception("The corresponding e1 for " + i + " does not exist!!!") + e1 = np.array(cv2.imread(e1_addr, cv2.IMREAD_UNCHANGED)).astype("float32") + e1 = e1 / e1_factor + e1 = torch.tensor(e1, dtype=torch.float32, device=device) + e1 = torch.permute(e1, (2, 0, 1)) + + e2_addr = os.path.join(e2_dir, i) + if e2_addr.endswith('.png'): + e2_addr = e2_addr.replace('.png', '.exr') + elif e2_addr.endswith('.jpg'): + e2_addr = e2_addr.replace('.jpg', '.exr') + else: + raise Exception("Image format for " + i + " is not supported!!!") + e2_addr = e2_addr.replace('.exr', '_e2.exr') + if not os.path.isfile(e2_addr): + raise Exception("The corresponding e1 for " + i + " does not exist!!!") + # e2 = np.array(cv2.imread(e2_addr, cv2.IMREAD_UNCHANGED)).astype("float32") + e2 = np.array(cv2.imread(e2_addr, cv2.IMREAD_UNCHANGED) * 255.0).astype("uint8") + e2 = torch.tensor(e2, dtype=torch.float32, device=device) + e2 = e2 / 255.0 + e2 = e2 / e2_factor + e2 = torch.permute(e2, (2, 0, 1)) + + count = 0 + for u_deg, v_deg in yaw_pitch_pairs: + v_rad = torch.pi * v_deg / 180.0 + u_rad = torch.pi * u_deg / 180.0 + pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) + pers_mask_e1 = equi2pers(mask_e1, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) + pers_mask_e2 = equi2pers(mask_e2, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) + pers_e1 = equi2pers(e1, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}, clip_output=False) + pers_e2 = equi2pers(e2, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}, clip_output=False) + + assert isinstance(pers_image, torch.Tensor) + if i.lower().endswith((".exr")): + # normalize alpha channel + pers_image = (pers_image.permute(1, 2, 0)).type(torch.float32).to("cpu").numpy() + cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.exr", pers_image) + else: + pers_image *= 255.0 + pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy() + cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.png", pers_image) + + assert isinstance(pers_mask_e1, torch.Tensor) + pers_mask_e1 = (pers_mask_e1.permute(1, 2, 0))[:, :, 0].type(torch.uint8).to("cpu").numpy() + cv2.imwrite(f"{output_mask_e1_dir}/{i[:-4]}_{count}.png", pers_mask_e1) + + assert isinstance(pers_mask_e2, torch.Tensor) + pers_mask_e2 = (pers_mask_e2.permute(1, 2, 0))[:, :, 0].type(torch.uint8).to("cpu").numpy() + cv2.imwrite(f"{output_mask_e2_dir}/{i[:-4]}_{count}.png", pers_mask_e2) + + assert isinstance(pers_e1, torch.Tensor) + pers_e1 = (pers_e1.permute(1, 2, 0)).type(torch.float32).to("cpu").numpy() + cv2.imwrite(f"{output_e1_dir}/{i[:-4]}_{count}.exr", pers_e1) + + assert isinstance(pers_e2, torch.Tensor) + pers_e2 = (pers_e2.permute(1, 2, 0)).type(torch.float32).to("cpu").numpy() + cv2.imwrite(f"{output_e2_dir}/{i[:-4]}_{count}.exr", pers_e2) + + count += 1 + + return output_dir, output_e1_dir, output_e2_dir, output_mask_e1_dir, output_mask_e2_dir + + + def generate_planar_projections_from_equirectangular_GT( metadata_path: Path, image_dir: Path, diff --git a/nerfstudio/process_data/lantern_images_to_nerfstudio_dataset.py b/nerfstudio/process_data/lantern_images_to_nerfstudio_dataset.py new file mode 100644 index 000000000..6f8fa0574 --- /dev/null +++ b/nerfstudio/process_data/lantern_images_to_nerfstudio_dataset.py @@ -0,0 +1,181 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. +# +# 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. + +"""Processes an image sequence to a nerfstudio compatible dataset.""" + +from dataclasses import dataclass +from typing import Optional + +from nerfstudio.process_data import equirect_utils, process_data_utils +from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import ColmapConverterToNerfstudioDataset + +from nerfstudio.utils.rich_utils import CONSOLE + +from pathlib import Path + + +@dataclass +class LanternImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): + """Process images into a nerfstudio dataset. + + 1. Scales images to a specified size. + 2. Calculates the camera poses for each image using `COLMAP `_. + """ + + percent_radius_crop: float = 1.0 + """Create circle crop mask. The radius is the percent of the image diagonal.""" + + e1_dir: Path = None + """Path to the E1 directory.""" + e2_dir: Path = None + """Path to the E2 directory.""" + + mask_e1_dir: Path = None + """Path to the E1 mask directory.""" + mask_e2_dir: Path = None + """Path to the E2 mask directory.""" + + exposure1: float = 1.0 + """Path to the E2 mask directory.""" + + exposure2: float = 0.009 + """Path to the E2 mask directory.""" + + def main(self) -> None: + """Process images into a nerfstudio dataset.""" + + self.is_HDR = True + + require_cameras_exist = False + if self.colmap_model_path != ColmapConverterToNerfstudioDataset.default_colmap_path(): + if not self.skip_colmap: + raise RuntimeError("The --colmap-model-path can only be used when --skip-colmap is not set.") + if not (self.output_dir / self.colmap_model_path).exists(): + raise RuntimeError(f"The colmap-model-path {self.output_dir / self.colmap_model_path} does not exist.") + require_cameras_exist = True + + image_rename_map: Optional[dict[str, str]] = None + + # Generate planar projections if equirectangular + if self.camera_type == "equirectangular": + if self.eval_data is not None: + raise ValueError("Cannot use eval_data with camera_type equirectangular.") + pers_size = equirect_utils.compute_resolution_from_equirect(self.data, self.images_per_equirect) + CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image") + self.data, perspective_e1, perspective_e2, perspective_mask1, perspective_mask2 = equirect_utils.generate_planar_projections_from_equirectangular_with_two_exposures( + self.data, pers_size, self.images_per_equirect, self.mask_e1_dir, self.mask_e2_dir, self.e1_dir, self.e2_dir, self.exposure1, self.exposure2, crop_factor=self.crop_factor, + ) + self.camera_type = "perspective" + + summary_log = [] + + # Copy and downscale images + if not self.skip_image_processing: + # Copy images to output directory + + image_rename_map_paths = process_data_utils.copy_images( + self.data, + image_dir=self.image_dir, + crop_factor=self.crop_factor, + verbose=self.verbose, + num_downscales=self.num_downscales, + ) + + if self.eval_data is not None: + eval_image_rename_map_paths = process_data_utils.copy_images( + self.eval_data, + image_dir=self.image_dir, + crop_factor=self.crop_factor, + image_prefix="frame_eval_", + verbose=self.verbose, + num_downscales=self.num_downscales, + ) + image_rename_map_paths.update(eval_image_rename_map_paths) + + (self.output_dir / "masks1").mkdir(parents=True, exist_ok=True) + mask_rename_map_paths = process_data_utils.copy_images( + perspective_mask1, + image_dir=self.output_dir / "masks1", + num_downscales=self.num_downscales, + verbose=self.verbose, + crop_factor=self.crop_factor + ) + + (self.output_dir / "masks2").mkdir(parents=True, exist_ok=True) + mask_rename_map_paths = process_data_utils.copy_images( + perspective_mask2, + image_dir=self.output_dir / "masks2", + num_downscales=self.num_downscales, + verbose=self.verbose, + crop_factor=self.crop_factor + ) + + (self.output_dir / "exposure1").mkdir(parents=True, exist_ok=True) + copied_image_paths = process_data_utils.copy_images( + perspective_e1, + image_dir=self.output_dir / "exposure1", + verbose=self.verbose, + num_downscales=self.num_downscales, + crop_factor=self.crop_factor, + is_HDR = True + ) + + (self.output_dir / "exposure2").mkdir(parents=True, exist_ok=True) + copied_image_paths = process_data_utils.copy_images( + perspective_e2, + image_dir=self.output_dir / "exposure2", + verbose=self.verbose, + num_downscales=self.num_downscales, + crop_factor=self.crop_factor, + is_HDR = True + ) + + image_rename_map = dict((a.name, b.name) for a, b in image_rename_map_paths.items()) + num_frames = len(image_rename_map) + summary_log.append(f"Starting with {num_frames} images") + + else: + num_frames = len(process_data_utils.list_images(self.data)) + if num_frames == 0: + raise RuntimeError("No usable images in the data folder.") + summary_log.append(f"Starting with {num_frames} images") + + image_rename_map = None + # Run COLMAP + if not self.skip_colmap: + require_cameras_exist = True + self._run_colmap() + # Colmap uses renamed images + image_rename_map = None + + # # Export depth maps + image_id_to_depth_path, log_tmp = self._export_depth() + summary_log += log_tmp + + if require_cameras_exist and not (self.absolute_colmap_model_path / "cameras.bin").exists(): + raise RuntimeError(f"Could not find existing COLMAP results ({self.colmap_model_path / 'cameras.bin'}).") + + summary_log += self._save_transforms_multiple_exposures( + num_frames, + (self.output_dir / "masks1"), + (self.output_dir / "masks2"), + (self.output_dir / "exposure1"), + (self.output_dir / "exposure2"), + image_rename_map + ) + + CONSOLE.log("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") + + for summary in summary_log: + CONSOLE.log(summary) diff --git a/nerfstudio/scripts/process_data.py b/nerfstudio/scripts/process_data.py index d273cdb19..1d2568703 100644 --- a/nerfstudio/scripts/process_data.py +++ b/nerfstudio/scripts/process_data.py @@ -39,6 +39,9 @@ from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import ( from nerfstudio.process_data.images_to_nerfstudio_dataset import ( ImagesToNerfstudioDataset, ) +from nerfstudio.process_data.lantern_images_to_nerfstudio_dataset import ( + LanternImagesToNerfstudioDataset, +) from nerfstudio.process_data.aligned_pano2plane import ( ProcessAlignedPano ) @@ -402,6 +405,7 @@ Commands = Union[ Annotated[ProcessRealityCapture, tyro.conf.subcommand(name="realitycapture")], Annotated[ProcessRecord3D, tyro.conf.subcommand(name="record3d")], Annotated[ProcessAlignedPano, tyro.conf.subcommand(name="aligned_pano2plane")], + Annotated[LanternImagesToNerfstudioDataset, tyro.conf.subcommand(name="lantern")], ]