From 2838cfc8a57629778abbda9670798c7f1b08ecfe Mon Sep 17 00:00:00 2001 From: Rohan Mathur Date: Tue, 14 Mar 2023 12:04:28 -0700 Subject: [PATCH] Adding the crop code made for videos to images (#1598) * added equirect support * added cropping for planar image input * random comment * bug fix and running command * added crop-factor to docstrings --- nerfstudio/process_data/process_data_utils.py | 30 ++++++++++++++++--- scripts/process_data.py | 6 ++-- 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 14ae33662..a73127fcb 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -159,7 +159,7 @@ def convert_video_to_images( ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB{crop_cmd} -r 1" else: CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.") - ffmpeg_cmd += " -pix_fmt bgr8" + ffmpeg_cmd += f" -pix_fmt bgr8 -vf {crop_cmd[1:]}" ffmpeg_cmd += f" {out_filename}" run_command(ffmpeg_cmd, verbose=verbose) @@ -174,13 +174,18 @@ def convert_video_to_images( def copy_images_list( - image_paths: List[Path], image_dir: Path, crop_border_pixels: Optional[int] = None, verbose: bool = False + image_paths: List[Path], + image_dir: Path, + crop_border_pixels: Optional[int] = None, + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), + verbose: bool = False, ) -> List[Path]: """Copy all images in a list of Paths. Useful for filtering from a directory. Args: image_paths: List of Paths of images to copy to a new directory. image_dir: Path to the output directory. crop_border_pixels: If not None, crops each edge by the specified number of pixels. + crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right) verbose: If True, print extra logging. Returns: A list of the copied image Paths. @@ -207,6 +212,18 @@ def copy_images_list( crop = f"crop=iw-{crop_border_pixels*2}:ih-{crop_border_pixels*2}" ffmpeg_cmd = f'ffmpeg -y -noautorotate -i "{image_dir / filename}" -q:v 2 -vf {crop} "{image_dir / filename}"' run_command(ffmpeg_cmd, verbose=verbose) + elif crop_factor != (0.0, 0.0, 0.0, 0.0): + file_type = image_paths[0].suffix + filename = f"frame_%05d{file_type}" + height = 1 - crop_factor[0] - crop_factor[1] + width = 1 - crop_factor[2] - crop_factor[3] + start_x = crop_factor[2] + start_y = crop_factor[0] + crop_cmd = f',"crop=w=iw*{width}:h=ih*{height}:x=iw*{start_x}:y=ih*{start_y}"' + ffmpeg_cmd = ( + f'ffmpeg -y -noautorotate -i "{image_dir / filename}" -q:v 2 -vf {crop_cmd[1:]} "{image_dir / filename}"' + ) + run_command(ffmpeg_cmd, verbose=verbose) num_frames = len(image_paths) @@ -266,13 +283,16 @@ def copy_and_upscale_polycam_depth_maps_list( return copied_depth_map_paths -def copy_images(data: Path, image_dir: Path, verbose) -> OrderedDict[Path, Path]: +def copy_images( + data: Path, image_dir: Path, verbose, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) +) -> OrderedDict[Path, Path]: """Copy images from a directory to a new directory. Args: data: Path to the directory of images. image_dir: Path to the output directory. verbose: If True, print extra logging. + crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right) Returns: The mapping from the original filenames to the new ones. """ @@ -283,7 +303,9 @@ def copy_images(data: Path, image_dir: Path, verbose) -> OrderedDict[Path, Path] CONSOLE.log("[bold red]:skull: No usable images in the data folder.") sys.exit(1) - copied_images = copy_images_list(image_paths=image_paths, image_dir=image_dir, verbose=verbose) + copied_images = copy_images_list( + image_paths=image_paths, image_dir=image_dir, crop_factor=crop_factor, verbose=verbose + ) return OrderedDict((original_path, new_path) for original_path, new_path in zip(image_paths, copied_images)) diff --git a/scripts/process_data.py b/scripts/process_data.py index 2e08bf9a9..2b0647434 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -91,6 +91,8 @@ class ProcessImages: """Number of samples per image to take from each equirectangular image. Used only when camera-type is equirectangular. """ + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) + """Portion of the image to crop. All values should be in [0,1]. (top, bottom, left, right)""" gpu: bool = True """If True, use GPU.""" use_sfm_depth: bool = False @@ -129,7 +131,7 @@ class ProcessImages: pers_size = equirect_utils.compute_resolution_from_equirect(self.data, self.images_per_equirect) CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image") self.data = equirect_utils.generate_planar_projections_from_equirectangular( - self.data, pers_size, self.images_per_equirect + self.data, pers_size, self.images_per_equirect, crop_factor=self.crop_factor ) summary_log = [] @@ -138,7 +140,7 @@ class ProcessImages: if not self.skip_image_processing: # Copy images to output directory image_rename_map_paths = process_data_utils.copy_images( - self.data, image_dir=image_dir, verbose=self.verbose + self.data, image_dir=image_dir, crop_factor=self.crop_factor, verbose=self.verbose ) image_rename_map = dict((a.name, b.name) for a, b in image_rename_map_paths.items()) num_frames = len(image_rename_map)