mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 02:45:34 +08:00
Ethan Weber
e38040ee4d
* working data loader code * adding gitignore * fixing linting errors * set black line length in settings file * update * fixed filepath * black format * tqdm
76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
"""
|
|
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()
|