From 391fd9324798bad59737cda2a2d2096fd29de7a9 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 22 Apr 2022 16:56:50 -0700 Subject: [PATCH] Fix camera index bug --- mattport/nerf/dataset/image_dataset.py | 4 +++- mattport/nerf/field_modules/ray_generator.py | 6 +++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index da8fd233f..0acdc2405 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/mattport/nerf/dataset/image_dataset.py @@ -80,7 +80,7 @@ def collate_batch(batch, num_rays_per_batch, keep_full_image: bool = False): """_summary_ Args: - batch (_type_): _description_ + batch (_type_): _description_ """ # TODO(ethan): handle sampling even when in different processes # we don't want the same samples for all spawned processed when @@ -96,6 +96,8 @@ def collate_batch(batch, num_rays_per_batch, keep_full_image: bool = False): mask = batch["mask"][c, y, x] assert pixels.shape == (num_rays_per_batch, 3), pixels.shape + indices[:, 0] = batch["image_idx"][c] + collated_batch = { "indices": indices, "pixels": pixels, diff --git a/mattport/nerf/field_modules/ray_generator.py b/mattport/nerf/field_modules/ray_generator.py index 8bfedfbf7..8abb01a1c 100644 --- a/mattport/nerf/field_modules/ray_generator.py +++ b/mattport/nerf/field_modules/ray_generator.py @@ -37,8 +37,8 @@ class RayGenerator(nn.Module): x (_type_): _description_ """ c = ray_indices[:, 0] # camera indices - i = ray_indices[:, 1] # row indices - j = ray_indices[:, 2] # col indices + y = ray_indices[:, 1] # row indices + x = ray_indices[:, 2] # col indices intrinsics = self.intrinsics[c] camera_to_world = self.camera_to_world[c] # NOTE(ethan): we currently assume all images have the same height and width @@ -46,7 +46,7 @@ class RayGenerator(nn.Module): camera_class = get_camera_model(self.num_intrinsics_params) camera = camera_class(*self.intrinsics[camera_index].tolist()) image_coords = camera.get_image_coords().to(ray_indices.device) - coords = image_coords[i, j] + coords = image_coords[y, x] ray_bundle = camera_class.generate_rays(intrinsics=intrinsics, camera_to_world=camera_to_world, coords=coords) ray_bundle.camera_indices = c