Fixed friends dataset w/o caching (#365)

* fixing friends dataloading and nerfw running

* put dataset back to normal

nerfactory/cameras/cameras.py

@@ -220,7 +220,17 @@ class Cameras:
         dy = torch.sqrt(torch.sum((directions - directions_stack[2]) ** 2, dim=-1))
         pixel_area = dx * dy
 
-        return RayBundle(origins=origins, directions=directions, pixel_area=pixel_area[..., None])
+        if not isinstance(camera_indices, torch.Tensor):
+            ray_bundle_camera_indices = torch.Tensor([camera_indices]).broadcast_to((self._num_cameras)).to(self.device)
+        else:
+            ray_bundle_camera_indices = camera_indices
+
+        return RayBundle(
+            origins=origins,
+            directions=directions,
+            pixel_area=pixel_area[..., None],
+            camera_indices=ray_bundle_camera_indices,
+        )
 
     def to_json(
         self,

nerfactory/configs/base.py

@@ -199,6 +199,7 @@ class FriendsDataParserConfig(DataParserConfig):
 
     _target: Type = Friends
     data_directory: Path = Path("data/friends/TBBT-big_living_room")
+    include_semantics: bool = True
 
 
 @dataclass
@@ -302,6 +303,9 @@ class NerfWModelConfig(ModelConfig):
     num_coarse_samples: int = 64
     num_importance_samples: int = 64
     uncertainty_min: float = 0.03
+    num_images: int = 10000  # TODO: don't hardcode this
+    appearance_embedding_dim: int = 48
+    transient_embedding_dim: int = 16
 
 
 @dataclass

nerfactory/datamanagers/base.py

@@ -100,9 +100,9 @@ class DataManager(nn.Module):
         super().__init__()
         self.train_count = 0
         self.eval_count = 0
-        if self.train_input_dataset and self.train_input_dataset.inputs:
+        if self.train_input_dataset and self.train_input_dataset.dataset_inputs:
             self.setup_train()
-        if self.eval_input_dataset and self.eval_input_dataset.inputs:
+        if self.eval_input_dataset and self.eval_input_dataset.dataset_inputs:
             self.setup_eval()
 
     def forward(self):
@@ -250,7 +250,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             )
         self.iter_train_image_dataloader = iter(self.train_image_dataloader)
         self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch)
-        self.train_ray_generator = RayGenerator(self.train_input_dataset.inputs.cameras.to(self.device))
+        self.train_ray_generator = RayGenerator(self.train_input_dataset.dataset_inputs.cameras.to(self.device))
 
     def setup_eval(self):
         """Sets up the data loader for evaluation"""

nerfactory/datamanagers/dataloaders.py

@@ -117,7 +117,7 @@ class EvalDataloader(DataLoader):
         **kwargs,
     ):
         self.input_dataset = input_dataset
-        self.cameras = input_dataset.inputs.cameras.to(device)
+        self.cameras = input_dataset.dataset_inputs.cameras.to(device)
         self.num_rays_per_chunk = num_rays_per_chunk
         self.device = device
         self.kwargs = kwargs

nerfactory/datamanagers/dataparsers/base.py

@@ -16,16 +16,11 @@
 
 from __future__ import annotations
 
-import dataclasses
-import logging
 from abc import abstractmethod
 from dataclasses import dataclass
-from pathlib import Path
 
-from nerfactory.configs import base as cfg
+import nerfactory.configs.base as cfg
 from nerfactory.datamanagers.structs import DatasetInputs
-from nerfactory.utils.io import get_absolute_path, load_from_pkl, write_to_pkl
-from nerfactory.utils.misc import get_hash_str_from_dict
 
 
 @dataclass
@@ -47,60 +42,14 @@ class DataParser:
             DatasetInputs
         """
 
-    def get_dataset_inputs(self, split: str = "train", use_preprocessing_cache: bool = False) -> DatasetInputs:
+    def get_dataset_inputs(self, split: str = "train") -> DatasetInputs:
         """Returns the dataset inputs for the given split.
 
         Args:
             split: Which dataset split to generate.
-            use_preprocessing_cache: Whether to use the cached dataset inputs. Defaults to False.
 
         Returns:
             DatasetInputs
         """
-        if use_preprocessing_cache:
-            dataset_inputs = self._get_dataset_inputs_from_cache(split)
-        else:
-            dataset_inputs = self._generate_dataset_inputs(split)
-        return dataset_inputs
-
-    def _get_cache_filename(self, split: str) -> Path:
-        """Creates a cache filename from the dataset inputs arguments.
-
-        Args:
-            split: Which dataset split to generate filename for.
-
-        Returns:
-            filename for cache.
-        """
-        dataset_config_hash = get_hash_str_from_dict(dataclasses.asdict(self))
-        dataset_config_hash_filename = get_absolute_path(
-            f"cache/dataset_inputs/{dataset_config_hash}-{split}.pkl"
-        ).parent.mkdir()
-        assert dataset_config_hash_filename is not None, "dataset hash is None"
-        return dataset_config_hash_filename
-
-    def save_dataset_inputs_to_cache(self, split: str):
-        """Saves the dataset inputs to cache.
-
-        Args:
-            split: Which dataset split to save.
-        """
-        dataset_inputs = self.get_dataset_inputs(split=split)
-        dataset_config_hash_filename = self._get_cache_filename(split)
-        write_to_pkl(dataset_config_hash_filename, dataset_inputs)
-
-    def _get_dataset_inputs_from_cache(self, split: str) -> DatasetInputs:
-        """Loads the dataset inputs from cache. If the cache does not exist, it will be created.
-
-        Args:
-            split: Which dataset split to load.
-        """
-        dataset_config_hash_filename = self._get_cache_filename(split)
-        if dataset_config_hash_filename.exists():
-            logging.info("Loading dataset from cache.")
-            dataset_inputs = load_from_pkl(dataset_config_hash_filename)
-        else:
-            logging.info("Cache file not found. Generating and saving dataset to cache.")
-            dataset_inputs = self._generate_dataset_inputs(split=split)
-            self.save_dataset_inputs_to_cache(split)
+        dataset_inputs = self._generate_dataset_inputs(split)
         return dataset_inputs

nerfactory/datamanagers/dataparsers/friends_parser.py

@@ -21,34 +21,27 @@ from pathlib import Path
 import numpy as np
 import torch
 from PIL import Image
+from rich.console import Console
 
+import nerfactory.configs.base as cfg
 from nerfactory.cameras.cameras import Cameras, CameraType
 from nerfactory.datamanagers.dataparsers.base import DataParser
-from nerfactory.datamanagers.datasets import InputDataset
-from nerfactory.datamanagers.structs import (
-    DatasetInputs,
-    PointCloud,
-    SceneBounds,
-    Semantics,
-)
-from nerfactory.utils.colmap_utils import (
-    read_cameras_binary,
-    read_images_binary,
-    read_points3d_binary,
-)
+from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds, Semantics
 from nerfactory.utils.io import get_absolute_path, load_from_json
 
+CONSOLE = Console()
 
-def get_semantics_and_masks(self: InputDataset, image_idx: int):
+
+def get_semantics_and_masks(semantics: Semantics, image_idx: int):
     """function to process additional semantics and mask information"""
     # handle mask
-    person_index = self.inputs.additional_inputs["semantics"].thing_classes.index("person")
-    thing_image_filename = self.inputs.additional_inputs["semantics"].thing_filenames[image_idx]
+    person_index = semantics.thing_classes.index("person")
+    thing_image_filename = semantics.thing_filenames[image_idx]
     pil_image = Image.open(thing_image_filename)
     thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None]
     mask = (thing_semantics != person_index).to(torch.float32)  # 1 where valid
     # handle semantics
-    stuff_image_filename = self.inputs.additional_inputs["semantics"].stuff_filenames[image_idx]
+    stuff_image_filename = semantics.stuff_filenames[image_idx]
     pil_image = Image.open(stuff_image_filename)
     stuff_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None]
     return {"mask": mask, "semantics": stuff_semantics}
@@ -56,17 +49,9 @@ def get_semantics_and_masks(self: InputDataset, image_idx: int):
 
 @dataclass
 class Friends(DataParser):
-    """Friends Dataset
+    """Friends Dataset"""
 
-    Args:
-        data_directory: Location of data
-        include_semantics: whether or not to include the semantics. Defaults to False.
-        include_point_cloud: whether or not to include the point cloud. Defaults to False.
-    """
-
-    data_directory: Path
-    include_semantics: bool = True
-    include_point_cloud: bool = False
+    config: cfg.FriendsDataParserConfig
 
     @classmethod
     def _get_aabb_and_transform(cls, basedir):
@@ -90,62 +75,53 @@ class Friends(DataParser):
         aabb = bbox  # rename to aabb because it's an axis-aligned bounding box
         return torch.from_numpy(aabb).float(), torch.from_numpy(transposed_point_cloud_transform).float()
 
-    def _generate_dataset_inputs(self, split="train"):  # pylint: disable=too-many-statements
+    def _generate_dataset_inputs(self, split="train"):  # pylint: disable=unused-argument
 
-        abs_dir = get_absolute_path(self.data_directory)
+        abs_dir = get_absolute_path(self.config.data_directory)
 
-        images_data = read_images_binary(abs_dir / "colmap" / "images.bin")
-        # `image_path` is only the end of the filename, including the extension e.g., `.jpg`
-        image_paths = sorted((abs_dir / "images").iterdir())
+        cameras_json = load_from_json(abs_dir / "cameras.json")
+        frames = cameras_json["frames"]
+        bbox = torch.tensor(cameras_json["bbox"])
 
-        image_path_to_image_id = {}
-        image_id_to_image_path = {}
-        for v in images_data.values():
-            image_path_to_image_id[v.name] = v.id
-            image_id_to_image_path[v.id] = v.name
-        # TODO: handle the splits differently
-        image_filenames = [abs_dir / "images" / image_path for image_path in image_paths]
+        image_filenames = []
+        fx = []
+        fy = []
+        cx = []
+        cy = []
+        camera_to_worlds = []
+        for frame in frames:
+            # unpack data
+            image_filename = abs_dir / "images" / frame["image_name"]
+            intrinsics = torch.tensor(frame["intrinsics"])
+            camtoworld = torch.tensor(frame["camtoworld"])[:3]
+            # append data
+            image_filenames.append(image_filename)
+            fx.append(intrinsics[0, 0])
+            fy.append(intrinsics[1, 1])
+            cx.append(intrinsics[0, 2])
+            cy.append(intrinsics[1, 2])
+            camera_to_worlds.append(camtoworld)
+        fx = torch.stack(fx)
+        fy = torch.stack(fy)
+        cx = torch.stack(cx)
+        cy = torch.stack(cy)
+        camera_to_worlds = torch.stack(camera_to_worlds)
+
+        # rotate the cameras and box 90 degrees about the x axis to put the z axis up
+        rotation = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]], dtype=torch.float32)
+        camera_to_worlds[:, :3] = rotation @ camera_to_worlds[:, :3]
+        bbox = (rotation @ bbox.T).T
 
         # -- set the bounding box ---
-        aabb, transposed_point_cloud_transform = self._get_aabb_and_transform(abs_dir)
-        scene_bounds_original = SceneBounds(aabb=aabb)
-        # for shifting and rescale accoding to scene bounds
-        box_center = scene_bounds_original.get_center()
-        box_scale_factor = 5.0 / scene_bounds_original.get_diagonal_length()  # the target diagonal length
-        scene_bounds = scene_bounds_original.get_centered_and_scaled_scene_bounds(box_scale_factor)
-
-        # --- intrinsics ---
-        cameras_data = read_cameras_binary(abs_dir / "colmap" / "cameras.bin")
-        focal_lengths = []
-        for image_path in image_paths:
-            cam = cameras_data[image_path_to_image_id[image_path]]
-            assert len(cam.params) == 3
-            focal_lengths.append(cam.params[0])  # f (fx and fy)
-        focal_lengths = torch.stack(focal_lengths, dim=0)
-
-        cam = cameras_data[image_path_to_image_id[image_paths[0]]]
-        cx = cam.params[1]  # cx
-        cy = cam.params[2]  # cy
-
-        # --- camera_to_world (extrinsics) ---
-        camera_to_world = []
-        bottom_row = np.array([0, 0, 0, 1.0]).reshape(1, 4)
-        for image_path in image_paths:
-            image_data = images_data[image_path_to_image_id[image_path]]
-            rot = image_data.qvec2rotmat()
-            trans = image_data.tvec.reshape(3, 1)
-            c2w = np.concatenate([np.concatenate([rot, trans], 1), bottom_row], 0)
-            camera_to_world.append(c2w)
-        camera_to_world = torch.tensor(np.array(camera_to_world)).float()
-        camera_to_world = torch.inverse(camera_to_world)
-        camera_to_world[..., 1:3] *= -1
-        camera_to_world = transposed_point_cloud_transform @ camera_to_world
-        camera_to_world = camera_to_world[:, :3]
-        camera_to_world[..., 3] = (camera_to_world[..., 3] - box_center) * box_scale_factor  # center and rescale
+        scene_bounds = SceneBounds(aabb=bbox)
+        # # for shifting and rescale accoding to scene bounds
+        # box_center = scene_bounds_original.get_center()
+        # box_scale_factor = 5.0 / scene_bounds_original.get_diagonal_length()  # the target diagonal length
+        # scene_bounds = scene_bounds_original.get_centered_and_scaled_scene_bounds(box_scale_factor)
 
         # --- semantics ---
         semantics = None
-        if self.include_semantics:
+        if self.config.include_semantics:
             thing_filenames = [
                 Path(str(image_filename).replace("/images/", "/segmentations/thing/").replace(".jpg", ".png"))
                 for image_filename in image_filenames
@@ -168,32 +144,21 @@ class Friends(DataParser):
                 thing_filenames=thing_filenames,
             )
 
-        # Possibly include the sparse point cloud from COLMAP in the dataset inputs.
-        # NOTE(ethan): this will be common across the different splits.
-        point_cloud = PointCloud()
-        if self.include_point_cloud:
-            points_3d = read_points3d_binary(abs_dir / "colmap" / "points3D.bin")
-            xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])).float()
-            rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])).float()
-            xyz_h = torch.cat([xyz, torch.ones_like(xyz[..., :1])], -1)
-            xyz = (xyz_h @ transposed_point_cloud_transform.T)[..., :3]
-            xyz = (xyz - box_center) * box_scale_factor  # center and rescale
-            point_cloud.xyz = xyz
-            point_cloud.rgb = rgb
+        assert torch.all(cx[0] == cx), "Not all cameras have the same cx. Our Cameras class does not support this."
+        assert torch.all(cy[0] == cy), "Not all cameras have the same cy. Our Cameras class does not support this."
 
         cameras = Cameras(
-            fx=focal_lengths,
-            fy=focal_lengths,
-            cx=cx,
-            cy=cy,
-            camera_to_worlds=camera_to_world,
+            fx=fx,
+            fy=fy,
+            cx=float(cx[0]),
+            cy=float(cy[0]),
+            camera_to_worlds=camera_to_worlds,
             camera_type=CameraType.PERSPECTIVE,
         )
 
         dataset_inputs = DatasetInputs(
             image_filenames=image_filenames,
             cameras=cameras,
-            point_cloud=point_cloud,
             scene_bounds=scene_bounds,
             additional_inputs={"semantics": {"data": semantics, "func": get_semantics_and_masks}},
         )

nerfactory/datamanagers/datasets.py

@@ -42,10 +42,10 @@ class InputDataset(Dataset):
             alpha_color (TensorType[3], optional): Sets transparent regions to specified color, otherwise black.
         """
         super().__init__()
-        self.inputs: DatasetInputs = config.setup().get_dataset_inputs(split=split)
+        self.dataset_inputs: DatasetInputs = config.setup().get_dataset_inputs(split=split)
 
     def __len__(self):
-        return len(self.inputs.image_filenames)
+        return len(self.dataset_inputs.image_filenames)
 
     def get_numpy_image(self, image_idx: int) -> npt.NDArray[np.uint8]:
         """Returns the image.
@@ -56,7 +56,7 @@ class InputDataset(Dataset):
         Returns:
             np.uint8: an image of shape (H, W, 3 or 4)
         """
-        image_filename = self.inputs.image_filenames[image_idx]
+        image_filename = self.dataset_inputs.image_filenames[image_idx]
         pil_image = Image.open(image_filename)
         image = np.array(pil_image, dtype="uint8")  # shape is (h, w, 3 or 4)
         assert len(image.shape) == 3
@@ -67,9 +67,9 @@ class InputDataset(Dataset):
     def get_image(self, image_idx: int) -> TensorType["image_height", "image_width", "num_channels"]:
         """Returns a 3 channel image."""
         image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32") / 255.0)
-        if self.inputs.alpha_color is not None and image.shape[-1] == 4:
+        if self.dataset_inputs.alpha_color is not None and image.shape[-1] == 4:
             assert image.shape[-1] == 4
-            image = image[:, :, :3] * image[:, :, -1:] + self.inputs.alpha_color * (1.0 - image[:, :, -1:])
+            image = image[:, :, :3] * image[:, :, -1:] + self.dataset_inputs.alpha_color * (1.0 - image[:, :, -1:])
         else:
             image = image[:, :, :3]
         return image
@@ -80,11 +80,11 @@ class InputDataset(Dataset):
         data = {"image_idx": image_idx}
         assert is_not_none(image)
         data["image"] = image
-        for _, data_func_dict in self.inputs.additional_inputs.items():
+        for _, data_func_dict in self.dataset_inputs.additional_inputs.items():
             assert "func" in data_func_dict, "Missing function to process data: specify `func` in `additional_inputs`"
             func = data_func_dict["func"]
             assert "data" in data_func_dict, "No data to process: specify `data` in `additional_inputs`"
-            data.update(func(data_func_dict["data"]))
+            data.update(func(data_func_dict["data"], image_idx))
         return data
 
     def __getitem__(self, image_idx):

nerfactory/fields/nerfw_field.py

@@ -130,11 +130,12 @@ class VanillaNerfWField(Field):
         encoded_dir = self.direction_encoding(ray_samples.frustums.directions)
         if ray_samples.camera_indices is None:
             raise AttributeError("Camera indices are not provided.")
-        embedded_appearance = self.embedding_appearance(ray_samples.camera_indices.squeeze())
+        camera_indices = ray_samples.camera_indices.squeeze().to(ray_samples.frustums.origins.device)
+        embedded_appearance = self.embedding_appearance(camera_indices)
         mlp_in = torch.cat([density_embedding, encoded_dir, embedded_appearance], dim=-1)  # type: ignore
         mlp_head_out = self.mlp_head(mlp_in)
         outputs[self.field_head_rgb.field_head_name] = self.field_head_rgb(mlp_head_out)  # static rgb
-        embedded_transient = self.embedding_transient(ray_samples.camera_indices.squeeze())
+        embedded_transient = self.embedding_transient(camera_indices)
         transient_mlp_in = torch.cat([density_embedding, embedded_transient], dim=-1)  # type: ignore
         transient_mlp_out = self.mlp_transient(transient_mlp_in)
         outputs[self.field_head_transient_uncertainty.field_head_name] = self.field_head_transient_uncertainty(

nerfactory/models/nerfw.py

@@ -63,9 +63,6 @@ class NerfWModel(Model):
         """
         self.field_coarse = None
         self.field_fine = None
-        self.num_images = 10000  # TODO(ethan): fix this
-        self.appearance_embedding_dim = 48
-        self.transient_embedding_dim = 16
         super().__init__(config=config, **kwargs)
 
     def populate_modules(self):
@@ -82,11 +79,11 @@ class NerfWModel(Model):
 
         self.field_coarse = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding)
         self.field_fine = VanillaNerfWField(
-            num_images=self.num_images,
+            num_images=self.config.num_images,
             position_encoding=position_encoding,
             direction_encoding=direction_encoding,
-            appearance_embedding_dim=self.appearance_embedding_dim,
-            transient_embedding_dim=self.transient_embedding_dim,
+            appearance_embedding_dim=self.config.appearance_embedding_dim,
+            transient_embedding_dim=self.config.transient_embedding_dim,
         )
 
         # samplers
@@ -123,7 +120,7 @@ class NerfWModel(Model):
         if ray_bundle.camera_indices is not None:
             # TODO(ethan): remove this check
             assert (
-                torch.max(ray_bundle.camera_indices) < self.num_images
+                torch.max(ray_bundle.camera_indices) < self.config.num_images
             ), "num_images must be greater than the largest camera index"
 
         if self.field_coarse is None or self.field_fine is None:

nerfactory/pipelines/base.py

@@ -84,7 +84,9 @@ class Pipeline(nn.Module):
         self.datamanager.to(device)
         # TODO(ethan): get rid of scene_bounds from the model
         assert self.datamanager.train_input_dataset is not None, "Missing input dataset"
-        self.model: Model = config.model.setup(scene_bounds=self.datamanager.train_input_dataset.inputs.scene_bounds)
+        self.model: Model = config.model.setup(
+            scene_bounds=self.datamanager.train_input_dataset.dataset_inputs.scene_bounds
+        )
         self.model.to(device)
 
         self.world_size = world_size

nerfactory/viewer/server/viewer_utils.py

@@ -244,11 +244,11 @@ class VisualizerState:
         for idx in image_indices:
             image = dataset[idx]["image"]
             bgr = image[..., [2, 1, 0]]
-            camera_json = dataset.inputs.cameras.to_json(camera_idx=idx, image=bgr, resize_shape=(100, 100))
+            camera_json = dataset.dataset_inputs.cameras.to_json(camera_idx=idx, image=bgr, resize_shape=(100, 100))
             self.vis[f"sceneState/cameras/{idx:06d}"].write(camera_json)
 
         # draw the scene bounds (i.e., the bounding box)
-        json_ = dataset.inputs.scene_bounds.to_json()
+        json_ = dataset.dataset_inputs.scene_bounds.to_json()
         self.vis["sceneState/sceneBounds"].write(json_)
 
         # set the properties of the camera