diff --git a/README.md b/README.md index 40ebbac4..95db5213 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/examples/datasets/dnerf_synthetic.py b/examples/datasets/dnerf_synthetic.py index 84384228..564c2e62 100644 --- a/examples/datasets/dnerf_synthetic.py +++ b/examples/datasets/dnerf_synthetic.py @@ -86,6 +86,7 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( [0, 0, 1], ], dtype=torch.float32, + device=device, ) # (3, 3) assert self.images.shape[1:3] == (self.HEIGHT, self.WIDTH) diff --git a/examples/datasets/nerf_360_v2.py b/examples/datasets/nerf_360_v2.py index d04c99b9..895861db 100644 --- a/examples/datasets/nerf_360_v2.py +++ b/examples/datasets/nerf_360_v2.py @@ -169,6 +169,7 @@ def __init__( 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 @@ def __init__( 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): diff --git a/examples/datasets/nerf_synthetic.py b/examples/datasets/nerf_synthetic.py index 11542924..84bbb342 100644 --- a/examples/datasets/nerf_synthetic.py +++ b/examples/datasets/nerf_synthetic.py @@ -79,6 +79,7 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( [0, 0, 1], ], dtype=torch.float32, + device=device, ) # (3, 3) assert self.images.shape[1:3] == (self.HEIGHT, self.WIDTH) diff --git a/examples/train_mlp_dnerf.py b/examples/train_mlp_dnerf.py index 9702ee35..a0289b63 100644 --- a/examples/train_mlp_dnerf.py +++ b/examples/train_mlp_dnerf.py @@ -4,7 +4,7 @@ import argparse import math -import os +import pathlib import time import imageio @@ -24,6 +24,12 @@ 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 @@ 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 @@ 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() diff --git a/examples/train_mlp_nerf.py b/examples/train_mlp_nerf.py index aa114a44..c651369e 100644 --- a/examples/train_mlp_nerf.py +++ b/examples/train_mlp_nerf.py @@ -4,7 +4,7 @@ import argparse import math -import os +import pathlib import time import imageio @@ -23,6 +23,12 @@ 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 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 @@ 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 @@ 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() diff --git a/examples/train_ngp_nerf.py b/examples/train_ngp_nerf.py index 584cc13a..afa952ad 100644 --- a/examples/train_ngp_nerf.py +++ b/examples/train_ngp_nerf.py @@ -4,7 +4,7 @@ import argparse import math -import os +import pathlib import time import imageio @@ -23,6 +23,12 @@ 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 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 @@ 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 @@ def occ_eval_fn(x): 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()