Cleanup

nerfstudio/cameras/cameras.py

@@ -275,8 +275,8 @@ class Cameras:
             coords = self.get_image_coords().to(self.device)
 
         assert coords is not None
-        y = coords[..., 0]  # (..., 1)
-        x = coords[..., 1]  # (..., 1)
+        y = coords[..., 0]
+        x = coords[..., 1]
         fx, fy = self.fx[camera_indices], self.fy[camera_indices]
         cx, cy = self.cx[camera_indices], self.cy[camera_indices]
 

nerfstudio/data/datamanagers.py

@@ -38,7 +38,7 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
 )
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
 from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig
-from nerfstudio.data.pixel_samplers import EquirectangularSampler, PixelSampler
+from nerfstudio.data.pixel_samplers import PixelSampler
 from nerfstudio.data.utils.dataloaders import (
     CacheDataloader,
     FixedIndicesEvalDataloader,

nerfstudio/data/pixel_samplers.py

@@ -110,93 +110,96 @@ class PixelSampler:  # pylint: disable=too-few-public-methods
         )
         return pixel_batch
 
-def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False):
-    """
-    Operates on a batch of images and samples pixels to use for generating rays.
-    Returns a collated batch which is input to the Graph.
-    It will sample only within the valid 'mask' if it's specified.
 
-    Args:
-        batch: batch of images to sample from
-        num_rays_per_batch: number of rays to sample per batch
-        keep_full_image: whether or not to include a reference to the full image in returned batch
-    """
-    device = batch["image"].device
-    num_images, image_height, image_width, _ = batch["image"].shape
+# TODO(kevinddchen): Implement equirect pixel sampler
 
-    # only sample within the mask, if the mask is in the batch
-    # if "mask" in batch:
-    #     nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False)
-    #     chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch)
-    #     indices = nonzero_indices[chosen_indices]
-    # else:
-    #     indices = torch.floor(
-    #         torch.rand((num_rays_per_batch, 3), device=device)
-    #         * torch.tensor([num_images, image_height, image_width], device=device)
-    #     ).long()
+# def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False):
+#     """
+#     Operates on a batch of images and samples pixels to use for generating rays.
+#     Returns a collated batch which is input to the Graph.
+#     It will sample only within the valid 'mask' if it's specified.
 
-    number_of_images = torch.rand(num_rays_per_batch, device=device)
-    phi = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / torch.pi
-    theta = torch.rand(num_rays_per_batch, device=device)
-    indices = torch.floor(
-            torch.stack((number_of_images, phi, theta), dim=-1)
-            * torch.tensor([num_images, image_height, image_width], device=device)
-        ).long()
+#     Args:
+#         batch: batch of images to sample from
+#         num_rays_per_batch: number of rays to sample per batch
+#         keep_full_image: whether or not to include a reference to the full image in returned batch
+#     """
+#     device = batch["image"].device
+#     num_images, image_height, image_width, _ = batch["image"].shape
 
-    c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1))
-    image = batch["image"][c, y, x]
-    # mask, semantics = None, None
-    # if "mask" in batch:
-    #     mask = batch["mask"][c, y, x]
-    # if "semantics" in batch:
-    #     semantics = batch["semantics"][c, y, x]
-    # assert image.shape == (num_rays_per_batch, 3), image.shape
+#     # only sample within the mask, if the mask is in the batch
+#     # if "mask" in batch:
+#     #     nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False)
+#     #     chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch)
+#     #     indices = nonzero_indices[chosen_indices]
+#     # else:
+#     #     indices = torch.floor(
+#     #         torch.rand((num_rays_per_batch, 3), device=device)
+#     #         * torch.tensor([num_images, image_height, image_width], device=device)
+#     #     ).long()
 
-    # Needed to correct the random indices to their actual camera idx locations.
-    local_indices = indices.clone()
-    indices[:, 0] = batch["image_idx"][c]
-    collated_batch = {
-        "local_indices": local_indices,  # local to the batch returned
-        "indices": indices,  # with the abs camera indices
-        "image": image,
-    }
-    # if mask is not None:
-    #     collated_batch["mask"] = mask
-    # if semantics is not None:
-    #     collated_batch["semantics"] = semantics
+#     number_of_images = torch.rand(num_rays_per_batch, device=device)
+#     phi = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / torch.pi
+#     theta = torch.rand(num_rays_per_batch, device=device)
+#     indices = torch.floor(
+#             torch.stack((number_of_images, phi, theta), dim=-1)
+#             * torch.tensor([num_images, image_height, image_width], device=device)
+#         ).long()
 
-    # if keep_full_image:
-    #     collated_batch["full_image"] = batch["image"]
+#     c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1))
+#     image = batch["image"][c, y, x]
+#     # mask, semantics = None, None
+#     # if "mask" in batch:
+#     #     mask = batch["mask"][c, y, x]
+#     # if "semantics" in batch:
+#     #     semantics = batch["semantics"][c, y, x]
+#     # assert image.shape == (num_rays_per_batch, 3), image.shape
 
-    return collated_batch
+#     # Needed to correct the random indices to their actual camera idx locations.
+#     local_indices = indices.clone()
+#     indices[:, 0] = batch["image_idx"][c]
+#     collated_batch = {
+#         "local_indices": local_indices,  # local to the batch returned
+#         "indices": indices,  # with the abs camera indices
+#         "image": image,
+#     }
+#     # if mask is not None:
+#     #     collated_batch["mask"] = mask
+#     # if semantics is not None:
+#     #     collated_batch["semantics"] = semantics
 
-class EquirectangularSampler:  # pylint: disable=too-few-public-methods
-    """Samples 'pixel_batch's from 'image_batch's.
+#     # if keep_full_image:
+#     #     collated_batch["full_image"] = batch["image"]
 
-    Args:
-        num_rays_per_batch: number of rays to sample per batch
-        keep_full_image: whether or not to include a reference to the full image in returned batch
-    """
+#     return collated_batch
 
-    def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False) -> None:
-        self.num_rays_per_batch = num_rays_per_batch
-        self.keep_full_image = keep_full_image
+# class EquirectangularSampler:  # pylint: disable=too-few-public-methods
+#     """Samples 'pixel_batch's from 'image_batch's.
 
-    def set_num_rays_per_batch(self, num_rays_per_batch: int):
-        """Set the number of rays to sample per batch.
+#     Args:
+#         num_rays_per_batch: number of rays to sample per batch
+#         keep_full_image: whether or not to include a reference to the full image in returned batch
+#     """
 
-        Args:
-            num_rays_per_batch: number of rays to sample per batch
-        """
-        self.num_rays_per_batch = num_rays_per_batch
+#     def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False) -> None:
+#         self.num_rays_per_batch = num_rays_per_batch
+#         self.keep_full_image = keep_full_image
 
-    def sample(self, image_batch: Dict):
-        """Sample an image batch and return a pixel batch.
+#     def set_num_rays_per_batch(self, num_rays_per_batch: int):
+#         """Set the number of rays to sample per batch.
 
-        Args:
-            image_batch: batch of images to sample from
-        """
-        pixel_batch = collate_image_dataset_batch_equirectangular(
-            image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image
-        )
-        return pixel_batch
+#         Args:
+#             num_rays_per_batch: number of rays to sample per batch
+#         """
+#         self.num_rays_per_batch = num_rays_per_batch
+
+#     def sample(self, image_batch: Dict):
+#         """Sample an image batch and return a pixel batch.
+
+#         Args:
+#             image_batch: batch of images to sample from
+#         """
+#         pixel_batch = collate_image_dataset_batch_equirectangular(
+#             image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image
+#         )
+#         return pixel_batch

scripts/process_data.py

@@ -260,7 +260,7 @@ def copy_images(data: Path, image_dir: Path, verbose: bool, rename: bool = True)
         data: Path to the directory of images.
         image_dir: Path to the output directory.
         verbose: If True, print extra logging.
-        rename: If True, rename the images to be sequential.
+        rename: If True, rename the images sequentially to `frame_{num}`.
     Returns:
         The number of images copied.
     """
@@ -578,7 +578,7 @@ def opensfm_to_json(reconstruction_path: Path, output_dir: Path, camera_model: C
         aa_vec = np.array(shot_info["rotation"])  # 3D axis angle repr
         angle = np.linalg.norm(aa_vec)
         if angle > 1e-8:
-            # normalilze the axis-angle repr if angle is large enough
+            # normalize the axis-angle repr if angle is large enough
             aa_vec /= angle
             qx = aa_vec[0] * np.sin(angle / 2)
             qy = aa_vec[1] * np.sin(angle / 2)
@@ -594,6 +594,9 @@ def opensfm_to_json(reconstruction_path: Path, output_dir: Path, camera_model: C
         # Convert camera +z forwards (OpenSfM) convention to -z forwards (NeRF) convention
         # Equivalent to right-multiplication by diag([1, -1, -1, 1])
         c2w[0:3, 1:3] *= -1
+        # Rotation around global z-axis by -90 degrees
+        c2w = c2w[np.array([1, 0, 2, 3]), :]
+        c2w[2, :] *= -1
 
         name = Path(f"./images/{shot_name}")