diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml index 9dd89fe5a..a3e383a33 100644 --- a/.github/workflows/code_checks.yml +++ b/.github/workflows/code_checks.yml @@ -2,36 +2,40 @@ name: Check code for formating and run tests. on: push: - branches: [ master ] + branches: [master] pull_request: - branches: [ master ] + branches: [master] permissions: contents: read jobs: build: - runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 - - name: Set up Python 3.8.12 - uses: actions/setup-python@v3 - with: - python-version: "3.8.12" - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install -r environment/requirements.txt - python setup.py develop - - name: Run Black - uses: psf/black@stable - with: - options: "--check --verbose --line-length=120" - - name: Python Pylint - run: | - 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 --ignore tests/nerf/test_ray_generator.py + - uses: actions/checkout@v3 + - name: Set up Python 3.8.12 + uses: actions/setup-python@v3 + with: + python-version: "3.8.12" + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r environment/requirements.txt + python setup.py develop + - name: Run Black + uses: psf/black@stable + with: + options: "--check --verbose --line-length=120" + - name: Python Pylint + run: | + pylint mattport tests \ + --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 \ + --disable=fixme + - name: Test with pytest + run: | + pytest tests --ignore tests/nerf/dataloader/test_ray_generator.py --typeguard-packages=mattport --torchtyping-patch-typeguard diff --git a/.gitignore b/.gitignore index 06bc11f0b..b8382ed16 100644 --- a/.gitignore +++ b/.gitignore @@ -166,5 +166,7 @@ outputs/ # Data data/ -# Old -old/ \ No newline at end of file +# Misc +old/ +temp* +.nfs* \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json index 3ce315795..751433657 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -4,7 +4,9 @@ "**/__pycache__/": true, "**/.ipynb_checkpoints": true, "**/*.ipynb": true, + "**/old*": true, "**/runs*": true, + "**/temp*": true, }, // "editor.formatOnSave": true, "python.formatting.provider": "black", @@ -18,7 +20,9 @@ "--src=${workspaceFolder}" ], "python.testing.pytestArgs": [ - "tests" + "--ignore=tests/nerf/test_ray_generator.py", + "--typeguard-packages=mattport", + "--torchtyping-patch-typeguard" ], "python.testing.unittestEnabled": false, "python.testing.pytestEnabled": true, @@ -26,6 +30,7 @@ "--generated-members=numpy.*,torch.*,cv2.*,cv.*", "--max-line-length=120", "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W", - "--disable=too-many-arguments" + "--disable=too-many-arguments", + "--disable=fixme" ], } \ No newline at end of file diff --git a/README.md b/README.md index 31735309d..fe12bd2c2 100644 --- a/README.md +++ b/README.md @@ -17,9 +17,6 @@ cd external/goat python setup.py develop # then go back, i.e., `cd /path/to/mattport` -# For using with Jupyter -python -m ipykernel install --user --name mattport --display-name "mattport" - # Install mattport as a library python setup.py develop @@ -31,4 +28,11 @@ pytest tests ``` python scripts/run_train_nerf.py +``` + +# Setting up Jupyter + +``` +python -m jupyter lab build +bash environments/run_jupyter.sh ``` \ No newline at end of file diff --git a/environment/README.md b/environment/README.md deleted file mode 100644 index bf98aae6b..000000000 --- a/environment/README.md +++ /dev/null @@ -1,20 +0,0 @@ -# Environment - -# Installing docker - -#### Troubleshooting - -``` -# https://forums.docker.com/t/failing-to-start-dockerd-failed-to-create-nat-chain-docker/78269 -sudo update-alternatives --set iptables /usr/sbin/iptables-legacy -sudo update-alternatives --set ip6tables /usr/sbin/ip6tables-legacy -sudo service docker start - -VERSION_STRING=18.03.1~ce~3-0~ubuntu -sudo apt-get install docker-ce=$VERSION_STRING docker-ce-cli=$VERSION_STRING containerd.io - -# https://github.com/docker/for-linux/issues/123#issuecomment-346546953 -sudo ip link add name docker0 type bridge -sudo ip addr add dev docker0 172.17.0.1/16 -sudo dockerd --debug --iptables=false -``` \ No newline at end of file diff --git a/environment/jupyter.py b/environment/jupyter.py new file mode 100644 index 000000000..70dd75764 --- /dev/null +++ b/environment/jupyter.py @@ -0,0 +1,3 @@ +## https://github.com/jupyterlab/jupyterlab/issues/2049 +## Allow access to hidden files +c.ContentsManager.allow_hidden = True diff --git a/environment/mattport.dockerfile b/environment/mattport.dockerfile deleted file mode 100644 index aa8bad9b1..000000000 --- a/environment/mattport.dockerfile +++ /dev/null @@ -1,75 +0,0 @@ -FROM nvidia/cuda:10.2-devel-ubuntu18.04 - -# Set the shell to /bin/bash instead of /bin/sh so conda works correctly. -# https://towardsdatascience.com/conda-pip-and-docker-ftw-d64fe638dc45 -SHELL [ "/bin/bash", "--login", "-c" ] - -# Prepare and empty machine for building -RUN apt-get update && apt-get install -y \ - git \ - cmake \ - vim \ - build-essential \ - libboost-program-options-dev \ - libboost-filesystem-dev \ - libboost-graph-dev \ - libboost-system-dev \ - libboost-test-dev \ - libeigen3-dev \ - libsuitesparse-dev \ - libfreeimage-dev \ - libgoogle-glog-dev \ - libgflags-dev \ - libglew-dev \ - qtbase5-dev \ - libqt5opengl5-dev \ - libcgal-dev \ - libcgal-qt5-dev \ - wget - -# Build and install ceres solver -RUN apt-get -y install \ - libatlas-base-dev \ - libsuitesparse-dev -RUN git clone https://github.com/ceres-solver/ceres-solver.git --branch 1.14.0 -RUN cd ceres-solver && \ - mkdir build && \ - cd build && \ - cmake .. -DBUILD_TESTING=OFF -DBUILD_EXAMPLES=OFF && \ - make -j4 && \ - make install - -# Clear space # https://askubuntu.com/questions/179955/var-lib-apt-lists-is-huge -RUN rm -rf /var/lib/apt/lists/* - -# Build and install COLMAP - -# Note: This Dockerfile has been tested using COLMAP pre-release 3.6. -# Later versions of COLMAP (which will be automatically cloned as default) may -# have problems using the environment described thus far. If you encounter -# problems and want to install the tested release, then uncomment the branch -# specification in the line below -RUN git clone https://github.com/colmap/colmap.git #--branch 3.6 - -RUN cd colmap && \ - git checkout dev && \ - mkdir build && \ - cd build && \ - cmake .. && \ - make -j4 && \ - make install - -# Install Conda -ENV PATH="/root/miniconda3/bin:${PATH}" -ARG PATH="/root/miniconda3/bin:${PATH}" -RUN wget \ - https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \ - && mkdir /root/.conda \ - && bash Miniconda3-latest-Linux-x86_64.sh -b \ - && rm -f Miniconda3-latest-Linux-x86_64.sh -RUN conda --version - -# Install friends conda environment -COPY ./run_install_friends_conda.sh /tmp/run_install_friends_conda.sh -COPY ./requirements_friends.txt /tmp/requirements_friends.txt -RUN yes "Y" | bash /tmp/run_install_friends_conda.sh diff --git a/environment/requirements.txt b/environment/requirements.txt index bb93663fb..9110d4604 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -1,8 +1,14 @@ black hydra-core imageio==2.16.1 +ipywidgets>=7.6 +jupyterlab +kaleido +mediapy +plotly==5.7.0 pylint==2.13.4 pytest torch==1.11.0 torchtyping -torchvision==0.12.0 \ No newline at end of file +torchvision==0.12.0 +tqdm \ No newline at end of file diff --git a/environment/run_jupyter.sh b/environment/run_jupyter.sh new file mode 100644 index 000000000..0164bc6fd --- /dev/null +++ b/environment/run_jupyter.sh @@ -0,0 +1 @@ +python -m jupyter lab --ip 0.0.0.0 --port 9999 --no-browser --config environments/jupyter.py \ No newline at end of file diff --git a/mattport/nerf/dataset/base.py b/mattport/nerf/dataset/base.py deleted file mode 100644 index d93ca1133..000000000 --- a/mattport/nerf/dataset/base.py +++ /dev/null @@ -1,94 +0,0 @@ -""" -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 deleted file mode 100644 index 53b5fb9f0..000000000 --- a/mattport/nerf/dataset/blender_dataset.py +++ /dev/null @@ -1,55 +0,0 @@ -""" -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 deleted file mode 100644 index 401f3ebc4..000000000 --- a/mattport/nerf/dataset/build.py +++ /dev/null @@ -1,24 +0,0 @@ -""" -Builds the dataset -""" -import importlib -import logging - -from mattport.nerf.dataset.base import Dataset - - -def build_dataset(config: dict) -> Dataset: - """Returns the dataset according to the specified config. - # TODO(ethan): change the type from dict to be a dataclass - - Args: - config (dict): _description_ - """ - - # TODO(ethan): better handle this with configs - 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/image_dataset.py b/mattport/nerf/dataset/image_dataset.py new file mode 100644 index 000000000..2ef3ebfeb --- /dev/null +++ b/mattport/nerf/dataset/image_dataset.py @@ -0,0 +1,91 @@ +""" +Some dataset code. +""" + +from typing import List + +import numpy as np +import torch +from PIL import Image +from torch.utils.data import default_collate + +from mattport.nerf.dataset.utils import DotDict + + +class ImageDataset(torch.utils.data.Dataset): + """Dataset that returns images.""" + + def __init__(self, image_filenames: List[str], downscale_factor: float = 1.0): + super().__init__() + self.image_filenames = image_filenames + self.downscale_factor = downscale_factor + + def __len__(self): + return len(self.image_filenames) + + def get_image(self, image_idx: int) -> np.uint8: + """Returns the image. + + Args: + image_idx (int): The image index in the dataset. + + Returns: + np.uint8: an image of shape (H, W, 3 or 4) + """ + image_filename = self.image_filenames[image_idx] + image = np.array(Image.open(image_filename)) # shape is (h, w, 3 or 4) + assert len(image.shape) == 3 + assert image.shape[2] in [3, 4], f"Image shape of {image.shape} is in correct." + return image + + def __getitem__(self, image_idx): + # TODO(ethan): we should add code to deal with varying resolution and/or having a maximum pixel resolution + # e.g., we should be able to downscale the images or set a max height, width, etc. + + # the image might be RGB or RGBA, so we separate it + original_image = self.get_image(image_idx).astype("float32") / 255.0 + image = original_image[:, :, :3] + num_channels = original_image.shape[2] + if num_channels == 4: + mask = original_image[:, :, 3] + elif num_channels == 3: + mask = np.ones_like(original_image[:, :, 0]) + else: + raise ValueError(f"Image shape of {image.shape} is in correct.") + + return { + "image_idx": image_idx, + "image": image, # the pixels + "mask": mask, + } + + +def collate_batch(batch, num_rays_per_batch, keep_full_image: bool = False): + """_summary_ + + Args: + batch (_type_): _description_ + """ + # TODO(ethan): handle sampling even when in different processes + # we don't want the same samples for all spawned processed when + # using distributed training + + batch = default_collate(batch) + num_images, image_height, image_width, _ = batch["image"].shape + indices = torch.floor( + torch.rand((num_rays_per_batch, 3)) * torch.tensor([num_images, image_height, image_width]) + ).long() + c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)] + pixels = batch["image"][c, y, x] + mask = batch["mask"][c, y, x] + assert pixels.shape == (num_rays_per_batch, 3), pixels.shape + + collated_batch = { + "indices": indices, + "pixels": pixels, + "mask": mask, + } + if keep_full_image: + collated_batch["image"] = batch["image"] + + return DotDict(collated_batch) diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py new file mode 100644 index 000000000..ccb1a3f49 --- /dev/null +++ b/mattport/nerf/dataset/utils.py @@ -0,0 +1,100 @@ +""" +For loading the blender dataset format. +""" + +import os +from dataclasses import dataclass +from typing import List + +import imageio +import numpy as np +import torch + +from mattport.utils.io import load_from_json + + +class DotDict(dict): + """ + dot.notation access to dictionary attributes + """ + + def __getattr__(self, attr): + return self[attr] + + __setattr__ = dict.__setitem__ + __delattr__ = dict.__delitem__ + + +@dataclass +class DatasetInputs: + """Dataset inputs for the image dataset and the ray generator.""" + + image_filenames: List[str] + downscale_factor: float = 1.0 + intrinsics: torch.tensor = None + camera_to_world: torch.tensor = None + + +def load_blender_data(basedir, downscale_factor=1.0, split="train"): + """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_ + downscale_factor (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 + 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, + intrinsics=intrinsics, + camera_to_world=camera_to_world, + ) + + return dataset_inputs + + +def get_dataset_inputs(data_directory: str, dataset_type: str, downscale_factor: float = 1.0): + """Returns the dataset inputs, which will be used with an ImageDataset and RayGenerator. + + Args: + data_directory (str): _description_ + dataset_type (str): _description_ + downscale_factor (float, optional): _description_. Defaults to 1.0. + + Returns: + DatasetInputs: The inputs needed for generating rays. + """ + if dataset_type == "blender": + dataset_inputs = load_blender_data(data_directory, downscale_factor=downscale_factor) + else: + raise NotImplementedError() + + return dataset_inputs diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index d319ea74e..1e78e04aa 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -13,7 +13,7 @@ from mattport.nerf.field_modules import field_outputs class Renderer(nn.Module): """Base Renderer. Intended to be subclassed""" - def required_render_heads(self) -> Set[field_outputs.FieldOutput]: + def required_field_outputs(self) -> Set[field_outputs.FieldOutput]: """ Returns: Set[RenderHead]: RenderHeads required for this renderer. @@ -22,13 +22,13 @@ class Renderer(nn.Module): def forward( self, - render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]], + field_output_dict: Dict["str", TensorType[..., "num_samples", -1]], deltas: TensorType[..., "num_samples"], ) -> TensorType[..., "out_dim"]: """Composite samples along ray and render image Args: - in_tensor (TensorType[..., "num_samples", "features"]): feature for each sample + in_tensor (TensorType[..., "num_samples", -1]): feature for each sample deltas (TensorType[..., "num_samples"]): depth of each sample Returns: @@ -45,33 +45,33 @@ class RGB(Renderer): self.rgb_name = rgb_name self.density_name = density_name - def required_render_heads(self) -> Set[field_outputs.FieldOutput]: + def required_field_outputs(self) -> Set[field_outputs.FieldOutput]: return set(field_outputs.RGBFieldOutput, field_outputs.DensityFieldOutput) def forward( self, - render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]], + field_output_dict: Dict[str, TensorType[..., "num_samples", -1]], deltas: TensorType[..., "num_samples"], ) -> TensorType[..., "out_dim"]: """Composite samples along ray and render color image Args: - in_tensor (TensorType[..., "num_samples", 3]): RGB for each sample + in_tensor (TensorType[..., "num_samples", -1]): RGB for each sample deltas (TensorType[..., "num_samples"]): depth influence for each sample Returns: TensorType[..., "out_dim"]: Composited RGB ray """ - - delta_density = deltas * render_head_outputs[self.density_name] + delta_density = deltas * field_output_dict[self.density_name][..., 0] alphas = 1 - torch.exp(-delta_density) - transmittance = torch.cat(0, torch.cumsum(delta_density[..., :-1], dim=-1), dim=-1) + transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) + transmittance = torch.cat([torch.zeros((*transmittance.shape[:1], 1)), transmittance], axis=-1) transmittance = torch.exp(-transmittance) # [..., "num_samples"] weights = alphas * transmittance # [..., "num_samples"] - rgb = torch.sum(weights[..., None] * render_head_outputs[self.rgb_name], dim=-2) + rgb = torch.sum(weights[..., None] * field_output_dict[self.rgb_name], dim=-2) return rgb @@ -83,12 +83,12 @@ class Depth(Renderer): Renderer (_type_): _description_ """ - def required_render_heads(self) -> Set[field_outputs.FieldOutput]: + def required_field_outputs(self) -> Set[field_outputs.FieldOutput]: return set(field_outputs.DensityFieldOutput) def forward( self, - render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]], + field_output_dict: Dict["str", TensorType[..., "num_samples", -1]], deltas: TensorType[..., "num_samples"], ) -> TensorType[..., "out_dim"]: raise NotImplementedError diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 3c80c884b..b65f01062 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -24,8 +24,8 @@ class Trainer: def setup_dataset(self): """_summary_""" - self.train_dataset = build_dataset(self.config.dataset) - # self.test_dataset = build_dataset(self.config.dataset) + self.train_dataset = None + self.test_dataset = None def setup_graph(self): """_summary_""" diff --git a/mattport/structures/cameras.py b/mattport/structures/cameras.py index 88225f221..0eca34707 100644 --- a/mattport/structures/cameras.py +++ b/mattport/structures/cameras.py @@ -62,7 +62,7 @@ class Camera: """_summary_""" return - def get_image_coords(self) -> TensorType["image_height", "image_width", 2]: + def get_image_coords(self, pixel_offset: float = 0.5) -> TensorType["image_height", "image_width", 2]: """_summary_ Returns: @@ -71,7 +71,7 @@ class Camera: 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) + image_coords = torch.stack(image_coords, dim=-1) + pixel_offset # stored as (y, x) coordinates return image_coords @classmethod @@ -179,12 +179,13 @@ class PinholeCamera(Camera): 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 = torch.sum( + directions[:, None, :] * rotation, dim=-1 + ) # (num_rays, 1, 3) * (num_rays, 3, 3) -> (num_rays, 3) directions = normalize(directions, dim=-1) origins = camera_to_world[:, :3, 3] # (num_rays, 3) return RayBundle(origins=origins, directions=directions) diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index bb0ba5575..6ba272c21 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -1,20 +1,13 @@ """ Some ray datastructures. """ +import random 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_ @@ -27,7 +20,7 @@ class RayBundle: directions: TensorType["num_rays", 3] camera_indices: Optional[TensorType["num_rays"]] = None - def to_camera_ray_bundle(self, image_height, image_width) -> CameraRayBundle: + def to_camera_ray_bundle(self, image_height, image_width) -> "CameraRayBundle": """_summary_ Args: @@ -42,3 +35,33 @@ class RayBundle: directions=self.directions.view(image_height, image_width, 3), ) return camera_ray_bundle + + def __len__(self): + num_rays = self.origins.shape[0] + return num_rays + + def sample(self, num_rays: int): + """Returns a RayBundle as a subset of rays. + + Args: + num_rays (int): + + Returns: + RayBundle: _description_ + """ + assert num_rays <= len(self) + indices = random.sample(range(len(self)), k=num_rays) + return RayBundle( + origins=self.origins[indices], + directions=self.directions[indices], + camera_indices=self.camera_indices[indices], + ) + + +@dataclass +class CameraRayBundle: + """_summary_""" + + origins: TensorType["image_height", "image_width", 3] + directions: TensorType["image_height", "image_width", 3] + camera_index: int = None diff --git a/mattport/viewer/plotly.py b/mattport/viewer/plotly.py new file mode 100644 index 000000000..0f6ed1cc1 --- /dev/null +++ b/mattport/viewer/plotly.py @@ -0,0 +1,113 @@ +""" +Visualization code for plotly. +""" + +import numpy as np +import plotly.graph_objects as go +import torch + +from mattport.structures.rays import RayBundle + + +def color_str(color): + """Plotly color string. + + Args: + color: list [r, g, b] in [0, 1] range + + Returns: + str: plotly-formatted color string + """ + color = list((np.array(color) * 255.0).astype("int")) + return f"""rgb({color[0]}, {color[1]}, {color[2]})""" + + +def get_line_segments_from_lines( + lines, + color=color_str((1, 0, 0)), + marker_color=color_str((1, 0, 0)), + colors=None, + draw_marker=True, + draw_line=True, + marker_size=4, + line_width=10, +): + """Returns a list of Scatter3D objects for creating lines with plotly. + # TODO(ethan): make this function more efficient instead of having a list of objects. + + Args: + lines (_type_): list of [np.array(2, 3), ...], line is defined by a (2, 3) array + color (_type_, optional): _description_. Defaults to color_str((1, 0, 0)). + marker_color (_type_, optional): _description_. Defaults to color_str((1, 0, 0)). + colors (_type_, optional): _description_. Defaults to None. + draw_marker (bool, optional): _description_. Defaults to True. + draw_line (bool, optional): _description_. Defaults to True. + marker_size (int, optional): _description_. Defaults to 4. + line_width (int, optional): _description_. Defaults to 10. + + Returns: + _type_: _description_ + """ + data = [] + for idx, line in enumerate(lines): + thiscolor = color if draw_line else "rgba(0, 0, 0, 0)" + if colors is not None: + marker_color = colors[idx] + thiscolor = colors[idx] + data.append( + go.Scatter3d( + x=line[:, 0], + y=line[:, 1], + z=line[:, 2], + showlegend=False, + marker=dict( + size=marker_size, + color=marker_color, + colorscale="Viridis", + ) + if draw_marker + else dict(color="rgba(0, 0, 0, 0)"), + line=dict(color=thiscolor, width=line_width), + ) + ) + return data + + +def visualize_dataset(camera_origins, ray_bundle: RayBundle): + """Visualize a dataset with plotly using our cameras and generated rays.""" + + skip = 1 + size = 8 + assert len(ray_bundle) < 500, "Let's not break plotly by plotting too many rays!" + + data = [] + data += [ + go.Scatter3d( + x=camera_origins[::skip, 0], + y=camera_origins[::skip, 1], + z=camera_origins[::skip, 2], + mode="markers", + name="camera origins", + marker=dict(color="rgba(0, 0, 0, 1)", size=size), + ) + ] + + length = 2.0 + lines = torch.stack( + [ray_bundle.origins, ray_bundle.origins + ray_bundle.directions * length], dim=1 + ) # (num_rays, 2, 3) + + data += get_line_segments_from_lines(lines) + + layout = go.Layout( + autosize=False, + width=1000, + height=1000, + margin=go.layout.Margin(l=50, r=50, b=100, t=100, pad=4), + scene=go.layout.Scene( + aspectmode="data", + camera=dict(up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=0), eye=dict(x=1.25, y=1.25, z=1.25)), + ), + ) + fig = go.Figure(data=data, layout=layout) + return fig diff --git a/notebooks/notebook_imports.py b/notebooks/notebook_imports.py new file mode 100644 index 000000000..b71f57234 --- /dev/null +++ b/notebooks/notebook_imports.py @@ -0,0 +1,30 @@ +# image stuff +import cv2 +import imageio +import mediapy as media + +# ipython stuff +from IPython import get_ipython + +# plotly stuff +import plotly.graph_objects as go + + +def setup_ipynb(): + """ + https://stackoverflow.com/questions/15411967/how-can-i-check-if-code-is-executed-in-the-ipython-notebook + https://stackoverflow.com/questions/35595766/matplotlib-line-magic-causes-syntaxerror-in-python-script + This gets reference to the InteractiveShell instance + """ + try: + from IPython import get_ipython + + get_ipython().run_line_magic("load_ext", "autoreload") + get_ipython().run_line_magic("autoreload", "2") + get_ipython().run_line_magic("matplotlib", "inline") + return True + except: + return False + + +setup_ipynb() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..e34796ec5 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,2 @@ +[tool.black] +line-length = 120 \ No newline at end of file diff --git a/tests/nerf/dataloader/configs/blender_lego.yaml b/tests/nerf/dataloader/configs/blender_lego.yaml new file mode 100644 index 000000000..a4793a94e --- /dev/null +++ b/tests/nerf/dataloader/configs/blender_lego.yaml @@ -0,0 +1,21 @@ +# dataset options +dataset: + data_directory: data/lego + dataset_type: blender + downscale_factor: 1.0 + +# dataloader options +dataloader: + num_workers: 0 + shuffle: true + num_images_to_sample_from: 8 + num_rays_per_batch: 1024 + +# ray generation module and its option +network: + ray_generator: + class_name: RayGenerator + inputs: ['cxy'] # camera, x, and y indices + meta_data: + in_dim: 0 + learning_rate: 0.0 \ No newline at end of file diff --git a/tests/nerf/dataloader/test_ray_generator.py b/tests/nerf/dataloader/test_ray_generator.py new file mode 100644 index 000000000..a03d1c970 --- /dev/null +++ b/tests/nerf/dataloader/test_ray_generator.py @@ -0,0 +1,75 @@ +""" +Code to test ray generation. +""" + +import imageio +import torch +from hydra import compose, initialize +from torch.utils.data import DataLoader +from tqdm import tqdm + +from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch +from mattport.nerf.dataset.utils import get_dataset_inputs +from mattport.nerf.field_modules.ray_generator import RayGenerator +from mattport.utils.io import get_absolute_path +from mattport.viewer.plotly import visualize_dataset + + +def visualize_batch(batch): + """Returns an image from a batch.""" + assert "image" in batch, "You must call collate_batch with keep_full_image=True for this to work." + # set the color of the sampled rays + c, y, x = [i.flatten() for i in torch.split(batch.indices, 1, dim=-1)] + batch.image[c, y, x] = 1.0 + + # batch.image is num_images, h, w, 3 + images = torch.split(batch.image, 1, dim=0) + image_list = [image[0] for image in images] + image = torch.cat(image_list, dim=1) # cat along the width dimension + image = (image * 255.0).to(torch.uint8) + return image + + +def test_dataloader(): + """Testing for the dataloader from input dataset parameters to rays.""" + + with initialize(config_path="configs"): + cfg = compose(config_name="blender_lego.yaml") + + cfg.dataset.data_directory = get_absolute_path( + cfg.dataset.data_directory + ) # TODO(ethan): create a helper function for this + dataset_inputs = get_dataset_inputs(**cfg.dataset) + dataset = ImageDataset( + image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor + ) + + dataloader = DataLoader( + dataset, + batch_size=cfg.dataloader.num_images_to_sample_from, + collate_fn=lambda batch: collate_batch(batch, cfg.dataloader.num_rays_per_batch, keep_full_image=True), + num_workers=cfg.dataloader.num_workers, + shuffle=True, + ) + dataloader_iter = iter(dataloader) + + ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world) + + num_batches = 10 + for _ in tqdm(range(num_batches)): + batch = next(dataloader_iter) + ray_bundle = ray_generator.forward(batch.indices) + + # visualize the batch + image = visualize_batch(batch) + imageio.imwrite("temp0.png", image) + + # visualize the RayBundle + small_ray_bundle = ray_bundle.sample(100) + fig = visualize_dataset(camera_origins=dataset_inputs.camera_to_world[:, :3, 3], ray_bundle=small_ray_bundle) + fig.write_image("temp1.png") + fig.write_html("temp1.html") + + +if __name__ == "__main__": + test_dataloader() diff --git a/tests/nerf/test_ray_generator.py b/tests/nerf/test_ray_generator.py deleted file mode 100644 index e9d3bb36b..000000000 --- a/tests/nerf/test_ray_generator.py +++ /dev/null @@ -1,56 +0,0 @@ -""" -Test ray generation. -""" -import torch - -from mattport.nerf.dataset.blender_dataset import load_blender_data -from mattport.nerf.field_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/nerf/test_renderers.py b/tests/nerf/test_renderers.py new file mode 100644 index 000000000..e6897dad9 --- /dev/null +++ b/tests/nerf/test_renderers.py @@ -0,0 +1,40 @@ +""" +Test renderers +""" +import pytest +import torch + +from mattport.nerf import renderers + + +def test_rgb_renderer(): + """Test RGB volumetric rendering""" + rgb_name = "rgb" + density_name = "density" + num_samples = 10 + + rgb = torch.ones((3, num_samples, 3)) + densities = torch.ones((3, num_samples, 1)) + deltas = torch.ones((3, num_samples)) + + rgb_renderer = renderers.RGB(rgb_name, density_name) + + assert rgb_renderer.required_field_outputs is not None + + field_output_dict = {rgb_name: rgb, density_name: densities} + out = rgb_renderer(field_output_dict, deltas) + assert torch.max(out) > 0.9 + + field_output_dict = {rgb_name: rgb * 0, density_name: densities} + out = rgb_renderer(field_output_dict, deltas) + assert torch.max(out) == pytest.approx(0) + + field_output_dict = {rgb_name: rgb, density_name: densities * 0} + out = rgb_renderer(field_output_dict, deltas) + assert torch.max(out) == pytest.approx(0) + + # TODO Tancik: Add more precise tests + + +if __name__ == "__main__": + test_rgb_renderer() diff --git a/tests/structures/test_cameras.py b/tests/structures/test_cameras.py index c0d8143f1..e0d841a6b 100644 --- a/tests/structures/test_cameras.py +++ b/tests/structures/test_cameras.py @@ -1,12 +1,16 @@ +""" +Test the camera classes. +""" + import torch from mattport.structures.cameras import PinholeCamera -from mattport.structures.rays import RayBundle, CameraRayBundle +from mattport.structures.rays import CameraRayBundle def test_pinhole_camera(): + """Test that the pinhole camera model works.""" 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