diff --git a/.gitignore b/.gitignore index b8382ed16..8364d16fd 100644 --- a/.gitignore +++ b/.gitignore @@ -165,8 +165,10 @@ outputs/ # Data data/ +!configs/data # Misc old/ temp* -.nfs* \ No newline at end of file +.nfs* +external/ \ No newline at end of file diff --git a/README.md b/README.md index 666ec8db1..11eea60de 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # MattPort - # Quickstart + ## Installation: Setup the environment ``` @@ -34,13 +34,22 @@ Download the original NeRF dataset and put it in the following format. ## Training a model ``` -# Run with the default config +# Run with default config python scripts/run_train.py + # Run with config changes python scripts/run_train.py machine_config.num_gpus=1 + +# Run with different datasets +python scripts/run_train.py data/dataset=blender_lego +python scripts/run_train.py data/dataset=friends_TBBT-big_living_room + +# Run with different datasets and config changes +python scripts/run_train.py data/dataset=friends_TBBT-big_living_room graph.network.far_plane=14 ``` # Logging and profiling features + We provide many logging functionalities for timing and/or tracking losses during training. All of these loggers are configurable via `configs/logging.yml` 1. **Writer**: Logs losses and generated images during training to a specified output stream. Specify the type of writer (Tensorboard, local directory, Weights and Biases), and how often to log in the config. @@ -49,7 +58,6 @@ We provide many logging functionalities for timing and/or tracking losses during 3. **Profiler**: Computes the average total time of execution for any function with the `@profiler.time_function` decorator. Prints out the full profile at the termination of training or the program. - # Setting up Jupyter ``` @@ -59,7 +67,7 @@ bash environments/run_jupyter.sh # Tooling -We use [autoenv](https://github.com/hyperupcall/autoenv) to make setting up the environment and environment variables easier. This will run the `.envrc` file upon entering the `/path/to/mattport` folder. It will also remove the set environment parameters upon leaving. +One can use [autoenv](https://github.com/hyperupcall/autoenv) to make setting up the environment and environment variables easier. This will run the `.envrc` file upon entering the `/path/to/mattport` folder. It will also remove the environment parameters upon leaving. ``` # Install direnv. @@ -71,7 +79,4 @@ eval "$(direnv hook bash)" # Populate your .envrc with commands you want to run. Then, run the following to allow updates. cd /path/to/mattport direnv allow . -``` - - - \ No newline at end of file +``` \ No newline at end of file diff --git a/configs/data/dataset/blender_lego.yaml b/configs/data/dataset/blender_lego.yaml new file mode 100644 index 000000000..6b80a1cff --- /dev/null +++ b/configs/data/dataset/blender_lego.yaml @@ -0,0 +1,4 @@ +data_directory: data/blender/lego +dataset_type: blender +downscale_factor: 4 +alpha_color: white \ No newline at end of file diff --git a/configs/data/dataset/friends_TBBT-big_living_room.yaml b/configs/data/dataset/friends_TBBT-big_living_room.yaml new file mode 100644 index 000000000..152d90958 --- /dev/null +++ b/configs/data/dataset/friends_TBBT-big_living_room.yaml @@ -0,0 +1,3 @@ +data_directory: data/friends/TBBT-big_living_room +dataset_type: friends +downscale_factor: 8 \ No newline at end of file diff --git a/configs/data/default_data.yaml b/configs/data/default_data.yaml index 4f0e661e2..e67475ec6 100644 --- a/configs/data/default_data.yaml +++ b/configs/data/default_data.yaml @@ -1,9 +1,6 @@ -# dataset options -dataset: - data_directory: data/blender/lego - dataset_type: blender - downscale_factor: 4 - alpha_color: white +defaults: + # dataset options + - dataset: blender_lego # validation options validation_image_indices: [0, 10] # Which images to evaluate on. diff --git a/configs/default.yaml b/configs/default.yaml index cd23f0aca..192a5ac79 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -5,6 +5,7 @@ defaults: - logging: default_logging.yaml - override hydra/job_logging: colorlog - override hydra/hydra_logging: colorlog + - _self_ hydra: run: diff --git a/configs/graph/default_graph.yaml b/configs/graph/default_graph.yaml index 9fd43adf5..5cb95c6d1 100644 --- a/configs/graph/default_graph.yaml +++ b/configs/graph/default_graph.yaml @@ -1,8 +1,8 @@ # saving model specifications model_dir: /tmp/mattport_models/ steps_per_save: null -steps_per_test: 100 -max_num_iterations: 1000000 +steps_per_test: 1000 +max_num_iterations: 100000 # pretrained loading options resume_train: @@ -15,14 +15,14 @@ network: near_plane: 1.0 far_plane: 6.0 num_coarse_samples: 64 - num_importance_samples: 64 + num_importance_samples: 128 # optimizer options for network param_groups: fields: optimizer: type: Adam - lr: 1e-4 + lr: 5e-4 scheduler: type: StepLR step_size: 30 # Period of learning rate decay. diff --git a/configs/test/test_dryrun.yaml b/configs/tests/test_dryrun.yaml similarity index 92% rename from configs/test/test_dryrun.yaml rename to configs/tests/test_dryrun.yaml index a19f72df2..df121c80e 100644 --- a/configs/test/test_dryrun.yaml +++ b/configs/tests/test_dryrun.yaml @@ -9,6 +9,7 @@ machine_config: data: dataset: data_directory: data/blender/lego_test + downscale_factor: 8 validation_image_indices: [0] dataloader: num_images_to_sample_from: 1 diff --git a/environment/requirements.txt b/environment/requirements.txt index e5050a742..d26e36758 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -1,4 +1,4 @@ -black +black==22.3.0 hydra-colorlog hydra-core imageio==2.16.1 diff --git a/mattport/nerf/dataset/blender.py b/mattport/nerf/dataset/blender.py new file mode 100644 index 000000000..c8e0e9d12 --- /dev/null +++ b/mattport/nerf/dataset/blender.py @@ -0,0 +1,64 @@ +""" +Code to handle loading blender datasets. +""" + +import os +from typing import Optional + +import imageio +import numpy as np +import torch +from torchtyping import TensorType + +from mattport.nerf.dataset.structs import DatasetInputs +from mattport.utils.io import load_from_json + + +def load_blender_data( + basedir: str, downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None, split: str = "train" +) -> DatasetInputs: + """Processes the a blender dataset directory. + Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37. + + Args: + data_directory (str): Location of data + downscale_factor (int, optional): How much to downscale images. Defaults to 1.0. + alpha_color (TensorType[3], optional): Sets transparent regions to specified color, otherwise black. + split (str, optional): Which dataset split to generate. + + Returns: + DatasetInputs + """ + # pylint: disable=unused-argument + # pylint: disable-msg=too-many-locals + meta = load_from_json(os.path.join(basedir, f"transforms_{split}.json")) + image_filenames = [] + poses = [] + for frame in meta["frames"]: + fname = os.path.join(basedir, frame["file_path"].replace("./", "") + ".png") + image_filenames.append(fname) + poses.append(np.array(frame["transform_matrix"])) + poses = np.array(poses).astype(np.float32) + + img_0 = imageio.imread(image_filenames[0]) + image_height, image_width = img_0.shape[:2] + camera_angle_x = float(meta["camera_angle_x"]) + focal_length = 0.5 * image_width / np.tan(0.5 * camera_angle_x) + + cx = image_width / 2.0 + cy = image_height / 2.0 + camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform + num_cameras = len(image_filenames) + num_intrinsics_params = 3 + intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) + intrinsics *= torch.tensor([cx, cy, focal_length]) + + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + downscale_factor=downscale_factor, + alpha_color=alpha_color, + intrinsics=intrinsics * 1.0 / downscale_factor, # downscaling the intrinsics here + camera_to_world=camera_to_world, + ) + + return dataset_inputs diff --git a/mattport/nerf/dataset/colmap_utils.py b/mattport/nerf/dataset/colmap_utils.py new file mode 100644 index 000000000..d4db0304a --- /dev/null +++ b/mattport/nerf/dataset/colmap_utils.py @@ -0,0 +1,513 @@ +# pylint: disable-all +""" +Here we have 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. +""" + +# Copyright (c) 2022, ETH Zurich and UNC Chapel Hill. +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of +# its contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. +# +# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de) + +import os +import collections +import numpy as np +import struct +import argparse + + +CameraModel = collections.namedtuple("CameraModel", ["model_id", "model_name", "num_params"]) +Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"]) +BaseImage = collections.namedtuple("Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]) +Point3D = collections.namedtuple("Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"]) + + +class Image(BaseImage): + def qvec2rotmat(self): + return qvec2rotmat(self.qvec) + + +CAMERA_MODELS = { + CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), + CameraModel(model_id=1, model_name="PINHOLE", num_params=4), + CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), + CameraModel(model_id=3, model_name="RADIAL", num_params=5), + CameraModel(model_id=4, model_name="OPENCV", num_params=8), + CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), + CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), + CameraModel(model_id=7, model_name="FOV", num_params=5), + CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), + CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), + CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), +} +CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]) +CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS]) + + +def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): + """Read and unpack the next bytes from a binary file. + :param fid: + :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc. + :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. + :param endian_character: Any of {@, =, <, >, !} + :return: Tuple of read and unpacked values. + """ + data = fid.read(num_bytes) + 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. + :param fid: + :param data: data to send, if multiple elements are sent at the same time, + they should be encapsuled either in a list or a tuple + :param 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 + :param 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 + void Reconstruction::WriteCamerasText(const std::string& path) + void Reconstruction::ReadCamerasText(const std::string& path) + """ + cameras = {} + with open(path, "r") as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + camera_id = int(elems[0]) + model = elems[1] + width = int(elems[2]) + height = int(elems[3]) + params = np.array(tuple(map(float, elems[4:]))) + cameras[camera_id] = Camera(id=camera_id, model=model, width=width, height=height, params=params) + return cameras + + +def read_cameras_binary(path_to_model_file): + """ + see: src/base/reconstruction.cc + void Reconstruction::WriteCamerasBinary(const std::string& path) + void Reconstruction::ReadCamerasBinary(const std::string& path) + """ + cameras = {} + with open(path_to_model_file, "rb") as fid: + num_cameras = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_cameras): + camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ") + camera_id = camera_properties[0] + model_id = camera_properties[1] + model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name + width = camera_properties[2] + height = camera_properties[3] + num_params = CAMERA_MODEL_IDS[model_id].num_params + params = read_next_bytes(fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params) + cameras[camera_id] = Camera( + id=camera_id, model=model_name, width=width, height=height, params=np.array(params) + ) + assert len(cameras) == num_cameras + 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" + + "# Number of cameras: {}\n".format(len(cameras)) + ) + 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 + void Reconstruction::ReadImagesText(const std::string& path) + void Reconstruction::WriteImagesText(const std::string& path) + """ + images = {} + with open(path, "r") as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + image_id = int(elems[0]) + qvec = np.array(tuple(map(float, elems[1:5]))) + tvec = np.array(tuple(map(float, elems[5:8]))) + camera_id = int(elems[8]) + image_name = elems[9] + elems = fid.readline().split() + xys = np.column_stack([tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))]) + point3D_ids = np.array(tuple(map(int, elems[2::3]))) + images[image_id] = Image( + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) + return images + + +def read_images_binary(path_to_model_file): + """ + see: src/base/reconstruction.cc + void Reconstruction::ReadImagesBinary(const std::string& path) + void Reconstruction::WriteImagesBinary(const std::string& path) + """ + images = {} + with open(path_to_model_file, "rb") as fid: + num_reg_images = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_reg_images): + binary_image_properties = read_next_bytes(fid, num_bytes=64, format_char_sequence="idddddddi") + image_id = binary_image_properties[0] + qvec = np.array(binary_image_properties[1:5]) + tvec = np.array(binary_image_properties[5:8]) + camera_id = binary_image_properties[8] + image_name = "" + current_char = read_next_bytes(fid, 1, "c")[0] + while current_char != b"\x00": # look for the ASCII 0 entry + image_name += current_char.decode("utf-8") + current_char = read_next_bytes(fid, 1, "c")[0] + num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] + x_y_id_s = read_next_bytes(fid, num_bytes=24 * num_points2D, format_char_sequence="ddq" * num_points2D) + xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))]) + point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) + images[image_id] = Image( + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) + 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" + + "# Number of images: {}, mean observations per image: {}\n".format(len(images), mean_observations) + ) + + 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 + void Reconstruction::ReadPoints3DText(const std::string& path) + void Reconstruction::WritePoints3DText(const std::string& path) + """ + points3D = {} + with open(path, "r") as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + point3D_id = int(elems[0]) + xyz = np.array(tuple(map(float, elems[1:4]))) + rgb = np.array(tuple(map(int, elems[4:7]))) + error = float(elems[7]) + image_ids = np.array(tuple(map(int, elems[8::2]))) + point2D_idxs = np.array(tuple(map(int, elems[9::2]))) + points3D[point3D_id] = Point3D( + id=point3D_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2D_idxs=point2D_idxs + ) + return points3D + + +def read_pointsTD_binary(path_to_model_file): + """ + see: src/base/reconstruction.cc + void Reconstruction::ReadPoints3DBinary(const std::string& path) + void Reconstruction::WritePoints3DBinary(const std::string& path) + """ + points3D = {} + with open(path_to_model_file, "rb") as fid: + num_points = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_points): + binary_point_line_properties = read_next_bytes(fid, num_bytes=43, format_char_sequence="QdddBBBd") + point3D_id = binary_point_line_properties[0] + xyz = np.array(binary_point_line_properties[1:4]) + rgb = np.array(binary_point_line_properties[4:7]) + error = np.array(binary_point_line_properties[7]) + track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] + track_elems = read_next_bytes(fid, num_bytes=8 * track_length, format_char_sequence="ii" * track_length) + image_ids = np.array(tuple(map(int, track_elems[0::2]))) + point2D_idxs = np.array(tuple(map(int, track_elems[1::2]))) + points3D[point3D_id] = Point3D( + id=point3D_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2D_idxs=point2D_idxs + ) + 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" + + "# Number of points: {}, mean track length: {}\n".format(len(points3D), mean_track_length) + ) + + 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)) + and os.path.isfile(os.path.join(path, "images" + ext)) + and os.path.isfile(os.path.join(path, "points3D" + ext)) + ): + print("Detected model format: '" + ext + "'") + return True + + return False + + +def read_model(path, ext=""): + # try to detect the extension automatically + if ext == "": + if detect_model_format(path, ".bin"): + ext = ".bin" + elif detect_model_format(path, ".txt"): + ext = ".txt" + else: + print("Provide model format: '.bin' or '.txt'") + return + + if ext == ".txt": + cameras = read_cameras_text(os.path.join(path, "cameras" + ext)) + images = read_images_text(os.path.join(path, "images" + ext)) + points3D = read_pointsTD_text(os.path.join(path, "points3D") + ext) + else: + cameras = read_cameras_binary(os.path.join(path, "cameras" + ext)) + images = read_images_binary(os.path.join(path, "images" + ext)) + points3D = read_pointsTD_binary(os.path.join(path, "points3D") + 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( + [ + [ + 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], + 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], + ], + [ + 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], + ], + [ + 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], + 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, + ], + ] + ) + + +def rotmat2qvec(R): + Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat + K = ( + np.array( + [ + [Rxx - Ryy - Rzz, 0, 0, 0], + [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], + [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], + [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz], + ] + ) + / 3.0 + ) + eigvals, eigvecs = np.linalg.eigh(K) + qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] + 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/mattport/nerf/dataset/friends.py b/mattport/nerf/dataset/friends.py new file mode 100644 index 000000000..33bfce485 --- /dev/null +++ b/mattport/nerf/dataset/friends.py @@ -0,0 +1,90 @@ +""" +Code to handle loading friends datasets. +""" + +import os +import numpy as np + +import torch + +from mattport.nerf.dataset.colmap_utils import read_cameras_binary, read_images_binary, read_pointsTD_binary +from mattport.nerf.dataset.structs import DatasetInputs, PointCloud + + +def load_friends_data(basedir, downscale_factor=1, split="train", include_point_cloud=False): + """_summary_ + + Args: + basedir (_type_): _description_ + downscale_factor (float, optional): _description_. Defaults to 1.0. + split (str, optional): _description_. Defaults to "train". + include_point_cloud (bool): whether or not to include the point cloud + """ + # pylint: disable=unused-argument + + # --- image filenames ---- + images_data = read_images_binary(os.path.join(basedir, "colmap", "images.bin")) + # `image_path` is only the end of the filename, including the extension e.g., `.jpg` + image_paths = sorted(os.listdir(os.path.join(basedir, "images"))) + + image_path_to_image_id = {} + image_id_to_image_path = {} + for v in images_data.values(): + image_path_to_image_id[v.name] = v.id + image_id_to_image_path[v.id] = v.name + # TODO: handle the splits differently + image_filenames = [os.path.join(basedir, "images", image_path) for image_path in image_paths] + + # --- intrinsics --- + cameras_data = read_cameras_binary(os.path.join(basedir, "colmap", "cameras.bin")) + intrinsics = [] + for image_path in image_paths: + cam = cameras_data[image_path_to_image_id[image_path]] + print(cam.params) + assert len(cam.params) == 3 + focal_length = cam.params[0] # f (fx and fy) + cx = cam.params[1] # cx + cy = cam.params[2] # cy + intrinsics.append([cx, cy, focal_length]) + intrinsics = torch.tensor(intrinsics).float() + + # --- camera_to_world (extrinsics) --- + camera_to_world = [] + bottom_row = np.array([0, 0, 0, 1.0]).reshape(1, 4) + for image_path in image_paths: + image_data = images_data[image_path_to_image_id[image_path]] + rot = image_data.qvec2rotmat() + trans = image_data.tvec.reshape(3, 1) + c2w = np.concatenate([np.concatenate([rot, trans], 1), bottom_row], 0) + camera_to_world.append(c2w) + camera_to_world = torch.tensor(np.array(camera_to_world)).float() + camera_to_world = torch.inverse(camera_to_world) + camera_to_world = camera_to_world[:, :3] + camera_to_world[..., 1:3] *= -1 + + # --- masks to mask out things (e.g., people) --- + # TODO: + + # --- possibly transform --- + # TODO: + + # -- set the bounding box --- + # TODO: + + # Possibly include the sparse point cloud from COLMAP in the dataset inputs. + # NOTE(ethan): this will be common across the different splits. + point_cloud = None + if include_point_cloud: + points_3d = read_pointsTD_binary(os.path.join(basedir, "colmap", "points3D.bin")) + xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])) + rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])) + point_cloud = PointCloud(xyz=xyz, rgb=rgb) + + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + downscale_factor=downscale_factor, + intrinsics=intrinsics / downscale_factor, + camera_to_world=camera_to_world, + point_cloud=point_cloud, + ) + return dataset_inputs diff --git a/mattport/nerf/dataset/llff.py b/mattport/nerf/dataset/llff.py new file mode 100644 index 000000000..119dfabda --- /dev/null +++ b/mattport/nerf/dataset/llff.py @@ -0,0 +1,16 @@ +""" +Code for loading the LLFF data. +""" + + +def load_llff_data(basedir, downscale_factor=1, split="train", include_point_cloud=False): + """_summary_ + + Args: + basedir (_type_): _description_ + downscale_factor (float, optional): _description_. Defaults to 1.0. + split (str, optional): _description_. Defaults to "train". + include_point_cloud (bool): whether or not to include the point cloud + """ + # pylint: disable=unused-argument + raise NotImplementedError diff --git a/mattport/nerf/dataset/structs.py b/mattport/nerf/dataset/structs.py new file mode 100644 index 000000000..39ca9d371 --- /dev/null +++ b/mattport/nerf/dataset/structs.py @@ -0,0 +1,42 @@ +""" +Dataset input structures. +""" + +from dataclasses import dataclass +from typing import List, Optional + +import torch +from torchtyping import TensorType + + +@dataclass +class PointCloud: + """_summary_""" + + xyz: None + rgb: None + + +@dataclass +class DatasetInputs: + """Dataset inputs for the image dataset and the ray generator. + + Args: + image_filenames: Filenames for the images. + ... + """ + + image_filenames: List[str] + downscale_factor: int = 1 + intrinsics: torch.tensor = None + camera_to_world: torch.tensor = None + bounding_box: torch.tensor = None + depth_filenames: List[str] = None + semantic_filenames: List[str] = None + point_cloud: PointCloud = None # TODO(ethan): specify the type here + alpha_color: Optional[TensorType[3]] = None + + def check_inputs(self): + """Check the inputs to make sure everything is okay.""" + assert self.intrinsics.dtype == torch.float32 + assert self.camera_to_world.dtype == torch.float32 diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py index 69031846b..9c6425464 100644 --- a/mattport/nerf/dataset/utils.py +++ b/mattport/nerf/dataset/utils.py @@ -2,79 +2,15 @@ For loading the blender dataset format. """ -import os -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, Optional, Tuple, Union -import imageio -import numpy as np from omegaconf import ListConfig -import torch -from torchtyping import TensorType + +from mattport.nerf.dataset.blender import load_blender_data +from mattport.nerf.dataset.friends import load_friends_data +from mattport.nerf.dataset.structs import DatasetInputs from mattport.structures.colors import get_color - -from mattport.utils.io import get_absolute_path, load_from_json - - -@dataclass -class DatasetInputs: - """Dataset inputs for the image dataset and the ray generator.""" - - image_filenames: List[str] - downscale_factor: float = 1.0 - alpha_color: Optional[TensorType[3]] = None - intrinsics: torch.tensor = None - camera_to_world: torch.tensor = None - - -def load_blender_data( - basedir: str, downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None, split: str = "train" -) -> DatasetInputs: - """Processes the a blender dataset directory. - Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37. - - Args: - data_directory (str): Location of data - downscale_factor (int, optional): How much to downscale images. Defaults to 1.0. - alpha_color (TensorType[3], optional): Sets transparent regions to specified color, otherwise black. - split (str, optional): Which dataset split to generate. - - Returns: - DatasetInputs - """ - # pylint: disable=unused-argument - # pylint: disable-msg=too-many-locals - meta = load_from_json(os.path.join(basedir, f"transforms_{split}.json")) - image_filenames = [] - poses = [] - for frame in meta["frames"]: - fname = os.path.join(basedir, frame["file_path"].replace("./", "") + ".png") - image_filenames.append(fname) - poses.append(np.array(frame["transform_matrix"])) - poses = np.array(poses).astype(np.float32) - - img_0 = imageio.imread(image_filenames[0]) - image_height, image_width = img_0.shape[:2] - camera_angle_x = float(meta["camera_angle_x"]) - focal_length = 0.5 * image_width / np.tan(0.5 * camera_angle_x) - - cx = image_width / 2.0 - cy = image_height / 2.0 - camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform - num_cameras = len(image_filenames) - num_intrinsics_params = 3 - intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) - intrinsics *= torch.tensor([cx, cy, focal_length]) - - dataset_inputs = DatasetInputs( - image_filenames=image_filenames, - downscale_factor=downscale_factor, - alpha_color=alpha_color, - intrinsics=intrinsics * 1.0 / downscale_factor, # downscaling the intrinsics here - camera_to_world=camera_to_world, - ) - - return dataset_inputs +from mattport.utils.io import get_absolute_path def get_dataset_inputs_dict( @@ -109,6 +45,10 @@ def get_dataset_inputs_dict( alpha_color=alpha_color, ) dataset_inputs_dict[split] = dataset_inputs + elif dataset_type == "friends": + for split in splits: + dataset_inputs = load_friends_data(get_absolute_path(data_directory), downscale_factor=downscale_factor) + dataset_inputs_dict[split] = dataset_inputs else: raise NotImplementedError(f"{dataset_type} is not a valid dataset type") diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py index 0bb180492..66ede7afc 100644 --- a/mattport/nerf/field_modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -1,6 +1,7 @@ """ Encoding functions """ + import numpy as np import torch from torchtyping import TensorType diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 37eaa260e..29a4d0278 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -29,7 +29,7 @@ logging.getLogger("PIL").setLevel(logging.WARNING) class Trainer: """Training class""" - def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1, cpu=False): + def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1): """_summary_ Args: @@ -57,7 +57,7 @@ class Trainer: self.stats = StatsTracker(config, self.is_main_thread) if not profiler.PROFILER and self.config.logging.enable_profiler: profiler.PROFILER = profiler.Profiler(config, self.is_main_thread) - self.device = f"cuda:{self.local_rank}" if not cpu else "cpu" + self.device = "cpu" if self.world_size == 0 else f"cuda:{self.local_rank}" @profiler.time_function def setup(self): @@ -112,7 +112,7 @@ class Trainer: intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world, ).to(self.device) - self.setup_optimizers() + self.setup_optimizers() # NOTE(ethan): can this be before DDP? if self.config.graph.resume_train.load_dir: self.load_checkpoint(self.config.graph.resume_train) @@ -206,7 +206,7 @@ class Trainer: self.test_image(image_idx=image_idx, step=step) self.stats.print_stats(i / num_iterations) - self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1) + self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1) # NOTE(ethan): why is step -1? self.stats.print_stats(-1) @profiler.time_function @@ -221,7 +221,9 @@ class Trainer: if hasattr(self.graph, "module") else self.graph.get_losses(batch, graph_outputs) ) - loss_sum, loss_dict = self.get_aggregated_loss(losses) + loss_sum, loss_dict = self.get_aggregated_loss( + losses + ) # TODO(ethan): move this into the graph itself. not in trainer self.optimizers.zero_grad_all() loss_sum.backward() self.optimizers.scheduler_step_all(step) # NOTE(ethan): I think the scheduler needs to know what step we are on @@ -237,7 +239,7 @@ class Trainer: pixel_coords = torch.stack(pixel_coords, dim=-1).long() all_ray_indices = torch.cat([torch.ones_like(pixel_coords[..., :1]) * image_idx, pixel_coords], dim=-1).view( -1, 3 - ) + ) # NOTE(ethan): this is ugly... clean up... with torch.no_grad(): num_rays = all_ray_indices.shape[0] chunk_size = 1024 diff --git a/mattport/structures/cameras.py b/mattport/structures/cameras.py index 6e8f6e5e1..f6922ba37 100644 --- a/mattport/structures/cameras.py +++ b/mattport/structures/cameras.py @@ -217,4 +217,6 @@ def get_camera_model(num_intrinsics_params): """ if num_intrinsics_params == 3: return SimplePinholeCamera + if num_intrinsics_params == 4: + return PinholeCamera raise NotImplementedError diff --git a/mattport/utils/stats_tracker.py b/mattport/utils/stats_tracker.py index 03461c050..58a181658 100644 --- a/mattport/utils/stats_tracker.py +++ b/mattport/utils/stats_tracker.py @@ -30,8 +30,8 @@ class StatsTracker: def __init__(self, config: DictConfig, is_main_thread: bool): self.config = config + self.is_main_thread = is_main_thread if self.config.logging.enable_stats: - self.is_main_thread = is_main_thread self.max_history = self.config.logging.stats_tracker.max_history self.step = 0 self.stats_dict = {} diff --git a/scripts/run_train.py b/scripts/run_train.py index 28c587a15..dec47e315 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -105,7 +105,7 @@ def _distributed_worker( return output -def _train(local_rank: int, world_size: int, config: DictConfig, cpu: bool = False) -> Any: +def _train(local_rank: int, world_size: int, config: DictConfig) -> Any: """Main training function that sets up and runs the trainer per process Args: @@ -116,7 +116,7 @@ def _train(local_rank: int, world_size: int, config: DictConfig, cpu: bool = Fal Returns: Any: TODO(): determine the return type """ - trainer = Trainer(config, local_rank, world_size, cpu=cpu) + trainer = Trainer(config, local_rank, world_size) trainer.setup() trainer.train() return 0 @@ -145,7 +145,7 @@ def launch( world_size = num_machines * num_gpus_per_machine if world_size == 0: # Using only CPU and one process. - main_func(local_rank=0, world_size=0, config=config, cpu=True) + main_func(local_rank=0, world_size=0, config=config) elif world_size == 1: # Using one gpu and one process. try: diff --git a/tests/test_run_train.py b/tests/test_run_train.py index 879f9ecdb..b8a82d7a7 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -10,6 +10,6 @@ from scripts.run_train import main @pytest.mark.filterwarnings("ignore::DeprecationWarning") def test_run_train(): """test run train script works properly""" - initialize(config_path="../configs/test/") + initialize(config_path="../configs/tests/") conf = compose("test_dryrun.yaml") main(conf)