diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 01f2db4e5..bde0e0d0b 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -31,6 +31,7 @@ from nerfactory.dataloaders.data_parsers import ( Friends, InstantNGP, Mipnerf360, + Nerfactory, Record3D, ) from nerfactory.models.base import Model @@ -169,6 +170,17 @@ class DataParserConfig(InstantiateConfig): _target: Type = DataParser +@dataclass +class NerfactoryDataParserConfig(DataParserConfig): + """Nerfactory dataset config""" + + _target: Type = Nerfactory + data_directory: Path = Path("data/ours/posterv2") + scale_factor: float = 1.0 + downscale_factor: int = 1 + scene_scale: float = 0.33 + + @dataclass class BlenderDataParserConfig(DataParserConfig): """Blender dataset config""" @@ -202,7 +214,7 @@ class MipNerf360DataParserConfig(DataParserConfig): @dataclass class InstantNGPDataParserConfig(DataParserConfig): - """Mipnerf 360 dataset config""" + """Instant-NGP dataset config""" _target: Type = InstantNGP data_directory: Path = Path("data/ours/posterv2") diff --git a/nerfactory/dataloaders/data_parsers.py b/nerfactory/dataloaders/data_parsers.py index 139ba8cef..36d6d6d44 100644 --- a/nerfactory/dataloaders/data_parsers.py +++ b/nerfactory/dataloaders/data_parsers.py @@ -30,17 +30,17 @@ from scipy.spatial.transform import Rotation from nerfactory.cameras import utils as camera_utils from nerfactory.cameras.cameras import Cameras, CameraType from nerfactory.configs import base as cfg -from nerfactory.dataloaders.colmap_utils import ( - read_cameras_binary, - read_images_binary, - read_pointsTD_binary, -) from nerfactory.dataloaders.structs import ( DatasetInputs, PointCloud, SceneBounds, Semantics, ) +from nerfactory.utils.colmap_utils import ( + read_cameras_binary, + read_images_binary, + read_pointsTD_binary, +) from nerfactory.utils.colors import get_color from nerfactory.utils.io import ( get_absolute_path, @@ -129,6 +129,83 @@ class DataParser: return dataset_inputs +@dataclass +class Nerfactory(DataParser): + """Nerfactory Dataset + + Args: + data_directory: Location of data + scale_factor: How much to scale the camera origins by. + downscale_factor: How much to downscale images. Defaults to 1. + scene_scale: How much to scale the scene. Defaults to 0.33 + """ + + config: cfg.InstantNGPDataParserConfig + + def _generate_dataset_inputs(self, split="train"): + + abs_dir = get_absolute_path(self.config.data_directory) + + meta = load_from_json(abs_dir / "transforms.json") + image_filenames = [] + poses = [] + num_skipped_image_filenames = 0 + for frame in meta["frames"]: + fname = abs_dir / Path(frame["file_path"]) + if not fname: + num_skipped_image_filenames += 1 + else: + image_filenames.append(fname) + poses.append(np.array(frame["transform_matrix"])) + if num_skipped_image_filenames >= 0: + logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split) + assert ( + len(image_filenames) != 0 + ), """ + No image files found. + You should check the file_paths in the transforms.json file to make sure they are correct. + """ + poses = np.array(poses).astype(np.float32) + poses[:3, 3] *= self.config.scene_scale + + camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform + + distortion_params = camera_utils.get_distortion_params( + k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"]) + ) + + # in x,y,z order + # assumes that the scene is centered at the origin + aabb_scale = 16 + scene_bounds = SceneBounds( + aabb=torch.tensor( + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 + ) + ) + + cameras = Cameras( + fx=float(meta["fl_x"]), + fy=float(meta["fl_y"]), + cx=float(meta["cx"]), + cy=float(meta["cy"]), + distortion_params=distortion_params, + height=int(meta["h"]), + width=int(meta["w"]), + camera_to_worlds=camera_to_world, + camera_type=CameraType.PERSPECTIVE, + ) + + cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) + + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + downscale_factor=self.config.downscale_factor, + cameras=cameras, + scene_bounds=scene_bounds, + ) + return dataset_inputs + + @dataclass class Blender(DataParser): """Blender Dataset @@ -238,20 +315,11 @@ class InstantNGP(DataParser): poses = np.array(poses).astype(np.float32) poses[:3, 3] *= self.config.scene_scale - fx = float(meta["fl_x"]) - fy = float(meta["fl_y"]) - cx = float(meta["cx"]) - cy = float(meta["cy"]) - k1 = float(meta["k1"]) - k2 = float(meta["k2"]) - p1 = float(meta["p1"]) - p2 = float(meta["p2"]) - image_height = int(meta["h"]) - image_width = int(meta["w"]) - camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform - distortion_params = camera_utils.get_distortion_params(k1=k1, k2=k2, p1=p1, p2=p2) + distortion_params = camera_utils.get_distortion_params( + k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"]) + ) # in x,y,z order # assumes that the scene is centered at the origin @@ -263,13 +331,13 @@ class InstantNGP(DataParser): ) cameras = Cameras( - fx=fx, - fy=fy, - cx=cx, - cy=cy, + fx=float(meta["fl_x"]), + fy=float(meta["fl_y"]), + cx=float(meta["cx"]), + cy=float(meta["cy"]), distortion_params=distortion_params, - height=image_height, - width=image_width, + height=int(meta["h"]), + width=int(meta["w"]), camera_to_worlds=camera_to_world, camera_type=CameraType.PERSPECTIVE, ) diff --git a/nerfactory/dataloaders/colmap_utils.py b/nerfactory/utils/colmap_utils.py similarity index 63% rename from nerfactory/dataloaders/colmap_utils.py rename to nerfactory/utils/colmap_utils.py index 6cc8b2e6d..af7853685 100644 --- a/nerfactory/dataloaders/colmap_utils.py +++ b/nerfactory/utils/colmap_utils.py @@ -1,6 +1,6 @@ # pylint: disable-all """ -Here we have code taken from COLMAP for parsing data in the COLMAP format. +Here we have modified code taken from COLMAP for parsing data in the COLMAP format. Original file at https://github.com/colmap/colmap/blob/1a4d0bad2e90aa65ce997c9d1779518eaed998d5/scripts/python/read_write_model.py. """ @@ -84,24 +84,6 @@ def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): return struct.unpack(endian_character + format_char_sequence, data) -def write_next_bytes(fid, data, format_char_sequence, endian_character="<"): - """pack and write to a binary file. - - Args: - fid: - data: data to send, if multiple elements are sent at the same time, - they should be encapsuled either in a list or a tuple - format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. - should be the same length as the data list or tuple - endian_character: Any of {@, =, <, >, !} - """ - if isinstance(data, (list, tuple)): - bytes = struct.pack(endian_character + format_char_sequence, *data) - else: - bytes = struct.pack(endian_character + format_char_sequence, data) - fid.write(bytes) - - def read_cameras_text(path): """ see: src/base/reconstruction.cc @@ -151,42 +133,6 @@ def read_cameras_binary(path_to_model_file): return cameras -def write_cameras_text(cameras, path): - """ - see: src/base/reconstruction.cc - void Reconstruction::WriteCamerasText(const std::string& path) - void Reconstruction::ReadCamerasText(const std::string& path) - """ - HEADER = ( - "# Camera list with one line of data per camera:\n" - + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" - + f"# Number of cameras: {len(cameras)}\n" - ) - with open(path, "w") as fid: - fid.write(HEADER) - for _, cam in cameras.items(): - to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params] - line = " ".join([str(elem) for elem in to_write]) - fid.write(line + "\n") - - -def write_cameras_binary(cameras, path_to_model_file): - """ - see: src/base/reconstruction.cc - void Reconstruction::WriteCamerasBinary(const std::string& path) - void Reconstruction::ReadCamerasBinary(const std::string& path) - """ - with open(path_to_model_file, "wb") as fid: - write_next_bytes(fid, len(cameras), "Q") - for _, cam in cameras.items(): - model_id = CAMERA_MODEL_NAMES[cam.model].model_id - camera_properties = [cam.id, model_id, cam.width, cam.height] - write_next_bytes(fid, camera_properties, "iiQQ") - for p in cam.params: - write_next_bytes(fid, float(p), "d") - return cameras - - def read_images_text(path): """ see: src/base/reconstruction.cc @@ -258,57 +204,6 @@ def read_images_binary(path_to_model_file): return images -def write_images_text(images, path): - """ - see: src/base/reconstruction.cc - void Reconstruction::ReadImagesText(const std::string& path) - void Reconstruction::WriteImagesText(const std::string& path) - """ - if len(images) == 0: - mean_observations = 0 - else: - mean_observations = sum((len(img.point3D_ids) for _, img in images.items())) / len(images) - HEADER = ( - "# Image list with two lines of data per image:\n" - + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" - + "# POINTS2D[] as (X, Y, POINT3D_ID)\n" - + f"# Number of images: {len(images)}, mean observations per image: {mean_observations}\n" - ) - - with open(path, "w") as fid: - fid.write(HEADER) - for _, img in images.items(): - image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name] - first_line = " ".join(map(str, image_header)) - fid.write(first_line + "\n") - - points_strings = [] - for xy, point3D_id in zip(img.xys, img.point3D_ids): - points_strings.append(" ".join(map(str, [*xy, point3D_id]))) - fid.write(" ".join(points_strings) + "\n") - - -def write_images_binary(images, path_to_model_file): - """ - see: src/base/reconstruction.cc - void Reconstruction::ReadImagesBinary(const std::string& path) - void Reconstruction::WriteImagesBinary(const std::string& path) - """ - with open(path_to_model_file, "wb") as fid: - write_next_bytes(fid, len(images), "Q") - for _, img in images.items(): - write_next_bytes(fid, img.id, "i") - write_next_bytes(fid, img.qvec.tolist(), "dddd") - write_next_bytes(fid, img.tvec.tolist(), "ddd") - write_next_bytes(fid, img.camera_id, "i") - for char in img.name: - write_next_bytes(fid, char.encode("utf-8"), "c") - write_next_bytes(fid, b"\x00", "c") - write_next_bytes(fid, len(img.point3D_ids), "Q") - for xy, p3d_id in zip(img.xys, img.point3D_ids): - write_next_bytes(fid, [*xy, p3d_id], "ddq") - - def read_pointsTD_text(path): """ see: src/base/reconstruction.cc @@ -361,52 +256,6 @@ def read_pointsTD_binary(path_to_model_file): return points3D -def write_pointsTD_text(points3D, path): - """ - see: src/base/reconstruction.cc - void Reconstruction::ReadPoints3DText(const std::string& path) - void Reconstruction::WritePoints3DText(const std::string& path) - """ - if len(points3D) == 0: - mean_track_length = 0 - else: - mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items())) / len(points3D) - HEADER = ( - "# 3D point list with one line of data per point:\n" - + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n" - + f"# Number of points: {len(points3D)}, mean track length: {mean_track_length}\n" - ) - - with open(path, "w") as fid: - fid.write(HEADER) - for _, pt in points3D.items(): - point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error] - fid.write(" ".join(map(str, point_header)) + " ") - track_strings = [] - for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs): - track_strings.append(" ".join(map(str, [image_id, point2D]))) - fid.write(" ".join(track_strings) + "\n") - - -def write_pointsTD_binary(points3D, path_to_model_file): - """ - see: src/base/reconstruction.cc - void Reconstruction::ReadPoints3DBinary(const std::string& path) - void Reconstruction::WritePoints3DBinary(const std::string& path) - """ - with open(path_to_model_file, "wb") as fid: - write_next_bytes(fid, len(points3D), "Q") - for _, pt in points3D.items(): - write_next_bytes(fid, pt.id, "Q") - write_next_bytes(fid, pt.xyz.tolist(), "ddd") - write_next_bytes(fid, pt.rgb.tolist(), "BBB") - write_next_bytes(fid, pt.error, "d") - track_length = pt.image_ids.shape[0] - write_next_bytes(fid, track_length, "Q") - for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs): - write_next_bytes(fid, [image_id, point2D_id], "ii") - - def detect_model_format(path, ext): if ( os.path.isfile(os.path.join(path, "cameras" + ext)) @@ -440,18 +289,6 @@ def read_model(path, ext=""): return cameras, images, points3D -def write_model(cameras, images, points3D, path, ext=".bin"): - if ext == ".txt": - write_cameras_text(cameras, os.path.join(path, "cameras" + ext)) - write_images_text(images, os.path.join(path, "images" + ext)) - write_pointsTD_text(points3D, os.path.join(path, "points3D") + ext) - else: - write_cameras_binary(cameras, os.path.join(path, "cameras" + ext)) - write_images_binary(images, os.path.join(path, "images" + ext)) - write_pointsTD_binary(points3D, os.path.join(path, "points3D") + ext) - return cameras, images, points3D - - def qvec2rotmat(qvec): return np.array( [ @@ -492,25 +329,3 @@ def rotmat2qvec(R): if qvec[0] < 0: qvec *= -1 return qvec - - -def main(): - parser = argparse.ArgumentParser(description="Read and write COLMAP binary and text models") - parser.add_argument("--input_model", help="path to input model folder") - parser.add_argument("--input_format", choices=[".bin", ".txt"], help="input model format", default="") - parser.add_argument("--output_model", help="path to output model folder") - parser.add_argument("--output_format", choices=[".bin", ".txt"], help="outut model format", default=".txt") - args = parser.parse_args() - - cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format) - - print("num_cameras:", len(cameras)) - print("num_images:", len(images)) - print("num_points3D:", len(points3D)) - - if args.output_model is not None: - write_model(cameras, images, points3D, path=args.output_model, ext=args.output_format) - - -if __name__ == "__main__": - main() diff --git a/scripts/process_data.py b/scripts/process_data.py new file mode 100644 index 000000000..95998f0bd --- /dev/null +++ b/scripts/process_data.py @@ -0,0 +1,294 @@ +"""Processes a video or image sequence to a nerfactory compatible dataset.""" + +import json +import subprocess +import sys +from enum import Enum +from pathlib import Path +from shutil import which +from typing import Literal, Optional + +import dcargs +import numpy as np + +from nerfactory.utils import colmap_utils + + +class CameraType(Enum): + """Enum for camera types.""" + + OPENCV = "OPENCV" + OPENCV_FISHEYE = "OPENCV_FISHEYE" + + +CAMERA_MODELS = { + "perspective": CameraType.OPENCV, + "fisheye": CameraType.OPENCV_FISHEYE, +} + + +def _print_red(text): + """Prints text in red.""" + print(f"\033[91m{text}\033[00m") + + +def _print_green(text): + """Prints text in green.""" + print(f"\033[92m{text}\033[00m") + + +def _print_section_header(text): + """Prints a section header.""" + _print_green("=" * 80) + _print_green(text) + _print_green("=" * 80) + + +def check_ffmpeg_installed(): + """Checks if ffmpeg is installed.""" + ffmpeg_path = which("ffmpeg") + if ffmpeg_path is None: + _print_red("Could not find ffmpeg. Please install ffmpeg.") + print("See https://ffmpeg.org/download.html for installation instructions.") + print("ffmpeg is only necissary if using videos as input.") + sys.exit(1) + + +def check_colmap_installed(): + """Checks if colmap is installed.""" + colmap_path = which("colmap") + if colmap_path is None: + _print_red("Could not find COLMAP. Please install COLMAP.") + print("See https://colmap.github.io/install.html for installation instructions.") + sys.exit(1) + + +def get_colmap_version(default_version=3.8) -> float: + """Returns the version of COLMAP. + This code assumes that colmap returns a version string of the form + "COLMAP 3.8 ..." which may not be true for all versions of COLMAP. + + Args: + default_version: Default version to return if COLMAP version can't be determined. + Returns: + The version of COLMAP. + """ + output = run_command("colmap", return_output=True) + assert output is not None + for line in output.split("\n"): + if line.startswith("COLMAP"): + return float(line.split(" ")[1]) + _print_red(f"Could not find COLMAP version. Using default {default_version}") + return default_version + + +def run_command(cmd, return_output=False) -> Optional[str]: + """Runs a command and returns the output. + + Args: + cmd: Command to run. + return_output: If True, returns the output of the command. + Returns: + The output of the command if return_output is True, otherwise None. + """ + if return_output: + rc, output = subprocess.getstatusoutput(cmd) + if rc != 0: + _print_red(f"Error running command: {cmd}") + _print_red(f"Got: {output}") + sys.exit(1) + return output + rc = subprocess.run(cmd, shell=True, check=True).returncode + if rc != 0: + _print_red(f"Error running command: {cmd}") + sys.exit(1) + return None + + +def convert_video_to_images(video_path: Path, image_dir: Path, num_frames_target: int): + """Converts a video into a sequence of images. + + Args: + video_path: Path to the video. + output_dir: Path to the output directory. + fps: Frames per second. + """ + cmd = f"ffprobe -v error -select_streams v:0 -count_packets \ + -show_entries stream=nb_read_packets -of csv=p=0 {video_path}" + output = run_command(cmd, return_output=True) + assert output is not None + num_frames = int(output) + print("Number of frames in video:", num_frames) + + out_filename = image_dir / "frame_%05d.png" + + ffmpeg_cmd = f"ffmpeg -i {video_path}" + + spacing = num_frames // num_frames_target + + if spacing > 1: + ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB -r 1" + else: + _print_red("Can't satify requested number of frames. Extracting all frames.") + + ffmpeg_cmd += f" {out_filename}" + + run_command(ffmpeg_cmd) + + +def run_colmap(image_dir: Path, colmap_dir: Path, camera_model: CameraType, gpu: bool = True) -> None: + """Runs COLMAP on the images. + + Args: + image_dir: Path to the directory containing the images. + colmap_dir: Path to the output directory. + camera_model: Camera model to use. + gpu: If True, use GPU. + """ + + colmap_version = get_colmap_version() + + # Feature extraction + feature_extractor_cmd = [ + "colmap feature_extractor", + f"--database_path {colmap_dir}/database.db", + f"--image_path {image_dir}", + "--ImageReader.single_camera 1", + f"--ImageReader.camera_model {camera_model.value}", + f"--SiftExtraction.use_gpu {int(gpu)}", + ] + feature_extractor_cmd = " ".join(feature_extractor_cmd) + run_command(feature_extractor_cmd) + + # Feature matching + feature_matcher_cmd = [ + "colmap exhaustive_matcher", + f"--database_path {colmap_dir}/database.db", + f"--SiftMatching.use_gpu {int(gpu)}", + ] + feature_matcher_cmd = " ".join(feature_matcher_cmd) + run_command(feature_matcher_cmd) + + # Bundle adjustment + sparse_dir = colmap_dir / "sparse" + sparse_dir.mkdir(parents=True, exist_ok=True) + bundle_adjuster_cmd = [ + "colmap mapper", + f"--database_path {colmap_dir}/database.db", + f"--image_path {image_dir}", + f"--output_path {sparse_dir}", + ] + if colmap_version >= 3.7: + bundle_adjuster_cmd.append("--Mapper.ba_global_function_tolerance 1e-6") + + bundle_adjuster_cmd = " ".join(bundle_adjuster_cmd) + run_command(bundle_adjuster_cmd) + + +def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path) -> None: + """Converts COLMAP's cameras.bin and images.bin to a JSON file. + + Args: + cameras_path: Path to the cameras.bin file. + images_path: Path to the images.bin file. + output_dir: Path to the output directory. + """ + + cameras = colmap_utils.read_cameras_binary(cameras_path) + images = colmap_utils.read_images_binary(images_path) + + # Only supports one camera + camera_params = cameras[1].params + + frames = [] + for _, im_data in images.items(): + rotation = colmap_utils.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 to ours + c2w[0:3, 1:3] *= -1 + c2w = c2w[np.array([1, 0, 2, 3]), :] + c2w[2, :] *= -1 + + name = f"./images/{im_data.name}" + + frame = { + "file_path": name, + "transform_matrix": c2w.tolist(), + } + frames.append(frame) + + out = { + "fl_x": float(camera_params[0]), + "fl_y": float(camera_params[1]), + "k1": float(camera_params[4]), + "k2": float(camera_params[5]), + "p1": float(camera_params[6]), + "p2": float(camera_params[7]), + "cx": float(camera_params[2]), + "cy": float(camera_params[3]), + "w": cameras[1].width, + "h": cameras[1].height, + "frames": frames, + } + + with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: + json.dump(out, f, indent=4) + + +def main( + data: Path, + output_dir: Path, + num_frames_target: int, + camera_type: Literal["perspective", "fisheye"] = "perspective", + gpu: bool = True, +): + """Process images or videos into a Nerfactory dataset. + + This script does the following: + 1) Converts video into images (if video is provided). + 2) Scales images to a specified size. + 3) Calculates and stores the sharpness of each image. + 4) Calculates the camera poses for each image using `COLMAP `_. + + + Args: + data: Path the data, either a video file or a directory of images. + output_dir: Path to the output directory. + num_frames: Target number of frames to use for the dataset, results may not be exact. + camera_type: Camera model to use. + gpu: If True, use GPU. + """ + + print(data) + + check_ffmpeg_installed() + check_colmap_installed() + + output_dir.mkdir(parents=True, exist_ok=True) + image_dir = output_dir / "images" + image_dir.mkdir(parents=True, exist_ok=True) + + if data.is_file(): + _print_section_header("Converting video into images.") + convert_video_to_images(data, image_dir=image_dir, num_frames_target=num_frames_target) + + _print_section_header("Running Colmap.") + colmap_dir = output_dir / "colmap" + colmap_dir.mkdir(parents=True, exist_ok=True) + + camera_model = CAMERA_MODELS[camera_type] + run_colmap(image_dir=image_dir, colmap_dir=colmap_dir, camera_model=camera_model, gpu=gpu) + + _print_section_header("Saving output to JSON.") + colmap_to_json( + cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", + images_path=colmap_dir / "sparse" / "0" / "images.bin", + output_dir=output_dir, + ) + + +if __name__ == "__main__": + dcargs.cli(main)