mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 02:45:34 +08:00
Ethan Weber
fb440fedf3
* occupancy grid in progress * visualizing occupancy grid. currently restructing field and renderer head outputs * black and pylint. broke nerfw and semantic nerf but it's okay * get_project_root
89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
"""
|
|
Code to test ray generation.
|
|
"""
|
|
|
|
import imageio
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
from tqdm import tqdm
|
|
|
|
from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one
|
|
from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
|
|
from mattport.nerf.dataset.utils import get_dataset_inputs
|
|
from mattport.nerf.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["local_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(visualize=False):
|
|
"""Testing for the dataloader from input dataset parameters to rays."""
|
|
|
|
data_directory = "tests/data/lego_test"
|
|
dataset_format = "blender"
|
|
downscale_factor = 1
|
|
num_images_to_sample_from = 1
|
|
num_times_to_repeat_images = 40
|
|
num_rays_per_batch = 1024
|
|
num_workers = 0
|
|
|
|
data_directory = get_absolute_path(data_directory)
|
|
dataset_inputs = get_dataset_inputs(
|
|
data_directory=data_directory, dataset_format=dataset_format, downscale_factor=downscale_factor, split="train"
|
|
)
|
|
image_dataset = ImageDataset(
|
|
image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor
|
|
)
|
|
iter_dataset = CollateIterDataset(
|
|
image_dataset,
|
|
collate_fn=lambda batch_list: collate_batch(batch_list, num_rays_per_batch, keep_full_image=True),
|
|
num_samples_to_collate=num_images_to_sample_from,
|
|
num_times_to_repeat=num_times_to_repeat_images,
|
|
)
|
|
dataloader = DataLoader(
|
|
iter_dataset,
|
|
batch_size=1,
|
|
num_workers=num_workers,
|
|
collate_fn=collate_batch_size_one,
|
|
pin_memory=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)
|
|
if visualize:
|
|
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)
|
|
if visualize:
|
|
fig.write_image("temp1.png")
|
|
fig.write_html("temp1.html")
|
|
assert True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_dataloader(visualize=True)
|