diff --git a/lantern_scripts/deal_with_openSFM.py b/lantern_scripts/deal_with_openSFM.py new file mode 100644 index 000000000..a07531b5c --- /dev/null +++ b/lantern_scripts/deal_with_openSFM.py @@ -0,0 +1,41 @@ +import numpy as np +import argparse +import json + + +def read_json_file(path): + + with open(path) as f: + data = json.load(f) + camera_dict = {} + camera_dict_file = './transform.json' + + camera_dict["frames"] = [] + + for img_name in data[0]['shots']: + + rvec = np.array(tuple(map(float, data[0]['shots'][img_name]['rotation']))) + tvec = np.array(tuple(map(float, data[0]['shots'][img_name]['translation']))) + + from scipy.spatial.transform import Rotation as R + rot = R.from_rotvec(rvec).as_matrix() + W2C = np.eye(4) + W2C[:3, :3] = rot + W2C[:3, 3] = np.array(tvec) + + camera_dict["frames"].append( + { + "file_path": img_name, + "transform_matrix": W2C.tolist() + } + ) + + with open(camera_dict_file, 'w') as fp: + json.dump(camera_dict, fp, indent=2, sort_keys=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Read and write COLMAP binary and text models') + parser.add_argument('--input_model', help='path to input model folder') + args = parser.parse_args() + read_json_file(args.input_model) diff --git a/nerfstudio/process_data/openSFM_to_nerfstudio.py b/nerfstudio/process_data/openSFM_to_nerfstudio.py new file mode 100644 index 000000000..3f397856c --- /dev/null +++ b/nerfstudio/process_data/openSFM_to_nerfstudio.py @@ -0,0 +1,335 @@ +# 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 panorama sequence to a nerfstudio compatible dataset.""" + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Literal, Tuple +import os +import numpy as np +from numpy.linalg import inv +import cv2 +from tqdm import tqdm + + +from nerfstudio.process_data import equirect_utils, process_data_utils +from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import ( + BaseConverterToNerfstudioDataset, +) +from rich.progress import ( + BarColumn, + Progress, + TaskProgressColumn, + TextColumn, + TimeRemainingColumn, +) +from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn +from nerfstudio.utils import io + + +def apply_correction(coefs, img, CORRECTION_CURVE_TYPE) -> np.ndarray: + # check that img is in range 0-255 + assert img.dtype == np.uint8 + assert CORRECTION_CURVE_TYPE in ['gamma'] + img = ((img.astype(np.float32) / 255.0) ** np.array(coefs).T) # * 255.0 + # img = ((img.astype(np.float32) / 255.0) ** coefs[None, None, :]) * 255.0 + return img + + +def opensfm_to_opengl(shot_data): + rvec = np.array(tuple(map(float, shot_data['rotation']))) + tvec = np.array(tuple(map(float, shot_data['translation']))) + + # Converting opencv to opengl coordinates + from scipy.spatial.transform import Rotation as R + rot = R.from_rotvec(rvec).as_matrix() + + W2C = np.eye(4) + W2C[:3, :3] = rot + W2C[:3, 3] = np.array(tvec) + + C2W = inv(W2C) + C2W[0:3, 1:3] *= -1 + C2W = C2W[np.array([1, 0, 2, 3]), :] + C2W[2, :] *= -1 + + return C2W, rvec, tvec + + +@dataclass +class OpenSFMToNeRFStudioDataset(BaseConverterToNerfstudioDataset): + """Process panorama data into a nerfstudio dataset. Skip the COLMAP + + This script does the following: + + 1. Crops panoramas to planar images, and restore corresponding perspective camera poses. + 1. Scales images to a specified size. + 2. Converts perpective camera poses into the nerfstudio format. + """ + + metadata: Path = "" + """Path the metadata of the panoramas sequence.""" + num_downscales: int = 0 + """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 + will downscale the images by 2x, 4x, and 8x.""" + max_dataset_size: int = 300 + """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. If -1, + use all images.""" + images_per_equirect: Literal[8, 14] = 8 + """Number of samples per image to take from each equirectangular image. + Used only when camera-type is equirectangular. + """ + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) + """Portion of the image to crop. All values should be in [0,1]. (top, bottom, left, right)""" + skip_image_processing: bool = False + """If True, skips copying and downscaling of images and only runs COLMAP if possible and enabled""" + is_HDR: bool = False + """If True, process the .exr files as HDR images.""" + use_mask: bool = False + """If True, process the .exr files with mask.""" + use_exposure: bool = False + """If True, using different exposures.""" + + exposure1: float = None # t = 1/125 + """Well-exposed exposure value (left camera).""" + exposure2: float = None # t = 1/25000 + """Fast-exposed exposure value (right camera).""" + capture_settings_file_name: str = "capture_settings.json" + + skip_linearization: bool = False + + skip_saturation_mask: bool = False + + skip_histograms: bool = False + + def main(self) -> None: + """Process images into a nerfstudio dataset.""" + + self.output_dir.mkdir(parents=True, exist_ok=True) + image_dir = self.output_dir / "images" + image_dir.mkdir(parents=True, exist_ok=True) + + mask_dir = self.output_dir / "masks" + if self.use_mask: + mask_dir.mkdir(parents=True, exist_ok=True) + + # check camera settings + with open(self.data / self.capture_settings_file_name) as f: + # read shutter_left, shutter_right, iso_left, iso_right + data = json.load(f) + self.white_balance_setting = data["left"]["white_balance"] + self.exposure1 = 1.0 + self.exposure2 = data["right"]["shutter_speed"] / data["left"]["shutter_speed"] + + summary_log = [] + + if not self.skip_image_processing: + openSFM_reconstruction = io.load_from_json(self.metadata) + frames_names = openSFM_reconstruction[0]['shots'].keys() + + # Computing the basis trandoem matrix + rot_list = [] + trn_list = [] + for fname in frames_names: + if fname.startswith("left_sfm"): + shot_data = openSFM_reconstruction[0]['shots'][fname] + left_C2W, left_rvec, left_tvec = opensfm_to_opengl(shot_data) + + right_fname = fname.replace("left_sfm", "right_sfm") + if right_fname in frames_names: + shot_data = openSFM_reconstruction[0]['shots'][right_fname] + right_C2W, right_rvec, right_tvec = opensfm_to_opengl(shot_data) + + basis_change = np.linalg.inv(left_C2W) @ right_C2W + + from scipy.spatial.transform import Rotation as R + r = R.from_matrix(basis_change[0:3, 0:3]) + rot = r.as_euler('zyx', degrees=True) + rot_list.append(r.as_quat()) + trn_list.append(basis_change[0:3, 3]) + else: + print("The right one is not registered for " + fname) + + basis_change = np.zeros((4, 4)) + from sksurgerycore.algorithms.averagequaternions import average_quaternions + avg_quat = average_quaternions(np.asarray(rot_list)) + r = R.from_quat(avg_quat) + basis_change[0:3, 0:3] = r.as_matrix() + basis_change[0:3, 3] = np.mean(np.asarray(trn_list), 0) + basis_change[3, 3] = 1.0 + print("basis_change: ", basis_change) + + # Convert OpenSFM coordinates to the NeRFStudio + camera_dict = {} + camera_dict["frames"] = [] + for fname in frames_names: + # if fname.startswith("left_sfm"): + if fname.startswith("left_e1"): + shot_data = openSFM_reconstruction[0]['shots'][fname] + left_C2W, _, _ = opensfm_to_opengl(shot_data) + camera_dict["frames"].append( + { + "file_path": fname.split('.')[0], + "transform_matrix": left_C2W.tolist() + }) + + right_fname = fname.replace("left_e1", "right_e2") + # right_fname = fname.replace("left_sfm", "right_sfm") + right_C2W = left_C2W @ basis_change + camera_dict["frames"].append( + { + "file_path": right_fname.split('.')[0], + "transform_matrix": right_C2W.tolist() + }) + + with open(self.output_dir / 'panoramic_transforms.json', 'w') as f: + json.dump(camera_dict, f, indent=4) + + pers_size = equirect_utils.compute_resolution_from_equirect(self.data / 'left_e1', self.images_per_equirect) + # pers_size = equirect_utils.compute_resolution_from_equirect(self.data / 'left_sfm', self.images_per_equirect) + + cropped_images_filename = [] + cropped_masks_filename = [] + metadata_dict_all = {} + # for sequence_name in ['left_sfm', 'right_sfm']: + for sequence_name in ['left_e1', 'right_e2']: + CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image") + out_dir = equirect_utils.generate_planar_projections_from_equirectangular_GT( + self.output_dir / 'panoramic_transforms.json', + self.data / sequence_name, pers_size, self.images_per_equirect, + crop_factor=self.crop_factor, + clip_output = False, + use_mask = self.use_mask, + prefix = sequence_name + ) + self.camera_type = "perspective" + metadata_dict = io.load_from_json(out_dir / "transforms.json") + metadata_dict_all[sequence_name] = metadata_dict + # Copy images to output directory + for frame in metadata_dict["frames"]: + cropped_images_filename.append(Path(frame["file_path"])) + if self.use_mask: + cropped_masks_filename.append(Path(frame["mask_path"])) + + # Copy images to output directory + copied_image_paths = process_data_utils.copy_images_list( + cropped_images_filename, image_dir=image_dir, verbose=self.verbose, num_downscales=self.num_downscales + ) + copied_image_paths = [Path("images/" + copied_image_path.name) for copied_image_path in copied_image_paths] + if self.use_mask: + copied_mask_paths = process_data_utils.copy_images_list( + cropped_masks_filename, image_dir=mask_dir, verbose=self.verbose, num_downscales=self.num_downscales + ) + copied_mask_paths = [Path("masks/" + copied_mask_path.name) for copied_mask_path in copied_mask_paths] + num_frames = len(copied_image_paths) + + summary_log.append(f"Used {num_frames} images out of {num_frames} total") + if self.max_dataset_size > 0: + summary_log.append( + "To change the size of the dataset add the argument [yellow]--max_dataset_size[/yellow] to " + f"larger than the current value ({self.max_dataset_size}), or -1 to use all images." + ) + + metadata_path = self.output_dir / "transforms.json" + metadata_dict = {} + for key in metadata_dict_all["left_e1"].keys() - 'frames': + metadata_dict[key] = metadata_dict_all["left_e1"][key] + metadata_dict["frames"] = metadata_dict_all['left_e1']['frames'] + metadata_dict_all['right_e2']['frames'] + # for key in metadata_dict_all["left_sfm"].keys() - 'frames': + # metadata_dict[key] = metadata_dict_all["left_sfm"][key] + # metadata_dict["frames"] = metadata_dict_all['left_sfm']['frames'] + metadata_dict_all['right_sfm']['frames'] + + for i, frame in enumerate(metadata_dict["frames"]): + frame["file_path"] = str(copied_image_paths[i]) + if self.use_mask: + frame["mask_path"] = str(copied_mask_paths[i]) + + with open(metadata_path, "w", encoding="utf-8") as f: + json.dump(metadata_dict, f, indent=4) + + transform_data = io.load_from_json(self.output_dir / "transforms.json") + + # linearization + if not self.skip_linearization: + + with open(f'linearize-ricoh-z10-wb-{self.white_balance_setting}.json') as f: + data = json.load(f) + CORRECTION_CURVE_TYPE = 'gamma' + + for frame in tqdm(transform_data["frames"], desc="Linearization: "): + img = cv2.imread(os.path.join(self.output_dir, frame['file_path'])) + # import pdb; pdb.set_trace() + img_corrected = apply_correction((data["b"], data["g"], data["r"]), img, CORRECTION_CURVE_TYPE) # + # img_corrected /= 255 + # img_corrected = img_corrected ** 2.2 + # import pdb; pdb.set_trace() + img_name = frame['file_path'].split('/')[-1] + exposed_img = np.float32(img_corrected) # * self.exposure1 if img_name.startswith('left_e1') else img_corrected / self.exposure2 + # exposed_img = img_corrected * self.exposure1 + cv2.imwrite(os.path.join(self.output_dir, frame['file_path'].replace('.png', '_linear.exr')), exposed_img) + # import pdb; pdb.set_trace() + + # saturation mask + if not self.skip_saturation_mask: + for frame in tqdm(transform_data["frames"], desc="Computing the Saturation mask: "): + img = cv2.imread(os.path.join(self.output_dir, frame['file_path'])) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = img.astype(np.float32) / 255.0 + # left e1 threshold: valid from 0 to 0.95 + # right e2 threshold: valid from 0.2 to 1 + # to have mask be =1, all 3 channels must be valid + img_name = frame['file_path'].split('/')[-1] + if img_name.startswith('left_e1'): + min_thresh = 0.0 + max_thresh = 0.95 + else: + min_thresh = 0.2 + max_thresh = 1.0 + mask = np.greater_equal(img, min_thresh) & np.less_equal(img, max_thresh) + mask = np.all(mask, axis=2) + mask = mask.astype(np.uint8) * 255 + cv2.imwrite(os.path.join(self.output_dir, frame['mask_path'].replace('.png', '_saturation_mask.png')), mask) + # add to frame + frame['saturation_mask_path'] = frame['mask_path'].replace('.png', '_saturation_mask.png') + + # change frame file_path to linear exr + for frame in transform_data['frames']: + frame['file_path'] = frame['file_path'].replace('.png', '_linear.exr') + + # for HDR-Nerfacto, output also the exposures + exposures_content = {} + for frame in transform_data['frames']: + if 'right_e2' in frame['file_path']: + exposures_content[frame['file_path']] = self.exposure2 + else: + exposures_content[frame['file_path']] = self.exposure1 + + # add the exposures to the transforms.json + for frame in transform_data['frames']: + frame['exposure'] = exposures_content[frame['file_path']] + + with open(self.output_dir / 'exposures.json', "w", encoding="UTF-8") as file: + json.dump(exposures_content, file, indent=4) + + with open(self.output_dir / 'transforms.json', "w", encoding="UTF-8") as file: + json.dump(transform_data, file, indent=4) + + + CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") + + for summary in summary_log: + CONSOLE.print(summary, justify="center") + CONSOLE.rule() \ No newline at end of file