From 512a51f7147446a0da9e957ea4b7a43474bf98e2 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Wed, 7 Sep 2022 12:24:29 -0700 Subject: [PATCH] Fixed friends dataset w/o caching (#365) * fixing friends dataloading and nerfw running * put dataset back to normal --- nerfactory/cameras/cameras.py | 12 +- nerfactory/configs/base.py | 4 + nerfactory/datamanagers/base.py | 6 +- nerfactory/datamanagers/dataloaders.py | 2 +- nerfactory/datamanagers/dataparsers/base.py | 57 +------ .../dataparsers/friends_parser.py | 149 +++++++----------- nerfactory/datamanagers/datasets.py | 14 +- nerfactory/fields/nerfw_field.py | 5 +- nerfactory/models/nerfw.py | 11 +- nerfactory/pipelines/base.py | 4 +- nerfactory/viewer/server/viewer_utils.py | 4 +- 11 files changed, 98 insertions(+), 170 deletions(-) diff --git a/nerfactory/cameras/cameras.py b/nerfactory/cameras/cameras.py index eef5f2397..7e5415eb9 100644 --- a/nerfactory/cameras/cameras.py +++ b/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, diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 0616036d7..e82f57921 100644 --- a/nerfactory/configs/base.py +++ b/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 diff --git a/nerfactory/datamanagers/base.py b/nerfactory/datamanagers/base.py index 2d8bb0c4f..e8e842e47 100644 --- a/nerfactory/datamanagers/base.py +++ b/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""" diff --git a/nerfactory/datamanagers/dataloaders.py b/nerfactory/datamanagers/dataloaders.py index 0614fb800..9a7716f8c 100644 --- a/nerfactory/datamanagers/dataloaders.py +++ b/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 diff --git a/nerfactory/datamanagers/dataparsers/base.py b/nerfactory/datamanagers/dataparsers/base.py index 4f18b7679..a3101c896 100644 --- a/nerfactory/datamanagers/dataparsers/base.py +++ b/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 diff --git a/nerfactory/datamanagers/dataparsers/friends_parser.py b/nerfactory/datamanagers/dataparsers/friends_parser.py index 3dd10844f..3bfbf0e7f 100644 --- a/nerfactory/datamanagers/dataparsers/friends_parser.py +++ b/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}}, ) diff --git a/nerfactory/datamanagers/datasets.py b/nerfactory/datamanagers/datasets.py index e988aeddf..13b430320 100644 --- a/nerfactory/datamanagers/datasets.py +++ b/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): diff --git a/nerfactory/fields/nerfw_field.py b/nerfactory/fields/nerfw_field.py index a68451bb7..b5e6fecf4 100644 --- a/nerfactory/fields/nerfw_field.py +++ b/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( diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py index 93206ef3d..3b3d6f5cc 100644 --- a/nerfactory/models/nerfw.py +++ b/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: diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 1334195f7..15bae3c48 100644 --- a/nerfactory/pipelines/base.py +++ b/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 diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 2d29f12c3..fc13724ba 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/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