diff --git a/README-pandora.md b/README-pandora.md new file mode 100644 index 000000000..ca021530d --- /dev/null +++ b/README-pandora.md @@ -0,0 +1,96 @@ + + +# How it works? + +## 0. Setup + +``` + +pip install skylibs equilib +``` + +## 1. Data acquisition & pre-processing + +### OPTION 1: Real data + +1. Acquire data using the apparatus (stick with two Ricoh cameras) + + Make sure the cameras point up to avoid having lights in the seam. + + Do two recordings with both Ricoh Theta cameras. Use the flash of a flashlight both times, the falling edge will later be detected by a software. + + > Cameras: + > + > **camera 1**: left, well-exposed, serial: YN14100695 \ + > **camera 3**: right, under-exposed, serial: YN14111000 + +2. Take one video capture with the two cameras well-exposed. +3. Take one video capture with Camera 1 (left) well-exposed and Camera 2 (right) fast-exposed. +4. Run the script for automatic flash detection + + It will temporally align the frames from the left and right camera and remove the initial frames with the flash. + + It should output the following file structure + ``` + left_colmap_baseline/ + 00000.png + 00001.png + ... + 00XXX.png + right_colmap_baseline/ + 00000.png + 00001.png + ... + 00XXX.png + left_well_exposed/ + 00000.png + 00001.png + ... + 00XXX.png + right_fast_exposed/ + 00000.png + 00001.png + ... + 00XXX.png + ``` +5. + + +### OPTION 2: Synthetic data + +`` +export PYTHONPATH=$pwd:$PYTHONPATH +python lantern_scripts/split_synthetic_colmap.py --has_exposure --left_dir ./Dual_Cameras/left --right_dir ./Dual_Cameras/right --out_dir ./Dual_Cameras/split_colmap_out +`` + +## 2. Process data using COLMAP + +``` +ns-process-data lantern --data ./Dual_Cameras/split_colmap_subset --output-dir ./Dual_Cameras/split_colmap_subset/split_colmap_out5 --skip-image-processing --camera-type equirectangular --images-per-equirect 8 +``` + +> Thing to keep in mind: the way Nerfstudio splits train and test images may not be optimal: left might be in one dataset and right in another... that's a bit of a data leak. + +## 2. Pre-training NeRF + +Using both the well-exposed and fast-exposed frames, as well as their associated camera poses, a NeRF will be trained to predict both the RGB and mask to do volumetric rendering. + +After convergence, go to the next step. + +``` +ns-train lantern-nerfacto --help +``` + +## 3. Running MomoNet + +(Using which frames?) A pretrained UNet image-to-image network called MomoNet will infill the missing dynamic range and store these predictions in new images. +It will be trained on the Laval indoor HDR dataset to do this task. + +TODO: validate this part + +## 4. Train from scratch or fine-tune the NeRF model using MomoNet's predictions + + + +``` +``` \ No newline at end of file diff --git a/lantern_scripts/combine_left_right.py b/lantern_scripts/combine_left_right.py index 854fe544b..212afd045 100644 --- a/lantern_scripts/combine_left_right.py +++ b/lantern_scripts/combine_left_right.py @@ -1,14 +1,14 @@ import argparse -import numpy as np -import os -from tqdm import tqdm -from pathlib import Path -from PIL import Image import copy -import cv2 import json +import os from pathlib import Path +import cv2 +import numpy as np +from PIL import Image +from tqdm import tqdm + os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" @@ -83,7 +83,7 @@ def write_to_json(filename: Path, content: dict): def process_one_image(pixels, out_dir, name, has_exposure, exposure_value, zmin, zmax): - if args.has_exposure: + if has_exposure: e1_ldr, e1_hdr, weights_e1 = apply_the_exposure_get_weights(pixels, exposure_value) weights_e1_unit = make_weights_binary(e1_ldr, weights_e1, zmin, zmax) mask_e1 = weights_e1_unit[:, :, 0] & weights_e1_unit[:, :, 1] & weights_e1_unit[:, :, 2] @@ -94,6 +94,7 @@ def process_one_image(pixels, out_dir, name, has_exposure, exposure_value, zmin, e1_hdr = pixels # out_addr = os.path.join(args.out_dir, pano_name + '_e1.exr') out_addr = os.path.join(out_dir, name + '.exr') + os.makedirs(os.path.dirname(out_addr), exist_ok=True) cv2.imwrite(out_addr, e1_hdr) diff --git a/nerfstudio/process_data/images_to_nerfstudio_dataset.py b/nerfstudio/process_data/images_to_nerfstudio_dataset.py index 90e66a292..1b2083035 100644 --- a/nerfstudio/process_data/images_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/images_to_nerfstudio_dataset.py @@ -60,6 +60,11 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): is_HDR=self.is_HDR, HDR_planar_image_size=pers_size_hdr ) self.camera_type = "perspective" + else: + # make all none + perspective_masks = None + perspective_hdrs = None + summary_log = [] diff --git a/nerfstudio/process_data/lantern_images_to_nerfstudio_dataset.py b/nerfstudio/process_data/lantern_images_to_nerfstudio_dataset.py index 6f8fa0574..7cb643809 100644 --- a/nerfstudio/process_data/lantern_images_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/lantern_images_to_nerfstudio_dataset.py @@ -14,16 +14,25 @@ """Processes an image sequence to a nerfstudio compatible dataset.""" +import json +import os +import re +import shutil +import tempfile +from copy import deepcopy from dataclasses import dataclass +from pathlib import Path 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 +import numpy as np +from nerfstudio.process_data import colmap_utils, equirect_utils, process_data_utils +from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import ( + ColmapConverterToNerfstudioDataset, +) +from nerfstudio.utils.io import load_from_json from nerfstudio.utils.rich_utils import CONSOLE -from pathlib import Path - @dataclass class LanternImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): @@ -35,26 +44,33 @@ class LanternImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): 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.""" + # left_colmap_baseline_dir: Optional[Path] = None + # """Path to the left camera baseline directory.""" + # left_colmap_baseline_mask_dir: Optional[Path] = None + # """Path to the left camera baseline mask directory.""" + # right_colmap_baseline_dir: Optional[Path] = None + # """Path to the right camera baseline directory.""" + # right_colmap_baseline_mask_dir: Optional[Path] = None + # """Path to the right camera baseline mask directory.""" +# + # left_e1_dir: Optional[Path] = None + # """Path to the left camera e1 directory (well-exposed).""" + # left_e1_mask_dir: Optional[Path] = None + # """Path to the left camera e1 mask directory (well-exposed).""" + # right_e2_dir: Optional[Path] = None + # """Path to the right camera e2 directory (fast-exposed).""" + # right_e2_mask_dir: Optional[Path] = None + # """Path to the right camera e2 mask directory (fast-exposed).""" exposure1: float = 1.0 - """Path to the E2 mask directory.""" + """Well-exposed exposure value (left camera).""" exposure2: float = 0.009 - """Path to the E2 mask directory.""" + """Fast-exposed exposure value (right camera).""" def main(self) -> None: """Process images into a nerfstudio dataset.""" - + assert self.skip_image_processing, "Image processing not supported for now" self.is_HDR = True require_cameras_exist = False @@ -67,113 +83,193 @@ class LanternImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): image_rename_map: Optional[dict[str, str]] = None + sequence_names = ['left_colmap_baseline', 'right_colmap_baseline', 'left_e1', 'right_e2'] + sequence_mask_names = ['left_colmap_baseline_mask', 'right_colmap_baseline_mask', 'left_e1_mask', 'right_e2_mask'] + perspective_dirs = [] + perspective_mask_dirs = [] + # 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) + pers_size = equirect_utils.compute_resolution_from_equirect(self.data / sequence_names[0], 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, - ) + + for sequence_name, sequence_mask_name in zip(sequence_names, sequence_mask_names): + # check if mask exists + if not (self.data / sequence_mask_name).exists(): + mask_path = None + else: + mask_path = self.data / sequence_mask_name + output_dir, output_mask_dir, _ = equirect_utils.generate_planar_projections_from_equirectangular( + self.data / sequence_name, pers_size, self.images_per_equirect, mask_path, crop_factor=self.crop_factor + ) + perspective_dirs.append(output_dir) + CONSOLE.log(f"Generated {len(os.listdir(output_dir))} images in {output_dir}") + perspective_mask_dirs.append(output_mask_dir) + self.camera_type = "perspective" + else: + raise ValueError("Camera type must be equirectangular") 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 - ) + #TODO + image_rename_map_paths = {} + for i in range(len(perspective_dirs)): + image_rename_map_paths.update(process_data_utils.copy_images( + perspective_dirs[i], + 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: + raise NotImplementedError + 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) 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)) + num_frames = len(process_data_utils.list_images(self.data / sequence_names[0])) 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 + + with tempfile.TemporaryDirectory() as tmp_dir: + # copy files for colmap in tmp_dir + for i in range(len(perspective_dirs)): + # avoid right_e2 (fast-exposed), colmap doesn't work well on it + if i == 3: + continue + + shutil.copytree(perspective_dirs[i], os.path.join(tmp_dir, sequence_names[i], 'planar_projections')) + # TODO: mask + + # Run COLMAP + if not self.skip_colmap: + require_cameras_exist = True + self.absolute_colmap_path.mkdir(parents=True, exist_ok=True) + colmap_utils.run_colmap( + image_dir=Path(tmp_dir), + colmap_dir=self.absolute_colmap_path, + camera_model=process_data_utils.CAMERA_MODELS[self.camera_type], + camera_mask_path=None, #TODO + gpu=self.gpu, + verbose=self.verbose, + matching_method=self.matching_method, + colmap_cmd=self.colmap_cmd, + ) + # 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 + incomplete_transforms_file = self.output_dir / "transforms_incomplete.json" - 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'}).") + if not self.skip_colmap: + 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( + num_frames, + # todo: mask + ) - 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 - ) + shutil.move(str(self.output_dir / 'transforms.json'), incomplete_transforms_file) + CONSOLE.log(f"Saved initial transforms to {incomplete_transforms_file}") + original_transforms = load_from_json(incomplete_transforms_file) + + frames = original_transforms["frames"] + sorted_frames = sorted(frames, key=lambda x: x["file_path"]) + + # frames are splitted into 4 groups: left_calib, right_calib, left_e1, left_e2 + # the two first groups are used to infer the basis_change between left and right + # that basis_change is then applied to the third group (left_e1) to infer the fourth (left_e2) + + grouped_frames = {sequence_name: dict() for sequence_name in ['left_colmap_baseline', 'right_colmap_baseline', 'left_e1']} + for frame in sorted_frames: + # use regex + matches = re.match(r'images/(\w+)/.*_(\d+)_(\d+).png', frame["file_path"]) + if matches: + sequence_name = matches.group(1) + frame_idx = int(matches.group(2)) + crop_direction = int(matches.group(3)) + if crop_direction not in grouped_frames[sequence_name].keys(): + grouped_frames[sequence_name][crop_direction] = [] + grouped_frames[sequence_name][crop_direction].append(frame) + else: + # warning + print(f'Warning: {frame["file_path"]} does not match regex') + + + # group 3 is empty, we want to infer its transform + grouped_frames['right_e2'] = dict() + for crop_direction in range(self.images_per_equirect): + group_0_tranform = np.array(grouped_frames['left_colmap_baseline'][crop_direction][0]['transform_matrix']) + group_1_tranform = np.array(grouped_frames['right_colmap_baseline'][crop_direction][0]['transform_matrix']) + + # TODO: check which one is right + basis_change = np.linalg.inv(group_0_tranform) @ group_1_tranform + # basis_change = np.linalg.inv(group_1_tranform) @ group_0_tranform + + # sometimes, colmap omits frames, so if a frame doesn't have the left_e1 pose, we can't infer the right_e2 pose + if grouped_frames['left_e1'].get(crop_direction): + for frame in grouped_frames['left_e1'][crop_direction]: + + # apply group_3_tranform to frame, put it in group 3 + group_3_frame_transform = (np.array(frame['transform_matrix']) @ basis_change).tolist() + if crop_direction not in grouped_frames['right_e2'].keys(): + grouped_frames['right_e2'][crop_direction] = [] + grouped_frames['right_e2'][crop_direction].append({ + 'file_path': frame['file_path'].replace('left_e1', 'right_e2').replace('lhs', 'rhs'), # TODO validate file naming convention + 'transform_matrix': group_3_frame_transform + }) + else: + # warning + print(f'Warning: no left_e1 for crop_direction {crop_direction}') + + # write to text file grouped_frames + with open(self.output_dir / 'grouped_frames.json', 'w') as f: + json.dump(grouped_frames, f, indent=4) + + # generate new transforms.json + new_transforms = deepcopy(original_transforms) + new_transforms['frames'] = [] + for sequence_name in ['left_e1', 'right_e2']: + for i, crop_direction in enumerate(range(self.images_per_equirect)): + new_transforms['frames'] += grouped_frames[sequence_name][crop_direction] + + with open(self.output_dir / 'transforms.json', "w", encoding="UTF-8") as file: + json.dump(new_transforms, file, indent=4) + + # for HDR-Nerfacto, output also the exposures + exposures_content = {} + for frame in new_transforms['frames']: + if 'right_e2' in frame['file_path']: + exposures_content[frame['file_path']] = self.exposure2 + else: + exposures_content[frame['file_path']] = self.exposure1 + + with open(self.output_dir / 'exposures.json', "w", encoding="UTF-8") as file: + json.dump(exposures_content, file, indent=4) CONSOLE.log("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:")