From b61fea866687ce798385e884523f25d2f66a0132 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 5 Sep 2022 18:33:46 +0100 Subject: [PATCH] Remove downscaling from dataloading (#340) --- docs/tutorials/data/choosing_dataset.md | 2 -- docs/tutorials/data/custom_dataset_format.md | 3 -- nerfactory/configs/base.py | 3 -- nerfactory/dataloaders/base.py | 2 -- nerfactory/dataloaders/data_parsers.py | 34 ++++++-------------- nerfactory/dataloaders/image_dataset.py | 32 +----------------- nerfactory/dataloaders/structs.py | 1 - scripts/process_data.py | 30 +++++++++++++++++ tests/test_run_train.py | 6 ++-- 9 files changed, 43 insertions(+), 70 deletions(-) diff --git a/docs/tutorials/data/choosing_dataset.md b/docs/tutorials/data/choosing_dataset.md index 993014e07..ccb2d6389 100644 --- a/docs/tutorials/data/choosing_dataset.md +++ b/docs/tutorials/data/choosing_dataset.md @@ -9,12 +9,10 @@ data: dataset_inputs_train: data_directory: data/instant_ngp/fox dataset_format: instant_ngp - downscale_factor: 1 alpha_color: white dataset_inputs_eval: data_directory: data/instant_ngp/fox dataset_format: instant_ngp - downscale_factor: 1 alpha_color: white ``` diff --git a/docs/tutorials/data/custom_dataset_format.md b/docs/tutorials/data/custom_dataset_format.md index 1b42acb01..597ad288b 100644 --- a/docs/tutorials/data/custom_dataset_format.md +++ b/docs/tutorials/data/custom_dataset_format.md @@ -12,7 +12,6 @@ class DatasetInputs: """Dataset inputs are used to initialize datasets and the NeRF graph.""" image_filenames: List[str] - downscale_factor: int = 1 intrinsics: torch.tensor = None camera_to_world: torch.tensor = None mask_filenames: List[str] = None @@ -26,7 +25,6 @@ class DatasetInputs: # See `nerfactory/data/format/instant_ngp.py` for the code. def load_instant_ngp_data( basedir: str, - downscale_factor: int = 1, split: str = "train", camera_translation_scalar=0.33 ) -> DatasetInputs: """Returns a DatasetInputs struct.""" @@ -42,7 +40,6 @@ def get_dataset_inputs( data_directory: str, dataset_format: str, split: str, - downscale_factor: int = 1, alpha_color: Optional[Union[str, list, ListConfig]] = None, ) -> DatasetInputs: """Makes a call to `load__data` and returns a DatasetInputs struct.""" diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index bde0e0d0b..3baaed5bc 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -189,7 +189,6 @@ class BlenderDataParserConfig(DataParserConfig): data_directory: Path = Path("data/blender/lego") scale_factor: float = 1.0 alpha_color: str = "white" - downscale_factor: int = 1 @dataclass @@ -219,7 +218,6 @@ class InstantNGPDataParserConfig(DataParserConfig): _target: Type = InstantNGP data_directory: Path = Path("data/ours/posterv2") scale_factor: float = 1.0 - downscale_factor: int = 1 scene_scale: float = 0.33 @@ -229,7 +227,6 @@ class Record3DDataParserConfig(DataParserConfig): _target: Type = Record3D data_directory: Path = Path("data/record3d/garden") - downscale_factor: int = 1 val_skip: int = 8 aabb_scale = 4.0 max_dataset_size: int = 150 diff --git a/nerfactory/dataloaders/base.py b/nerfactory/dataloaders/base.py index 15c3b8185..d4df44aa8 100644 --- a/nerfactory/dataloaders/base.py +++ b/nerfactory/dataloaders/base.py @@ -253,7 +253,6 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method if self.config.image_dataset_type == "rgb": self.train_image_dataset = ImageDataset( image_filenames=self.train_datasetinputs.image_filenames, - downscale_factor=self.train_datasetinputs.downscale_factor, alpha_color=self.train_datasetinputs.alpha_color, ) elif self.config.image_dataset_type == "panoptic": @@ -290,7 +289,6 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method if self.config.image_dataset_type == "rgb": self.eval_image_dataset = ImageDataset( image_filenames=self.eval_datasetinputs.image_filenames, - downscale_factor=self.eval_datasetinputs.downscale_factor, alpha_color=self.eval_datasetinputs.alpha_color, ) elif self.config.image_dataset_type == "panoptic": diff --git a/nerfactory/dataloaders/data_parsers.py b/nerfactory/dataloaders/data_parsers.py index 36d6d6d44..0c4b895be 100644 --- a/nerfactory/dataloaders/data_parsers.py +++ b/nerfactory/dataloaders/data_parsers.py @@ -140,7 +140,7 @@ class Nerfactory(DataParser): scene_scale: How much to scale the scene. Defaults to 0.33 """ - config: cfg.InstantNGPDataParserConfig + config: cfg.NerfactoryDataParserConfig def _generate_dataset_inputs(self, split="train"): @@ -152,6 +152,10 @@ class Nerfactory(DataParser): num_skipped_image_filenames = 0 for frame in meta["frames"]: fname = abs_dir / Path(frame["file_path"]) + if self.config.downscale_factor > 1: + fname = abs_dir / f"images_{self.config.downscale_factor}" / Path(frame["file_path"]).name + else: + fname = abs_dir / Path(frame["file_path"]) if not fname: num_skipped_image_filenames += 1 else: @@ -199,7 +203,6 @@ class Nerfactory(DataParser): dataset_inputs = DatasetInputs( image_filenames=image_filenames, - downscale_factor=self.config.downscale_factor, cameras=cameras, scene_bounds=scene_bounds, ) @@ -215,7 +218,6 @@ class Blender(DataParser): data_directory: Location of data alpha_color: Sets transparent regions to specified color, otherwise black. scale_factor: How much to scale the camera origins by. - downscale_factor: How much to downscale images. Defaults to 1. """ def __init__(self, config: cfg.BlenderDataParserConfig): @@ -223,7 +225,6 @@ class Blender(DataParser): self.data_directory: Path = config.data_directory self.scale_factor: float = config.scale_factor self.alpha_color = config.alpha_color - self.downscale_factor: int = config.downscale_factor def _generate_dataset_inputs(self, split="train"): if self.alpha_color is not None: @@ -263,12 +264,9 @@ class Blender(DataParser): camera_type=CameraType.PERSPECTIVE, ) - cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor) - dataset_inputs = DatasetInputs( image_filenames=image_filenames, cameras=cameras, - downscale_factor=self.downscale_factor, alpha_color=alpha_color_tensor, scene_bounds=scene_bounds, ) @@ -283,7 +281,6 @@ class InstantNGP(DataParser): Args: data_directory: Location of data scale_factor: How much to scale the camera origins by. - downscale_factor: How much to downscale images. Defaults to 1. scene_scale: How much to scale the scene. Defaults to 0.33 """ @@ -342,12 +339,9 @@ class InstantNGP(DataParser): camera_type=CameraType.PERSPECTIVE, ) - cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) - # TODO(ethan): add alpha background color dataset_inputs = DatasetInputs( image_filenames=image_filenames, - downscale_factor=self.config.downscale_factor, cameras=cameras, scene_bounds=scene_bounds, ) @@ -428,7 +422,7 @@ class Mipnerf360(DataParser): image_height, image_width = img_0.shape[:2] poses[:, :2, 4] = np.array([image_height, image_width]) - poses[:, 2, 4] = poses[:, 2, 4] * 1.0 / self.config.downscale_factor + poses[:, 2, 4] = poses[:, 2, 4] # Reorder pose to match our convention poses = np.concatenate([poses[:, :, 1:2], -poses[:, :, 0:1], poses[:, :, 2:]], axis=-1) @@ -471,7 +465,6 @@ class Mipnerf360(DataParser): dataset_inputs = DatasetInputs( image_filenames=image_filenames, - downscale_factor=1, cameras=cameras, scene_bounds=scene_bounds, ) @@ -485,7 +478,6 @@ class Record3D(DataParser): Args: data_directory: Location of data - downscale_factor: How much to downscale images. Defaults to 1. val_skip: 1/val_skip images to use for validation. Defaults to 8. aabb_scale: Scene scale, Defaults to 4.0. max_dataset_size: Max number of images to train on. If the dataset has @@ -590,11 +582,8 @@ class Record3D(DataParser): camera_type=CameraType.PERSPECTIVE, ) - cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) - dataset_inputs = DatasetInputs( image_filenames=image_filenames, - downscale_factor=self.config.downscale_factor, cameras=cameras, scene_bounds=scene_bounds, ) @@ -608,13 +597,11 @@ class Friends(DataParser): Args: data_directory: Location of data - downscale_factor: How much to downscale images. Defaults to 1. 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 - downscale_factor: int = 1 include_semantics: bool = True include_point_cloud: bool = False @@ -732,17 +719,16 @@ class Friends(DataParser): point_cloud.rgb = rgb cameras = Cameras( - fx=focal_lengths / self.downscale_factor, - fy=focal_lengths / self.downscale_factor, - cx=cx / self.downscale_factor, - cy=cy / self.downscale_factor, + fx=focal_lengths, + fy=focal_lengths, + cx=cx, + cy=cy, camera_to_worlds=camera_to_world, camera_type=CameraType.PERSPECTIVE, ) dataset_inputs = DatasetInputs( image_filenames=image_filenames, - downscale_factor=self.downscale_factor, cameras=cameras, semantics=semantics, point_cloud=point_cloud, diff --git a/nerfactory/dataloaders/image_dataset.py b/nerfactory/dataloaders/image_dataset.py index 6e8fb1676..c266ffa0f 100644 --- a/nerfactory/dataloaders/image_dataset.py +++ b/nerfactory/dataloaders/image_dataset.py @@ -37,20 +37,16 @@ class ImageDataset(Dataset): def __init__( self, image_filenames: List[Path], - downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None, ): """_summary_ Args: image_filenames (List[str]): List of image filenames - downscale_factor (int, optional): How much to downscale the image. Defaults to 1. alpha_color (TensorType[3], optional): Sets transparent regions to specified color, otherwise black. """ super().__init__() - assert isinstance(downscale_factor, int) self.image_filenames = image_filenames - self.downscale_factor = downscale_factor self.alpha_color = alpha_color def __len__(self): @@ -67,19 +63,6 @@ class ImageDataset(Dataset): """ image_filename = self.image_filenames[image_idx] pil_image = Image.open(image_filename) - if self.downscale_factor != 1.0: - image_width, image_height = pil_image.size - if image_width % self.downscale_factor != 0: - raise ValueError( - f"Image width {image_width} is not divisible by downscale_factor {self.downscale_factor}" - ) - if image_height % self.downscale_factor != 0: - raise ValueError( - f"Image height {image_height} is not divisible by downscale_factor {self.downscale_factor}" - ) - pil_image = pil_image.resize( - (image_width // self.downscale_factor, image_height // self.downscale_factor), Image.BILINEAR - ) image = np.array(pil_image, dtype="uint8") # shape is (h, w, 3 or 4) assert len(image.shape) == 3 assert image.dtype == np.uint8 @@ -134,24 +117,17 @@ class PanopticImageDataset(ImageDataset): self, semantics: Semantics, image_filenames: List[Path], - downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None, **kwargs, ): self.semantics = semantics self.person_index = self.semantics.thing_classes.index("person") - super().__init__(image_filenames, downscale_factor, alpha_color, **kwargs) + super().__init__(image_filenames, alpha_color, **kwargs) def get_mask(self, image_idx): """Mask out the people. Valid only where there aren't people.""" thing_image_filename = self.semantics.thing_filenames[image_idx] pil_image = Image.open(thing_image_filename) - if self.downscale_factor != 1.0: - image_width, image_height = pil_image.size - # the use of NEAREST is important for semantic classes - pil_image = pil_image.resize( - (image_width // self.downscale_factor, image_height // self.downscale_factor), Image.NEAREST - ) thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] mask = (thing_semantics != self.person_index).to(torch.float32) # 1 where valid return mask @@ -159,11 +135,5 @@ class PanopticImageDataset(ImageDataset): def get_semantics(self, image_idx): stuff_image_filename = self.semantics.stuff_filenames[image_idx] pil_image = Image.open(stuff_image_filename) - if self.downscale_factor != 1.0: - image_width, image_height = pil_image.size - # the use of NEAREST is important for semantic classes - pil_image = pil_image.resize( - (image_width // self.downscale_factor, image_height // self.downscale_factor), Image.NEAREST - ) stuff_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] return stuff_semantics diff --git a/nerfactory/dataloaders/structs.py b/nerfactory/dataloaders/structs.py index be6c74b86..2d5fac18a 100644 --- a/nerfactory/dataloaders/structs.py +++ b/nerfactory/dataloaders/structs.py @@ -121,7 +121,6 @@ class DatasetInputs: image_filenames: List[Path] cameras: Cameras - downscale_factor: int = 1 mask_filenames: Optional[List[Path]] = None depth_filenames: Optional[List[Path]] = None scene_bounds: SceneBounds = SceneBounds() diff --git a/scripts/process_data.py b/scripts/process_data.py index 95998f0bd..8b3672b8a 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -238,11 +238,35 @@ def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path) -> N json.dump(out, f, indent=4) +def downscale_images(image_dir: Path, num_downscales: int) -> None: + """Downscales the images in the directory. + + Args: + image_dir: Path to the directory containing the images. + num_downscales: Number of times to downscale the images. Downscales by 2 each time. + """ + downscale_factors = [2**i for i in range(num_downscales + 1)[1:]] + for downscale_factor in downscale_factors: + # Downscale images with ffmpeg + assert downscale_factor > 1 + assert isinstance(downscale_factor, int) + downscale_dir = image_dir.parent / f"images_{downscale_factor}" + downscale_dir.mkdir(parents=True, exist_ok=True) + ffmpeg_cmd = [ + f"ffmpeg -i {image_dir}/frame_%05d.png ", + f"-vf scale=iw/{downscale_factor}:ih/{downscale_factor} ", + f"{downscale_dir}/frame_%05d.png", + ] + ffmpeg_cmd = " ".join(ffmpeg_cmd) + run_command(ffmpeg_cmd) + + def main( data: Path, output_dir: Path, num_frames_target: int, camera_type: Literal["perspective", "fisheye"] = "perspective", + num_downscales: int = 0, gpu: bool = True, ): """Process images or videos into a Nerfactory dataset. @@ -259,6 +283,8 @@ def main( output_dir: Path to the output directory. num_frames: Target number of frames to use for the dataset, results may not be exact. camera_type: Camera model to use. + num_downscales: Number of times to downscale the images. Downscales by 2 each time. For example a value of 3 + will downscale the images by 2x, 4x, and 8x. gpu: If True, use GPU. """ @@ -275,6 +301,10 @@ def main( _print_section_header("Converting video into images.") convert_video_to_images(data, image_dir=image_dir, num_frames_target=num_frames_target) + if num_downscales > 0: + _print_section_header("Downscaling images.") + downscale_images(image_dir, num_downscales) + _print_section_header("Running Colmap.") colmap_dir = output_dir / "colmap" colmap_dir.mkdir(parents=True, exist_ok=True) diff --git a/tests/test_run_train.py b/tests/test_run_train.py index 8fc7a172b..0e2a5774d 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -25,13 +25,11 @@ def set_reduced_config(config: cfg.Config): config.pipeline.data_manager.image_dataset_type = "rgb" config.pipeline.data_manager.train_dataset = cfg.BlenderDataParserConfig( - data_directory=Path("tests/data/lego_test"), downscale_factor=16 + data_directory=Path("tests/data/lego_test") ) config.pipeline.data_manager.train_num_images_to_sample_from = 1 config.pipeline.data_manager.train_num_rays_per_batch = 4 - config.pipeline.data_manager.eval_dataset = cfg.BlenderDataParserConfig( - data_directory=Path("tests/data/lego_test"), downscale_factor=16 - ) + config.pipeline.data_manager.eval_dataset = cfg.BlenderDataParserConfig(data_directory=Path("tests/data/lego_test")) # reduce model factors config.pipeline.model.num_coarse_samples = 4