tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-13 03:14:00 +08:00
Code Activity

Merge branch 'master' into eng/distributed_training

Browse Source
This commit is contained in:
Ethan Weber authored and GitHub committed 2022-04-17 09:40:30 -07:00
2 parents 563dba0bad 621a94e65c
commit e5bc28c4fb
26 files changed
+586 -385
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.github/workflows/code_checks.yml
+27 -23
View File
@@ -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
.gitignore
+4 -2
View File
@@ -166,5 +166,7 @@ outputs/
# Data
data/
# Old
old/
# Misc
old/
temp*
.nfs*
.vscode/settings.json
Vendored
+7 -2
View File
@@ -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"
],
}
README.md
+7 -3
View File
@@ -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
```
environment/README.md
-20
View File
@@ -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
```
environment/jupyter.py
+3
View File
@@ -0,0 +1,3 @@
## https://github.com/jupyterlab/jupyterlab/issues/2049
## Allow access to hidden files
c.ContentsManager.allow_hidden = True
environment/mattport.dockerfile
-75
View File
@@ -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
environment/requirements.txt
+7 -1
View File
@@ -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
torchvision==0.12.0
tqdm
environment/run_jupyter.sh
+1
View File
@@ -0,0 +1 @@
python -m jupyter lab --ip 0.0.0.0 --port 9999 --no-browser --config environments/jupyter.py
mattport/nerf/dataset/base.py
-94
View File
@@ -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?
mattport/nerf/dataset/blender_dataset.py
-55
View File
@@ -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
mattport/nerf/dataset/build.py
-24
View File
@@ -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
mattport/nerf/dataset/image_dataset.py
+91
View File
@@ -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)
mattport/nerf/dataset/utils.py
+100
View File
@@ -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
mattport/nerf/renderers.py
+12 -12
View File
@@ -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
mattport/nerf/trainer.py
+2 -2
View File
@@ -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_"""
mattport/structures/cameras.py
+5 -4
View File
@@ -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)
mattport/structures/rays.py
+33 -10
View File
@@ -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
mattport/viewer/plotly.py
+113
View File
@@ -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
notebooks/notebook_imports.py
+30
View File
@@ -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()
pyproject.toml
+2
View File
@@ -0,0 +1,2 @@
[tool.black]
line-length = 120
tests/nerf/dataloader/configs/blender_lego.yaml
+21
View File
@@ -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
tests/nerf/dataloader/test_ray_generator.py
+75
View File
@@ -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()
tests/nerf/test_ray_generator.py
-56
View File
@@ -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
tests/nerf/test_renderers.py
+40
View File
@@ -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()
tests/structures/test_cameras.py
+6 -2
View File
@@ -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