Skip to content

GitLab

Unverified Commit 82fd69c7 authored by Jingchen Ye's avatar Jingchen Ye Committed by GitHub
Browse files

Updates on examples (#174) 

* move torch data transfer into dataloader

* Update README

* use args.data_root

* Remove redundant check

* Fix isort

* Fix black
parent 15330b4c
Hide whitespace changes
Inline Side-by-side
Showing with 54 additions and 50 deletions
README.md
View file @ 82fd69c7
......@@ -132,7 +132,7 @@ optimizer.step()
## Examples:
Before running those example scripts, please check the script about which dataset it is needed, and download the dataset first.
Before running those example scripts, please check the script about which dataset is needed, and download the dataset first. You could use `--data_root` to specify the path.
```bash
# clone the repo with submodules.
......
examples/datasets/dnerf_synthetic.py
View file @ 82fd69c7
......@@ -86,6 +86,7 @@ class SubjectLoader(torch.utils.data.Dataset):
near: float = None,
far: float = None,
batch_over_images: bool = True,
device: str = "cuda:0",
):
super().__init__()
assert split in self.SPLITS, "%s" % split
......@@ -106,11 +107,15 @@ class SubjectLoader(torch.utils.data.Dataset):
self.focal,
self.timestamps,
) = _load_renderings(root_fp, subject_id, split)
self.images = torch.from_numpy(self.images).to(torch.uint8)
self.camtoworlds = torch.from_numpy(self.camtoworlds).to(torch.float32)
self.timestamps = torch.from_numpy(self.timestamps).to(torch.float32)[
:, None
]
self.images = torch.from_numpy(self.images).to(device).to(torch.uint8)
self.camtoworlds = (
torch.from_numpy(self.camtoworlds).to(device).to(torch.float32)
)
self.timestamps = (
torch.from_numpy(self.timestamps)
.to(device)
.to(torch.float32)[:, None]
)
self.K = torch.tensor(
[
[self.focal, 0, self.WIDTH / 2.0],
......@@ -118,6 +123,7 @@ class SubjectLoader(torch.utils.data.Dataset):
[0, 0, 1],
],
dtype=torch.float32,
device=device,
) # (3, 3)
assert self.images.shape[1:3] == (self.HEIGHT, self.WIDTH)
......
examples/datasets/nerf_360_v2.py
View file @ 82fd69c7
......@@ -169,6 +169,7 @@ class SubjectLoader(torch.utils.data.Dataset):
far: float = None,
batch_over_images: bool = True,
factor: int = 1,
device: str = "cuda:0",
):
super().__init__()
assert split in self.SPLITS, "%s" % split
......@@ -186,9 +187,11 @@ class SubjectLoader(torch.utils.data.Dataset):
self.images, self.camtoworlds, self.K = _load_colmap(
root_fp, subject_id, split, factor
)
self.images = torch.from_numpy(self.images).to(torch.uint8)
self.camtoworlds = torch.from_numpy(self.camtoworlds).to(torch.float32)
self.K = torch.tensor(self.K).to(torch.float32)
self.images = torch.from_numpy(self.images).to(device).to(torch.uint8)
self.camtoworlds = (
torch.from_numpy(self.camtoworlds).to(device).to(torch.float32)
)
self.K = torch.tensor(self.K).to(device).to(torch.float32)
self.height, self.width = self.images.shape[1:3]
def __len__(self):
......
examples/datasets/nerf_synthetic.py
View file @ 82fd69c7
......@@ -79,6 +79,7 @@ class SubjectLoader(torch.utils.data.Dataset):
near: float = None,
far: float = None,
batch_over_images: bool = True,
device: str = "cuda:0",
):
super().__init__()
assert split in self.SPLITS, "%s" % split
......@@ -109,8 +110,10 @@ class SubjectLoader(torch.utils.data.Dataset):
self.images, self.camtoworlds, self.focal = _load_renderings(
root_fp, subject_id, split
)
self.images = torch.from_numpy(self.images).to(torch.uint8)
self.camtoworlds = torch.from_numpy(self.camtoworlds).to(torch.float32)
self.images = torch.from_numpy(self.images).to(device).to(torch.uint8)
self.camtoworlds = (
torch.from_numpy(self.camtoworlds).to(device).to(torch.float32)
)
self.K = torch.tensor(
[
[self.focal, 0, self.WIDTH / 2.0],
......@@ -118,6 +121,7 @@ class SubjectLoader(torch.utils.data.Dataset):
[0, 0, 1],
],
dtype=torch.float32,
device=device,
) # (3, 3)
assert self.images.shape[1:3] == (self.HEIGHT, self.WIDTH)
......
examples/train_mlp_dnerf.py
View file @ 82fd69c7
......@@ -4,7 +4,7 @@ Copyright (c) 2022 Ruilong Li, UC Berkeley.
import argparse
import math
import os
import pathlib
import time
import imageio
......@@ -24,6 +24,12 @@ if __name__ == "__main__":
set_random_seed(42)
parser = argparse.ArgumentParser()
parser.add_argument(
"--data_root",
type=str,
default=str(pathlib.Path.cwd() / "data/dnerf"),
help="the root dir of the dataset",
)
parser.add_argument(
"--train_split",
type=str,
......@@ -91,31 +97,22 @@ if __name__ == "__main__":
gamma=0.33,
)
# setup the dataset
data_root_fp = "/home/ruilongli/data/dnerf/"
target_sample_batch_size = 1 << 16
grid_resolution = 128
train_dataset = SubjectLoader(
subject_id=args.scene,
root_fp=data_root_fp,
root_fp=args.data_root,
split=args.train_split,
num_rays=target_sample_batch_size // render_n_samples,
)
train_dataset.images = train_dataset.images.to(device)
train_dataset.camtoworlds = train_dataset.camtoworlds.to(device)
train_dataset.K = train_dataset.K.to(device)
train_dataset.timestamps = train_dataset.timestamps.to(device)
test_dataset = SubjectLoader(
subject_id=args.scene,
root_fp=data_root_fp,
root_fp=args.data_root,
split="test",
num_rays=None,
)
test_dataset.images = test_dataset.images.to(device)
test_dataset.camtoworlds = test_dataset.camtoworlds.to(device)
test_dataset.K = test_dataset.K.to(device)
test_dataset.timestamps = test_dataset.timestamps.to(device)
occupancy_grid = OccupancyGrid(
roi_aabb=args.aabb,
......@@ -191,7 +188,7 @@ if __name__ == "__main__":
f"n_rendering_samples={n_rendering_samples:d} | num_rays={len(pixels):d} |"
)
if step >= 0 and step % max_steps == 0 and step > 0:
if step > 0 and step % max_steps == 0:
# evaluation
radiance_field.eval()
......
examples/train_mlp_nerf.py
View file @ 82fd69c7
......@@ -4,7 +4,7 @@ Copyright (c) 2022 Ruilong Li, UC Berkeley.
import argparse
import math
import os
import pathlib
import time
import imageio
......@@ -23,6 +23,12 @@ if __name__ == "__main__":
set_random_seed(42)
parser = argparse.ArgumentParser()
parser.add_argument(
"--data_root",
type=str,
default=str(pathlib.Path.cwd() / "data/nerf_synthetic"),
help="the root dir of the dataset",
)
parser.add_argument(
"--train_split",
type=str,
......@@ -112,7 +118,6 @@ if __name__ == "__main__":
if args.scene == "garden":
from datasets.nerf_360_v2 import SubjectLoader
data_root_fp = "/home/ruilongli/data/360_v2/"
target_sample_batch_size = 1 << 16
train_dataset_kwargs = {"color_bkgd_aug": "random", "factor": 4}
test_dataset_kwargs = {"factor": 4}
......@@ -120,32 +125,24 @@ if __name__ == "__main__":
else:
from datasets.nerf_synthetic import SubjectLoader
data_root_fp = "/home/ruilongli/data/nerf_synthetic/"
target_sample_batch_size = 1 << 16
grid_resolution = 128
train_dataset = SubjectLoader(
subject_id=args.scene,
root_fp=data_root_fp,
root_fp=args.data_root,
split=args.train_split,
num_rays=target_sample_batch_size // render_n_samples,
**train_dataset_kwargs,
)
train_dataset.images = train_dataset.images.to(device)
train_dataset.camtoworlds = train_dataset.camtoworlds.to(device)
train_dataset.K = train_dataset.K.to(device)
test_dataset = SubjectLoader(
subject_id=args.scene,
root_fp=data_root_fp,
root_fp=args.data_root,
split="test",
num_rays=None,
**test_dataset_kwargs,
)
test_dataset.images = test_dataset.images.to(device)
test_dataset.camtoworlds = test_dataset.camtoworlds.to(device)
test_dataset.K = test_dataset.K.to(device)
occupancy_grid = OccupancyGrid(
roi_aabb=args.aabb,
......@@ -217,7 +214,7 @@ if __name__ == "__main__":
f"n_rendering_samples={n_rendering_samples:d} | num_rays={len(pixels):d} |"
)
if step >= 0 and step % max_steps == 0 and step > 0:
if step > 0 and step % max_steps == 0:
# evaluation
radiance_field.eval()
......
examples/train_ngp_nerf.py
View file @ 82fd69c7
......@@ -4,7 +4,7 @@ Copyright (c) 2022 Ruilong Li, UC Berkeley.
import argparse
import math
import os
import pathlib
import time
import imageio
......@@ -23,6 +23,12 @@ if __name__ == "__main__":
set_random_seed(42)
parser = argparse.ArgumentParser()
parser.add_argument(
"--data_root",
type=str,
default=str(pathlib.Path.cwd() / "data"),
help="the root dir of the dataset",
)
parser.add_argument(
"--train_split",
type=str,
......@@ -87,7 +93,6 @@ if __name__ == "__main__":
if args.unbounded:
from datasets.nerf_360_v2 import SubjectLoader
data_root_fp = "/home/ruilongli/data/360_v2/"
target_sample_batch_size = 1 << 20
train_dataset_kwargs = {"color_bkgd_aug": "random", "factor": 4}
test_dataset_kwargs = {"factor": 4}
......@@ -95,32 +100,24 @@ if __name__ == "__main__":
else:
from datasets.nerf_synthetic import SubjectLoader
data_root_fp = "/home/ruilongli/data/nerf_synthetic/"
target_sample_batch_size = 1 << 18
grid_resolution = 128
train_dataset = SubjectLoader(
subject_id=args.scene,
root_fp=data_root_fp,
root_fp=args.data_root,
split=args.train_split,
num_rays=target_sample_batch_size // render_n_samples,
**train_dataset_kwargs,
)
train_dataset.images = train_dataset.images.to(device)
train_dataset.camtoworlds = train_dataset.camtoworlds.to(device)
train_dataset.K = train_dataset.K.to(device)
test_dataset = SubjectLoader(
subject_id=args.scene,
root_fp=data_root_fp,
root_fp=args.data_root,
split="test",
num_rays=None,
**test_dataset_kwargs,
)
test_dataset.images = test_dataset.images.to(device)
test_dataset.camtoworlds = test_dataset.camtoworlds.to(device)
test_dataset.K = test_dataset.K.to(device)
if args.auto_aabb:
camera_locs = torch.cat(
......@@ -260,7 +257,7 @@ if __name__ == "__main__":
f"n_rendering_samples={n_rendering_samples:d} | num_rays={len(pixels):d} |"
)
if step >= 0 and step % max_steps == 0 and step > 0:
if step > 0 and step % max_steps == 0:
# evaluation
radiance_field.eval()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment