diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md index faae5bf60..009cb287c 100644 --- a/docs/quickstart/custom_dataset.md +++ b/docs/quickstart/custom_dataset.md @@ -310,15 +310,15 @@ For videos we reccomend taking a video with the camera held on top of your head. For a video, process the data using the following command: ```bash -ns-process-data video --camera-type equirectangular --images-per-equirect {8, or 14} --num-frames-target {num equirectangular frames to sample from} --percent-crop {top bottom left right} --data {data directory} --output-dir {output directory} +ns-process-data video --camera-type equirectangular --images-per-equirect {8, or 14} --num-frames-target {num equirectangular frames to sample from} --crop-factor {top bottom left right} --data {data directory} --output-dir {output directory} ``` See the equirectangular images section above for a description of the `--images-per-equirect` argument. The `num-frames-target` argument is optional but it is recomended to set it to 3*(seconds of video) frames. For example, if you have a 30 second video, you would use `--num-frames-target 90` (3*30=90). This number was chosen from a bit of experimentation and seems to work well for most videos. It is by no means a hard rule and you can experiment with different values. -The `percent-crop` argument is optional but often very helpful. This is because equirectangular videos taken by 360 cameras tend to have a portion of the bottom of the image that is the person who was holding the camera over their head. +The `crop-factor` argument is optional but often very helpful. This is because equirectangular videos taken by 360 cameras tend to have a portion of the bottom of the image that is the person who was holding the camera over their head. -This obscene mesh of human is obviously not helpful in training a nerf so we can remove it by cropping the bottom 20% of the image. This can be done by using the `--percent-crop 0 0.2 0 0` argument. +This obscene mesh of human is obviously not helpful in training a nerf so we can remove it by cropping the bottom 20% of the image. This can be done by using the `--crop-factor 0 0.2 0 0` argument. diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index 70690c4c3..340dd7884 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -15,6 +15,7 @@ """Helper utils for processing equirectangular data.""" import os +import sys from pathlib import Path from typing import List, Tuple @@ -22,6 +23,7 @@ import cv2 import numpy as np import torch from equilib import Equi2Pers +from rich.console import Console from rich.progress import ( BarColumn, Progress, @@ -32,19 +34,21 @@ from rich.progress import ( from nerfstudio.utils.rich_utils import ItersPerSecColumn +CONSOLE = Console(width=120) -def _crop_bottom(bound_arr: list, fov: int, percent_crop: float) -> List[float]: + +def _crop_bottom(bound_arr: list, fov: int, crop_factor: float) -> List[float]: """Returns a list of vertical bounds with the bottom cropped. Args: bound_arr (list): List of vertical bounds in ascending order. fov (int): Field of view of the camera. - percent_crop (float): Percent of the image to crop from the bottom. + crop_factor (float): Portion of the image to crop from the bottom. Returns: list: A new list of bounds with the bottom cropped. """ - degrees_chopped = 180 * percent_crop + degrees_chopped = 180 * crop_factor new_bottom_start = 90 - degrees_chopped - fov / 2 for i, el in reversed(list(enumerate(bound_arr))): if el > new_bottom_start + fov / 2: @@ -59,18 +63,18 @@ def _crop_bottom(bound_arr: list, fov: int, percent_crop: float) -> List[float]: return bound_arr -def _crop_top(bound_arr: list, fov: int, percent_crop: float) -> List[float]: +def _crop_top(bound_arr: list, fov: int, crop_factor: float) -> List[float]: """Returns a list of vertical bounds with the top cropped. Args: bound_arr (list): List of vertical bounds in ascending order. fov (int): Field of view of the camera. - percent_crop (float): Percent of the image to crop from the top. + crop_factor (float): Portion of the image to crop from the top. Returns: list: A new list of bounds with the top cropped. """ - degrees_chopped = 180 * percent_crop + degrees_chopped = 180 * crop_factor new_top_start = -90 + degrees_chopped + fov / 2 for i, el in enumerate(bound_arr): if el < new_top_start - fov / 2: @@ -86,22 +90,22 @@ def _crop_top(bound_arr: list, fov: int, percent_crop: float) -> List[float]: def _crop_bound_arr_vertical( - bound_arr: list, fov: int, percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) + bound_arr: list, fov: int, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) ) -> list: """Returns a list of vertical bounds adjusted for cropping. Args: bound_arr (list): Original list of vertical bounds in ascending order. fov (int): Field of view of the camera. - percent_crop (Tuple[float, float, float, float]): Crop arr (top, bottom, left, right). + crop_factor (Tuple[float, float, float, float]): Crop arr (top, bottom, left, right). Returns: - list: _description_ + list: Cropped bound arr """ - if percent_crop[1] > 0: - bound_arr = _crop_bottom(bound_arr, fov, percent_crop[1]) - if percent_crop[0] > 0: - bound_arr = _crop_top(bound_arr, fov, percent_crop[0]) + if crop_factor[1] > 0: + bound_arr = _crop_bottom(bound_arr, fov, crop_factor[1]) + if crop_factor[0] > 0: + bound_arr = _crop_top(bound_arr, fov, crop_factor[0]) return bound_arr @@ -109,7 +113,7 @@ def generate_planar_projections_from_equirectangular( image_dir: Path, planar_image_size: Tuple[int, int], samples_per_im: int, - percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), ) -> Path: """Generate planar projections from an equirectangular image. @@ -117,25 +121,31 @@ def generate_planar_projections_from_equirectangular( image_dir: The directory containing the equirectangular image. planar_image_size: The size of the planar projections [width, height]. samples_per_im: The number of samples to take per image. - percent_crop: The percent of the image to crop from the (top, bottom, left, and right). + crop_factor: The portion of the image to crop from the (top, bottom, left, and right). + Values should be in [0, 1]. returns: The path to the planar projections directory. """ + for i in crop_factor: + if i < 0 or i > 1: + CONSOLE.print("[bold red] Invalid crop factor. All values must be in [0,1].") + sys.exit(1) + device = torch.device("cuda") fov = 120 yaw_pitch_pairs = [] left_bound, right_bound = -180, 180 - if percent_crop[3] > 0: - left_bound = -180 + 360 * percent_crop[3] - if percent_crop[2] > 0: - right_bound = 180 - 360 * percent_crop[2] + if crop_factor[3] > 0: + left_bound = -180 + 360 * crop_factor[3] + if crop_factor[2] > 0: + right_bound = 180 - 360 * crop_factor[2] if samples_per_im == 8: fov = 120 bound_arr = [-45, 0, 45] - bound_arr = _crop_bound_arr_vertical(bound_arr, fov, percent_crop) + bound_arr = _crop_bound_arr_vertical(bound_arr, fov, crop_factor) if bound_arr[1] is not None: for i in np.arange(left_bound, right_bound, 90): yaw_pitch_pairs.append((i, bound_arr[1])) @@ -148,7 +158,7 @@ def generate_planar_projections_from_equirectangular( elif samples_per_im == 14: fov = 110 bound_arr = [-45, 0, 45] - bound_arr = _crop_bound_arr_vertical(bound_arr, fov, percent_crop) + bound_arr = _crop_bound_arr_vertical(bound_arr, fov, crop_factor) if bound_arr[1] is not None: for i in np.arange(left_bound, right_bound, 60): yaw_pitch_pairs.append((i, bound_arr[1])) @@ -172,7 +182,6 @@ def generate_planar_projections_from_equirectangular( TimeRemainingColumn(elapsed_when_finished=True, compact=True), ) - curr_im = 0 with progress: for i in progress.track(os.listdir(frame_dir), description="", total=num_ims): if i.lower().endswith((".jpg", ".png", ".jpeg")): @@ -187,7 +196,6 @@ def generate_planar_projections_from_equirectangular( pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy() cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.jpg", pers_image) count += 1 - curr_im += 1 return output_dir diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 2734883cb..14ae33662 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -103,7 +103,7 @@ def convert_video_to_images( video_path: Path, image_dir: Path, num_frames_target: int, - percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), verbose: bool = False, ) -> Tuple[List[str], int]: """Converts a video into a sequence of images. @@ -112,12 +112,17 @@ def convert_video_to_images( video_path: Path to the video. output_dir: Path to the output directory. num_frames_target: Number of frames to extract. - percent_crop: Percent of the image to crop. (top, bottom, left, right) + crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right) verbose: If True, logs the output of the command. Returns: A tuple containing summary of the conversion and the number of extracted frames. """ + for i in crop_factor: + if i < 0 or i > 1: + CONSOLE.print("[bold red]Error: Invalid crop factor. All crops must be in [0,1].") + sys.exit(1) + if video_path.is_dir(): CONSOLE.print(f"[bold red]Error: Video path is a directory, not a path: {video_path}") sys.exit(1) @@ -142,11 +147,11 @@ def convert_video_to_images( ffmpeg_cmd = f'ffmpeg -i "{video_path}"' crop_cmd = "" - if percent_crop != (0.0, 0.0, 0.0, 0.0): - height = 1 - percent_crop[0] - percent_crop[1] - width = 1 - percent_crop[2] - percent_crop[3] - start_x = percent_crop[2] - start_y = percent_crop[0] + if crop_factor != (0.0, 0.0, 0.0, 0.0): + 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}"' spacing = num_frames // num_frames_target @@ -411,25 +416,23 @@ def generate_circle_mask(height: int, width: int, percent_radius) -> Optional[np return mask -def generate_crop_mask( - height: int, width: int, percent_crop: Tuple[float, float, float, float] -) -> Optional[np.ndarray]: +def generate_crop_mask(height: int, width: int, crop_factor: Tuple[float, float, float, float]) -> Optional[np.ndarray]: """generate a crop mask of the given size. Args: height: The height of the mask. width: The width of the mask. - percent_crop: The percent of the image to crop in each direction [top, bottom, left, right]. + crop_factor: The percent of the image to crop in each direction [top, bottom, left, right]. Returns: The mask or None if no cropping is performed. """ - if np.all(np.array(percent_crop) == 0.0): + if np.all(np.array(crop_factor) == 0.0): return None - if np.any(np.array(percent_crop) < 0.0) or np.any(np.array(percent_crop) > 1.0): + if np.any(np.array(crop_factor) < 0.0) or np.any(np.array(crop_factor) > 1.0): CONSOLE.log("[bold red]Invalid crop percentage, must be between 0 and 1.") sys.exit(1) - top, bottom, left, right = percent_crop + top, bottom, left, right = crop_factor mask = np.zeros((height, width), dtype=np.uint8) top = int(top * height) bottom = int(bottom * height) @@ -440,20 +443,20 @@ def generate_crop_mask( def generate_mask( - height: int, width: int, percent_crop: Tuple[float, float, float, float], percent_radius: float + height: int, width: int, crop_factor: Tuple[float, float, float, float], percent_radius: float ) -> Optional[np.ndarray]: """generate a mask of the given size. Args: height: The height of the mask. width: The width of the mask. - percent_crop: The percent of the image to crop in each direction [top, bottom, left, right]. + crop_factor: The percent of the image to crop in each direction [top, bottom, left, right]. percent_radius: The radius of the circle as a percentage of the image diagonal size. Returns: The mask or None if no mask is needed. """ - crop_mask = generate_crop_mask(height, width, percent_crop) + crop_mask = generate_crop_mask(height, width, crop_factor) circle_mask = generate_circle_mask(height, width, percent_radius) if crop_mask is None: return circle_mask @@ -465,7 +468,7 @@ def generate_mask( def save_mask( image_dir: Path, num_downscales: int, - percent_crop: Tuple[float, float, float, float] = (0, 0, 0, 0), + crop_factor: Tuple[float, float, float, float] = (0, 0, 0, 0), percent_radius: float = 1.0, ) -> Optional[Path]: """Save a mask for each image in the image directory. @@ -473,7 +476,7 @@ def save_mask( Args: image_dir: The directory containing the images. num_downscales: The number of downscaling levels. - percent_crop: The percent of the image to crop in each direction [top, bottom, left, right]. + crop_factor: The percent of the image to crop in each direction [top, bottom, left, right]. percent_radius: The radius of the circle as a percentage of the image diagonal size. Returns: @@ -482,7 +485,7 @@ def save_mask( image_path = next(image_dir.glob("frame_*")) image = cv2.imread(str(image_path)) height, width = image.shape[:2] - mask = generate_mask(height, width, percent_crop, percent_radius) + mask = generate_mask(height, width, crop_factor, percent_radius) if mask is None: return None mask *= 255 diff --git a/scripts/process_data.py b/scripts/process_data.py index e2e40de03..570fda59e 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -292,8 +292,8 @@ class ProcessVideo: """ percent_radius_crop: float = 1.0 """Create circle crop mask. The radius is the percent of the image diagonal.""" - percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) - """Percent of the image to crop. (top, bottom, left, right)""" + 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)""" use_sfm_depth: bool = False """If True, export and use depth maps induced from SfM points.""" include_depth_debug: bool = False @@ -322,7 +322,7 @@ class ProcessVideo: self.data, image_dir=temp_image_dir, num_frames_target=self.num_frames_target, - percent_crop=(0.0, 0.0, 0.0, 0.0), + crop_factor=(0.0, 0.0, 0.0, 0.0), verbose=self.verbose, ) else: @@ -330,7 +330,7 @@ class ProcessVideo: self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, - percent_crop=self.percent_crop, + crop_factor=self.crop_factor, verbose=self.verbose, ) @@ -343,7 +343,7 @@ class ProcessVideo: self.output_dir / "temp_images", perspective_image_size, self.images_per_equirect, - percent_crop=self.percent_crop, + crop_factor=self.crop_factor, ) # copy the perspective images to the image directory @@ -361,7 +361,7 @@ class ProcessVideo: mask_path = process_data_utils.save_mask( image_dir=image_dir, num_downscales=self.num_downscales, - percent_crop=(0.0, 0.0, 0.0, 0.0), + crop_factor=(0.0, 0.0, 0.0, 0.0), percent_radius=self.percent_radius_crop, ) if mask_path is not None: