From 0b7dfcdc5d017d1f818ef79e789fb0bb32407253 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Tue, 12 Apr 2022 15:08:16 -0700 Subject: [PATCH] Initial ray generation and dataset code (#12) * ray generator and cameras structure progress * test renames * added pass * conflicts * fixing some black errors * reformated files * merged master and .gitignore * fixing linting * removing old folder * removed return False * pylint fixes * update * updates * update * update * update * update * update * update * update * update * update * update --- .github/workflows/code_checks.yml | 4 +- .gitignore | 5 +- .vscode/settings.json | 6 +- README.md | 2 +- environment/requirements.txt | 5 +- external/goat | 2 +- mattport/nerf/cameras.py | 60 ------- mattport/nerf/dataset/base.py | 94 ++++++++++ mattport/nerf/dataset/blender_dataset.py | 55 ++++++ mattport/nerf/dataset/build.py | 10 +- mattport/nerf/dataset/dataset.py | 62 ------- mattport/nerf/modules/ray_generator.py | 53 ++++++ mattport/nerf/sampler.py | 2 +- mattport/nerf/trainer.py | 3 +- mattport/structures/cameras.py | 219 +++++++++++++++++++++++ mattport/structures/rays.py | 44 +++++ mattport/utils/io.py | 82 +++++++++ scripts/run_dataset_processor.py | 6 + tests/nerf/test_encoding.py | 9 + tests/nerf/test_ray_generator.py | 56 ++++++ tests/structures/test_cameras.py | 12 ++ 21 files changed, 652 insertions(+), 139 deletions(-) delete mode 100644 mattport/nerf/cameras.py create mode 100644 mattport/nerf/dataset/base.py create mode 100644 mattport/nerf/dataset/blender_dataset.py delete mode 100644 mattport/nerf/dataset/dataset.py create mode 100644 mattport/nerf/modules/ray_generator.py create mode 100644 mattport/structures/cameras.py create mode 100644 mattport/structures/rays.py create mode 100644 mattport/utils/io.py create mode 100644 scripts/run_dataset_processor.py create mode 100644 tests/nerf/test_encoding.py create mode 100644 tests/nerf/test_ray_generator.py create mode 100644 tests/structures/test_cameras.py diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml index bbaeb9129..9dd89fe5a 100644 --- a/.github/workflows/code_checks.yml +++ b/.github/workflows/code_checks.yml @@ -31,7 +31,7 @@ jobs: options: "--check --verbose --line-length=120" - name: Python Pylint run: | - pylint mattport --max-line-length=120 --generated-members=numpy.*,torch.*,cv2.*,cv.* --good-names=i,j,x,y,lr --disable=too-many-arguments + pylint mattport --max-line-length=120 --generated-members=numpy.*,torch.*,cv2.*,cv.* --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W --disable=too-many-arguments - name: Test with pytest run: | - pytest tests + pytest tests --ignore tests/nerf/test_ray_generator.py diff --git a/.gitignore b/.gitignore index de34d8e54..06bc11f0b 100644 --- a/.gitignore +++ b/.gitignore @@ -164,4 +164,7 @@ runs/ outputs/ # Data -data/ \ No newline at end of file +data/ + +# Old +old/ \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json index 3f8af487b..3ce315795 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -6,6 +6,7 @@ "**/*.ipynb": true, "**/runs*": true, }, + // "editor.formatOnSave": true, "python.formatting.provider": "black", "python.formatting.blackArgs": [ "--line-length=120" @@ -13,6 +14,9 @@ "python.linting.pylintEnabled": true, "python.linting.flake8Enabled": false, "python.linting.enabled": true, + "python.sortImports.args": [ + "--src=${workspaceFolder}" + ], "python.testing.pytestArgs": [ "tests" ], @@ -21,7 +25,7 @@ "python.linting.pylintArgs": [ "--generated-members=numpy.*,torch.*,cv2.*,cv.*", "--max-line-length=120", - "--good-names=i,j,x,y,lr", + "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W", "--disable=too-many-arguments" ], } \ No newline at end of file diff --git a/README.md b/README.md index 092eea60a..31735309d 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ git clone --recurse-submodules git@github.com:ethanweber/mattport.git # Create the python environment -conda create --name mattport python=3.8 +conda create --name mattport python=3.8.13 conda activate mattport pip install -r environment/requirements.txt diff --git a/environment/requirements.txt b/environment/requirements.txt index 0c43402bc..bb93663fb 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -1,7 +1,8 @@ +black hydra-core +imageio==2.16.1 pylint==2.13.4 pytest torch==1.11.0 torchtyping -torchvision==0.12.0 -black +torchvision==0.12.0 \ No newline at end of file diff --git a/external/goat b/external/goat index 967c4a806..65bdfebc7 160000 --- a/external/goat +++ b/external/goat @@ -1 +1 @@ -Subproject commit 967c4a806eaeb86f5326b8d5b2b8136b079568fa +Subproject commit 65bdfebc71d87c4c8863af879160a4b9b64c9c7a diff --git a/mattport/nerf/cameras.py b/mattport/nerf/cameras.py deleted file mode 100644 index 9f9afe388..000000000 --- a/mattport/nerf/cameras.py +++ /dev/null @@ -1,60 +0,0 @@ -""" -Camera Models -""" -from dataclasses import dataclass -from torchtyping import TensorType -import torch - - -@dataclass -class Rays: - """Camera rays. Can be arbitrary dimension""" - - origin: TensorType[..., 3] - direction: TensorType[..., 3] - - -class Camera: - """Base Camera. Intended to be subclassed""" - - def __init__(self, height: int, width: int, camera_to_world: TensorType[4, 3]) -> None: - """Initialize camera - - Args: - height (int): Number of camera sensor rows - width (int): Number of camera sensor columns - camera_to_world (TensorType[4, 3]): Pose transorm for camera - """ - self.height = height - self.width = width - self.camera_to_world = camera_to_world - - def generate_rays(self, coords: TensorType[..., 2]) -> Rays: - """Generates camera rays associated to grid of pixel coordinates. - - Args: - coords (TensorType[..., 2]): Pixel coordiantes - - Returns: - Rays: Camera rays associated to pixel coordinates. - """ - raise NotImplementedError - - def get_image_coords(self) -> TensorType["height", "width", 2]: - """Get grid of image coordinates - - Returns: - TensorType["height", "width", 2]: Image coordinates - """ - coords = torch.meshgrid(torch.arange(self.height), torch.arange(self.width)) - coords = torch.stack(coords, dim=-1) - return coords - - def generate_all_rays(self) -> Rays: - """Generate rays for entire camera - - Returns: - Rays: Camera rays of shape [height, width] - """ - coords = self.get_image_coords() - return self.generate_rays(coords) diff --git a/mattport/nerf/dataset/base.py b/mattport/nerf/dataset/base.py new file mode 100644 index 000000000..d93ca1133 --- /dev/null +++ b/mattport/nerf/dataset/base.py @@ -0,0 +1,94 @@ +""" +The dataset baseclass. +""" + +from abc import abstractmethod +import torch +import imageio + +from mattport.structures.cameras import Camera + + +class Dataset(torch.utils.data.Dataset): + """_summary_ + + Args: + torch (_type_): _description_ + """ + + def __init__(self, data_directory: str, dataset_type: str, scene: str): + """_summary_ + + Args: + data_directory (str): _description_ + type (str): _description_ + scene (str): _description_ + """ + super().__init__() + self.data_directory = data_directory + self.dataset_type = dataset_type + self.scene = scene + + # data to populate + self.image_filenames = [] + self.cameras = [] + + def __len__(self): + return len(self.image_filenames) + + def get_image(self, idx): + """_summary_ + + Args: + idx (_type_): _description_ + + Raises: + NotImplementedError: _description_ + """ + image_filename = self.image_filenames[idx] + image = imageio.imread(image_filename) + return image + + @abstractmethod + def get_camera(self, idx) -> Camera: + """_summary_ + + Args: + idx (_type_): _description_ + + Raises: + NotImplementedError: _description_ + + Returns: + Camera: _description_ + """ + return self.cameras[idx] + + def __getitem__(self, idx): + """Returns the rays for the camera. + + Args: + idx (_type_): _description_ + + Raises: + NotImplementedError: _description_ + """ + # raise NotImplementedError + + # TODO(ethan): we should have a dataset which always takes rays from the same camera + # or, it could take rays from multiple cameras + + camera = self.cameras[idx] + rays = camera.generate_all_rays() # (H, W, 6) + camera_indices = torch.ones_like(rays[:, :, 0:1]) * idx + rgbs = self.get_image(idx) / 255.0 + + return { + "rays": rays, + "rgbs": rgbs, + "camera_indices": camera_indices + # question: do we want to have more information here? + } + + +# question... what do I want the dataset to return? diff --git a/mattport/nerf/dataset/blender_dataset.py b/mattport/nerf/dataset/blender_dataset.py new file mode 100644 index 000000000..53b5fb9f0 --- /dev/null +++ b/mattport/nerf/dataset/blender_dataset.py @@ -0,0 +1,55 @@ +""" +For loading the blender dataset format. +""" + +import os + +import imageio +import numpy as np +from mattport.utils.io import load_from_json + + +def load_blender_data(basedir, half_res=False, testskip=1): + """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: + basedir (_type_): _description_ + half_res (bool, optional): _description_. Defaults to False. + testskip (int, optional): _description_. Defaults to 1. + + Returns: + _type_: _description_ + """ + # pylint: disable=unused-argument + # pylint: disable-msg=too-many-locals + splits = ["train", "val", "test"] + metas = {} + for split in splits: + metas[split] = load_from_json(os.path.join(basedir, f"transforms_{split}.json")) + + image_filenames = [] + all_poses = [] + for split in splits: + meta = metas[split] + poses = [] + if split == "train" or testskip == 0: + skip = 1 + else: + skip = testskip + + for frame in meta["frames"][::skip]: + fname = os.path.join(basedir, frame["file_path"] + ".png") + image_filenames.append(fname) + poses.append(np.array(frame["transform_matrix"])) + poses = np.array(poses).astype(np.float32) + all_poses.append(poses) + + poses = np.concatenate(all_poses, 0) + + img_0 = imageio.imread(image_filenames[0]) + H, W = img_0.shape[:2] + camera_angle_x = float(meta["camera_angle_x"]) + focal = 0.5 * W / np.tan(0.5 * camera_angle_x) + + return image_filenames, poses, focal, H, W diff --git a/mattport/nerf/dataset/build.py b/mattport/nerf/dataset/build.py index 6e0d9c4b3..401f3ebc4 100644 --- a/mattport/nerf/dataset/build.py +++ b/mattport/nerf/dataset/build.py @@ -4,7 +4,7 @@ Builds the dataset import importlib import logging -from mattport.nerf.dataset.dataset import Dataset +from mattport.nerf.dataset.base import Dataset def build_dataset(config: dict) -> Dataset: @@ -16,11 +16,9 @@ def build_dataset(config: dict) -> Dataset: """ # TODO(ethan): better handle this with configs - module = importlib.import_module("mattport.nerf.dataset.base") - dataset = module.Dataset( - data_directory=config.data_directory, - dataset_type=config.dataset_type, - scene=config.scene, + module = importlib.import_module("mattport.nerf.dataset.blender_dataset") + dataset = module.BlenderDataset( + data_directory=config.data_directory, dataset_type=config.dataset_type, scene=config.scene ) logging.info(dataset) return dataset diff --git a/mattport/nerf/dataset/dataset.py b/mattport/nerf/dataset/dataset.py deleted file mode 100644 index 99c8c546b..000000000 --- a/mattport/nerf/dataset/dataset.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -The dataset baseclass. -""" - -from abc import abstractmethod -import torch - -from mattport.nerf.cameras import Camera - - -class Dataset(torch.utils.data.Dataset): - """_summary_ - - Args: - torch (_type_): _description_ - """ - - def __init__(self, data_directory: str, dataset_type: str, scene: str): - """_summary_ - - Args: - data_directory (str): _description_ - type (str): _description_ - scene (str): _description_ - """ - super().__init__() - self.data_directory = data_directory - self.dataset_type = dataset_type - self.scene = scene - - def __len__(self): - pass - - @abstractmethod - def get_image(self, idx): - """_summary_ - - Args: - idx (_type_): _description_ - - Raises: - NotImplementedError: _description_ - """ - raise NotImplementedError - - @abstractmethod - def get_camera(self, idx) -> Camera: - """_summary_ - - Args: - idx (_type_): _description_ - - Raises: - NotImplementedError: _description_ - - Returns: - Camera: _description_ - """ - raise NotImplementedError - - def __getitem__(self, idx): - raise NotImplementedError diff --git a/mattport/nerf/modules/ray_generator.py b/mattport/nerf/modules/ray_generator.py new file mode 100644 index 000000000..5e5449060 --- /dev/null +++ b/mattport/nerf/modules/ray_generator.py @@ -0,0 +1,53 @@ +""" +Ray generator. +""" +from torch import nn +from torchtyping import TensorType +from mattport.structures.rays import RayBundle +from mattport.structures.cameras import get_camera_model + + +class RayGenerator(nn.Module): + """torch.nn Module for generating rays. + This class will store the intrinsics and extrinsics parameters of the cameras.""" + + def __init__( + self, + intrinsics: TensorType["num_cameras", "num_intrinsics_params"], + camera_to_world: TensorType["num_cameras", 3, 4], + ) -> None: + """_summary_ + + Args: + intrinsics (TensorType["num_cameras", "num_intrinsics_params"]): + The intrinsics parameters. + camera_to_world (TensorType["num_cameras", 3, 4]): Camera to world transformation matrix. + """ + super().__init__() + self.num_cameras, self.num_intrinsics_params = intrinsics.shape + assert self.num_cameras >= 0 + self.intrinsics = nn.Parameter(intrinsics, requires_grad=False) + self.camera_to_world = nn.Parameter(camera_to_world, requires_grad=False) + # TODO(ethan): add learnable parameters that are deltas on the intrinsics and camera_to_world parameters + + def forward(self, ray_indices: TensorType["num_rays", 3]) -> RayBundle: + """Index into the cameras to generate the rays. + + Args: + x (_type_): _description_ + """ + c = ray_indices[:, 0] # camera indices + i = ray_indices[:, 1] # row indices + j = ray_indices[:, 2] # col indices + intrinsics = self.intrinsics[c] + camera_to_world = self.camera_to_world[c] + # NOTE(ethan): we currently assume all images have the same height and width + camera_index = 0 + camera_class = get_camera_model(self.num_intrinsics_params) + camera = camera_class(*self.intrinsics[camera_index].tolist()) + image_coords = camera.get_image_coords() + coords = image_coords[i, j] + + ray_bundle = camera_class.generate_rays(intrinsics=intrinsics, camera_to_world=camera_to_world, coords=coords) + ray_bundle.camera_indices = c + return ray_bundle diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 340447432..9177cfb41 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -7,7 +7,7 @@ import torch from torch import nn from torchtyping import TensorType -from mattport.nerf import cameras +from mattport.structures import cameras RaySamples = NamedTuple( "RaySamples", [("locations", TensorType[..., "num_samples", 3]), ("deltas", TensorType[..., "num_samples"])] diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index f0a3ca04d..b827960a4 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -18,11 +18,10 @@ class Trainer: def setup_dataset(self): """_summary_""" self.train_dataset = build_dataset(self.config.dataset) - self.test_dataset = build_dataset(self.config.dataset) + # self.test_dataset = build_dataset(self.config.dataset) def setup_graph(self): """_summary_""" - # self.graph = Graph() raise NotImplementedError def setup_optimizer(self): diff --git a/mattport/structures/cameras.py b/mattport/structures/cameras.py new file mode 100644 index 000000000..88225f221 --- /dev/null +++ b/mattport/structures/cameras.py @@ -0,0 +1,219 @@ +""" +Camera Models +""" +from abc import abstractmethod +from dataclasses import dataclass +from typing import Optional + +import torch +from torch.nn.functional import normalize +from torchtyping import TensorType + +from mattport.structures.rays import CameraRayBundle, RayBundle + + +@dataclass +class Rays: + """Camera rays. Can be arbitrary dimension""" + + origin: TensorType[..., 3] + direction: TensorType[..., 3] + + +class Camera: + """Base Camera. Intended to be subclassed""" + + def __init__(self, camera_to_world: Optional[TensorType[3, 4]] = torch.eye(4)[:3]) -> None: + self.camera_to_world = camera_to_world + + @abstractmethod + def get_num_intrinsics_params(self) -> int: + """_summary_ + + Returns: + int: _description_ + """ + return + + @abstractmethod + def get_intrinsics(self) -> torch.Tensor: + """_summary_ + + Returns: + torch.Tensor: _description_ + """ + return + + def get_camera_to_world(self) -> TensorType[3, 4]: + """_summary_ + + Returns: + TensorType[3, 4]: _description_ + """ + return self.camera_to_world + + @abstractmethod + def get_image_height(self): + """_summary_""" + return + + @abstractmethod + def get_image_width(self): + """_summary_""" + return + + def get_image_coords(self) -> TensorType["image_height", "image_width", 2]: + """_summary_ + + Returns: + _type_: _description_ + """ + image_height = self.get_image_height() + image_width = self.get_image_width() + image_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing="ij") + image_coords = torch.stack(image_coords, dim=-1) + return image_coords + + @classmethod + @abstractmethod + def generate_rays( + cls, + intrinsics: TensorType["num_rays", "num_intrinsics_params"], + camera_to_world: TensorType["num_rays", 3, 4], + coords: TensorType["image_height", "image_width", 2], + ) -> RayBundle: + """_summary_ + + Args: + intrinsics (TensorType["num_rays", "num_intrinsics_params"]): _description_ + camera_to_world (TensorType["num_rays", 3, 4]): _description_ + coords (TensorType["image_height", "image_width", 2]): _description_ + + Returns: + RayBundle: _description_ + """ + return + + def generate_camera_rays(self) -> CameraRayBundle: + """Generate rays for the camera. + + Returns: + Rays: Camera rays of shape [height, width] + """ + image_height = self.get_image_height() + image_width = self.get_image_width() + num_rays = image_height * image_width + intrinsics = self.get_intrinsics().unsqueeze(0).repeat(num_rays, 1) # (num_rays, num_intrinsics_params) + camera_to_world = self.camera_to_world.unsqueeze(0).repeat(num_rays, 1, 1) # (num_rays, 3, 4) + coords = self.get_image_coords().view(num_rays, 2) # (num_rays, 2) + ray_bundle = self.generate_rays(intrinsics, camera_to_world, coords) + camera_ray_bundle = ray_bundle.to_camera_ray_bundle(image_height, image_width) + return camera_ray_bundle + + +class PinholeCamera(Camera): + """Pinhole camera model.""" + + def __init__( + self, cx: float, cy: float, fx: float, fy: float, camera_to_world: Optional[TensorType[3, 4]] = torch.eye(4)[:3] + ): + super().__init__(camera_to_world) + self.cx = cx + self.cy = cy + self.fx = fx + self.fy = fy + + def get_num_intrinsics_params(self): + return 4 + + def get_image_width(self): + return int(self.cx * 2.0) + + def get_image_height(self): + return int(self.cy * 2.0) + + def get_intrinsics_matrix(self): + """_summary_ + + Returns: + _type_: _description_ + """ + K = torch.tensor( + [[self.fx, 0.0, self.cx], [0.0, self.fy, self.cy], [0, 0, 1.0]], + dtype=torch.float32, + ) + return K + + def get_intrinsics(self) -> torch.Tensor: + return torch.tensor([self.cx, self.cy, self.fx, self.fy]) + + @classmethod + def fx_index(cls): + """TODO(ethan): redo this in a better way. + Ideally we can dynamically grab the focal length parameters depending on + Simple vs. not Simple Pinhole Model. + + Returns: + _type_: _description_ + """ + return 2 + + @classmethod + def fy_index(cls): + """_summary_ + + Returns: + _type_: _description_ + """ + return 3 + + @classmethod + def generate_rays( + cls, + intrinsics: TensorType["num_rays", "num_intrinsics_params"], + camera_to_world: TensorType["num_rays", 3, 4], + coords: TensorType["num_rays", 2], + ) -> RayBundle: + + cx = intrinsics[:, 0:1] + cy = intrinsics[:, 1:2] + fx = intrinsics[:, cls.fx_index() : cls.fx_index() + 1] + fy = intrinsics[:, cls.fy_index() : cls.fy_index() + 1] + + y = coords[:, 0:1] # (num_rays, 1) + x = coords[:, 1:2] # (num_rays, 1) + directions = torch.cat([(x - cx) / fx, -(y - cy) / fy, -torch.ones_like(x)], -1) # (num_rays, 3) + rotation = camera_to_world[:, :3, :3] # (num_rays, 3, 3) + directions = torch.bmm(rotation, directions.unsqueeze(-1)).squeeze(-1) # (num_rays, 3) + directions = normalize(directions, dim=-1) + origins = camera_to_world[:, :3, 3] # (num_rays, 3) + return RayBundle(origins=origins, directions=directions) + + +class SimplePinholeCamera(PinholeCamera): + """Simple Pinhole Camera model.""" + + def __init__(self, cx: float, cy: float, f: float, camera_to_world: Optional[TensorType[3, 4]] = torch.eye(4)[:3]): + super().__init__(cx, cy, f, f, camera_to_world) + + @classmethod + def fx_index(cls): + return 2 + + @classmethod + def fy_index(cls): + return 2 + + +def get_camera_model(num_intrinsics_params): + """Returns the camera model given the specified number of intrinsics parameters. + + Args: + num_intrinsics_params (_type_): _description_ + + Returns: + _type_: _description_ + """ + if num_intrinsics_params == 3: + return SimplePinholeCamera + raise NotImplementedError diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py new file mode 100644 index 000000000..bb0ba5575 --- /dev/null +++ b/mattport/structures/rays.py @@ -0,0 +1,44 @@ +""" +Some ray datastructures. +""" +from dataclasses import dataclass +from typing import Optional +from torchtyping import TensorType + + +@dataclass +class CameraRayBundle: + """_summary_""" + + origins: TensorType["image_height", "image_width", 3] + directions: TensorType["image_height", "image_width", 3] + camera_index: int = None + + +@dataclass +class RayBundle: + """_summary_ + + Returns: + _type_: _description_ + """ + + origins: TensorType["num_rays", 3] + directions: TensorType["num_rays", 3] + camera_indices: Optional[TensorType["num_rays"]] = None + + def to_camera_ray_bundle(self, image_height, image_width) -> CameraRayBundle: + """_summary_ + + Args: + image_height (_type_): _description_ + image_width (_type_): _description_ + + Returns: + CameraRayBundle: _description_ + """ + camera_ray_bundle = CameraRayBundle( + origins=self.origins.view(image_height, image_width, 3), + directions=self.directions.view(image_height, image_width, 3), + ) + return camera_ray_bundle diff --git a/mattport/utils/io.py b/mattport/utils/io.py new file mode 100644 index 000000000..6ca89a126 --- /dev/null +++ b/mattport/utils/io.py @@ -0,0 +1,82 @@ +""" +Input/output utils. +""" + +import json +import os + + +def load_from_json(filename: str): + """_summary_ + + Args: + filename (str): _description_ + + Returns: + _type_: _description_ + """ + assert filename.endswith(".json") + with open(filename, "r", encoding="UTF-8") as file: + return json.load(file) + + +def write_to_json(filename: str, content: dict): + """_summary_ + + Args: + filename (str): _description_ + content (dict): _description_ + """ + assert filename.endswith(".json") + with open(filename, "w", encoding="UTF-8") as file: + json.dump(content, file) + + +def get_git_root(path, dirs=(".git",), default=None): + """_summary_ + "https://stackoverflow.com/questions/22081209/find-the-root-of-the-git-repository-where-the-file-lives + + Args: + path (_type_): _description_ + dirs (tuple, optional): _description_. Defaults to (".git",). + default (_type_, optional): _description_. Defaults to None. + + Returns: + _type_: _description_ + """ + prev, test = None, os.path.abspath(path) + while prev != test: + if any(os.path.isdir(os.path.join(test, directory)) for directory in dirs): + return test + prev, test = test, os.path.abspath(os.path.join(test, os.pardir)) + return default + + +def get_absolute_path(path, proj_root_func=get_git_root): + """ + Returns the full, absolute path. + Relative paths are assumed to start at the repo directory. + """ + if path == "": + return "" + absolute_path = path + if absolute_path[0] != "/": + absolute_path = os.path.join(proj_root_func(path), absolute_path) + return absolute_path + + +def make_dir(filename_or_folder): + """Make the directory for either the filename or folder. + Note that filename_or_folder currently needs to end in / for it to be recognized as a folder. + """ + if filename_or_folder[-1] != "/" and filename_or_folder.find(".") < 0: + folder = os.path.dirname(filename_or_folder + "/") + else: + folder = os.path.dirname(filename_or_folder) + if not os.path.exists(folder): + try: + os.makedirs(folder) + except Exception as e: # pylint: disable=broad-except + print(f"Couldn't create folder: {folder}. Maybe due to a parallel process?") + print(e) + return filename_or_folder diff --git a/scripts/run_dataset_processor.py b/scripts/run_dataset_processor.py new file mode 100644 index 000000000..aa510b4e7 --- /dev/null +++ b/scripts/run_dataset_processor.py @@ -0,0 +1,6 @@ +""" +Script to process different dataset types into our format. +""" + +# TODO(ethan): we should write a script here to convert from other datasets to our formats +# e.g., Blender, LLFF, ShapeNet, DTU diff --git a/tests/nerf/test_encoding.py b/tests/nerf/test_encoding.py new file mode 100644 index 000000000..99f26096c --- /dev/null +++ b/tests/nerf/test_encoding.py @@ -0,0 +1,9 @@ +import pytest + + +def test_positional_encoding(): + assert True + + +def test_lower_case(): + assert "MATTPORT".lower() == "mattport" diff --git a/tests/nerf/test_ray_generator.py b/tests/nerf/test_ray_generator.py new file mode 100644 index 000000000..22ba267cc --- /dev/null +++ b/tests/nerf/test_ray_generator.py @@ -0,0 +1,56 @@ +""" +Test ray generation. +""" +import torch + +from mattport.nerf.dataset.blender_dataset import load_blender_data +from mattport.nerf.modules.ray_generator import RayGenerator +from mattport.structures.rays import RayBundle +from mattport.utils.io import get_absolute_path + + +def test_ray_generator(): + """Test the ray generation works given camera, row, and col indices as input.""" + + # load data + datadir = get_absolute_path("data/lego") + image_filenames, poses, focal_length, image_height, image_width = load_blender_data( + datadir, half_res=False, testskip=1 + ) + 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([focal_length, cx, cy]) + + # Initialize the ray generator with the cameras. + ray_generator = RayGenerator(intrinsics, camera_to_world) + + # Random ray indices. + print("Random ray indexing.") + num_rays = 1024 + ray_indices = torch.floor(torch.rand((num_rays, 3)) * torch.tensor([num_cameras, image_height, image_width])).long() + ray_bundle = ray_generator(ray_indices) + assert isinstance(ray_bundle, RayBundle) + + # TODO(ethan): visualize the RayBundle object + # TODO(ethan): add other test cases? e.g., origin and camera_indices checks + + # check that the origins are correct + c = ray_indices[:, 0] + assert torch.allclose(ray_bundle.origins, camera_to_world[c, :3, 3]) + + # Ray indices for one camera. + coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing="ij") + coords = torch.stack(coords, dim=-1) + coords = coords.view(-1, 2) + num_rays = coords.shape[0] + camera_index = 0 + ray_indices = torch.cat([torch.ones_like(coords[:, :1]) * camera_index, coords], dim=-1).long() + ray_bundle = ray_generator(ray_indices) + + print(ray_bundle.camera_indices) + + # assert False diff --git a/tests/structures/test_cameras.py b/tests/structures/test_cameras.py new file mode 100644 index 000000000..c0d8143f1 --- /dev/null +++ b/tests/structures/test_cameras.py @@ -0,0 +1,12 @@ +import torch + +from mattport.structures.cameras import PinholeCamera +from mattport.structures.rays import RayBundle, CameraRayBundle + + +def test_pinhole_camera(): + pinhole_camera = PinholeCamera(cx=400.0, cy=400.0, fx=10.0, fy=10.0) + camera_ray_bundle = pinhole_camera.generate_camera_rays() + assert isinstance(camera_ray_bundle, CameraRayBundle) + assert torch.allclose(camera_ray_bundle.origins[0], torch.tensor([0.0, 0.0, 0.0])) + # assert False