Train and eval data splits (Nerfbusters) (#2207)

* nerfbusters dataparser

* nerfbuster dataparser and visualization of eval views

* resolved some errors from  pyright

* enable ns-process-data to take two videos as input and use one for eval and the other for train

* resolve some pyright issues

* instead of data being a tuple, I added eval_data as another argument

* implemented train eval split for all cases

* implemented train eval split for all cases

* reverting equi changes

* reverting reality capture changes

* raise valueerror

* revert comment

* added split override code back to dataparser

* reverted one change

* using image border now

* adding docs

---------

Co-authored-by: Ethan Weber <eweb0124@gmail.com>

nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py

@@ -19,16 +19,19 @@ Base class to process images or video into a nerfstudio dataset
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Optional
 
 
 @dataclass
 class BaseConverterToNerfstudioDataset(ABC):
-    """Base class to process images or video into a nerfstudio dataset"""
+    """Base class to process images or video into a nerfstudio dataset."""
 
     data: Path
     """Path the data, either a video file or a directory of images."""
     output_dir: Path
     """Path to the output directory."""
+    eval_data: Optional[Path] = None
+    """Path the eval data, either a video file or a directory of images. If set to None, the first will be used both for training and eval"""
     verbose: bool = False
     """If True, print extra logging."""
 

nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py

@@ -19,9 +19,7 @@ from pathlib import Path
 from typing import Dict, List, Literal, Optional, Tuple
 
 from nerfstudio.process_data import colmap_utils, hloc_utils, process_data_utils
-from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import (
-    BaseConverterToNerfstudioDataset,
-)
+from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import BaseConverterToNerfstudioDataset
 from nerfstudio.process_data.process_data_utils import CAMERA_MODELS
 from nerfstudio.utils import install_checks
 from nerfstudio.utils.rich_utils import CONSOLE

nerfstudio/process_data/images_to_nerfstudio_dataset.py

@@ -18,9 +18,7 @@ from dataclasses import dataclass
 from typing import Optional
 
 from nerfstudio.process_data import equirect_utils, process_data_utils
-from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import (
-    ColmapConverterToNerfstudioDataset,
-)
+from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import ColmapConverterToNerfstudioDataset
 from nerfstudio.utils.rich_utils import CONSOLE
 
 
@@ -47,11 +45,15 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset):
 
         # Generate planar projections if equirectangular
         if self.camera_type == "equirectangular":
+            if self.eval_data is not None:
+                raise ValueError("Cannot use eval_data with camera_type equirectangular.")
+
             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, crop_factor=self.crop_factor
             )
+
             self.camera_type = "perspective"
 
         summary_log = []
@@ -63,10 +65,25 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset):
                 self.data,
                 image_dir=self.image_dir,
                 crop_factor=self.crop_factor,
+                image_prefix="frame_train_" if self.eval_data is not None else "frame_",
                 verbose=self.verbose,
                 num_downscales=self.num_downscales,
                 same_dimensions=self.same_dimensions,
+                keep_image_dir=False,
             )
+            if self.eval_data is not None:
+                eval_image_rename_map_paths = process_data_utils.copy_images(
+                    self.eval_data,
+                    image_dir=self.image_dir,
+                    crop_factor=self.crop_factor,
+                    image_prefix="frame_eval_",
+                    verbose=self.verbose,
+                    num_downscales=self.num_downscales,
+                    same_dimensions=self.same_dimensions,
+                    keep_image_dir=True,
+                )
+                image_rename_map_paths.update(eval_image_rename_map_paths)
+
             image_rename_map = dict((a.name, b.name) for a, b in image_rename_map_paths.items())
             num_frames = len(image_rename_map)
             summary_log.append(f"Starting with {num_frames} images")

nerfstudio/process_data/process_data_utils.py

@@ -15,17 +15,16 @@
 """Helper utils for processing data into the nerfstudio format."""
 
 import math
+import re
 import shutil
 import sys
-import re
 from enum import Enum
 from pathlib import Path
 from typing import List, Literal, Optional, OrderedDict, Tuple, Union
 
 import cv2
-import numpy as np
-
 import imageio
+import numpy as np
 import rawpy
 
 from nerfstudio.utils.rich_utils import CONSOLE, status
@@ -114,6 +113,8 @@ def convert_video_to_images(
     num_downscales: int,
     crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
     verbose: bool = False,
+    image_prefix: str = "frame_",
+    keep_image_dir: bool = False,
 ) -> Tuple[List[str], int]:
     """Converts a video into a sequence of images.
 
@@ -124,10 +125,19 @@ def convert_video_to_images(
         num_downscales: Number of times to downscale the images. Downscales by 2 each time.
         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.
+        image_prefix: Prefix to use for the image filenames.
+        keep_image_dir: If True, don't delete the output directory if it already exists.
     Returns:
         A tuple containing summary of the conversion and the number of extracted frames.
     """
 
+    # If keep_image_dir is False, then remove the output image directory and its downscaled versions
+    if not keep_image_dir:
+        for i in range(num_downscales + 1):
+            dir_to_remove = image_dir if i == 0 else f"{image_dir}_{2**i}"
+            shutil.rmtree(dir_to_remove, ignore_errors=True)
+    image_dir.mkdir(exist_ok=True, parents=True)
+
     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].")
@@ -141,12 +151,6 @@ def convert_video_to_images(
         sys.exit(1)
 
     with status(msg="Converting video to images...", spinner="bouncingBall", verbose=verbose):
-        # delete existing images in folder
-        for img in image_dir.glob("*.png"):
-            if verbose:
-                CONSOLE.log(f"Deleting {img}")
-            img.unlink()
-
         num_frames = get_num_frames_in_video(video_path)
         if num_frames == 0:
             CONSOLE.print(f"[bold red]Error: Video has no frames: {video_path}")
@@ -168,7 +172,7 @@ def convert_video_to_images(
 
         downscale_chains = [f"[t{i}]scale=iw/{2**i}:ih/{2**i}[out{i}]" for i in range(num_downscales + 1)]
         downscale_dirs = [Path(str(image_dir) + (f"_{2**i}" if i > 0 else "")) for i in range(num_downscales + 1)]
-        downscale_paths = [downscale_dirs[i] / "frame_%05d.png" for i in range(num_downscales + 1)]
+        downscale_paths = [downscale_dirs[i] / f"{image_prefix}%05d.png" for i in range(num_downscales + 1)]
 
         for dir in downscale_dirs:
             dir.mkdir(parents=True, exist_ok=True)
@@ -200,7 +204,7 @@ def convert_video_to_images(
         num_final_frames = len(list(image_dir.glob("*.png")))
         summary_log = []
         summary_log.append(f"Starting with {num_frames} video frames")
-        summary_log.append(f"We extracted {num_final_frames} images")
+        summary_log.append(f"We extracted {num_final_frames} images with prefix '{image_prefix}'")
         CONSOLE.log("[bold green]:tada: Done converting video to images.")
 
         return summary_log, num_final_frames
@@ -210,9 +214,11 @@ def copy_images_list(
     image_paths: List[Path],
     image_dir: Path,
     num_downscales: int,
+    image_prefix: str = "frame_",
     crop_border_pixels: Optional[int] = None,
     crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
     verbose: bool = False,
+    keep_image_dir: bool = False,
     upscale_factor: Optional[int] = None,
     nearest_neighbor: bool = False,
     same_dimensions: bool = True,
@@ -222,19 +228,24 @@ def copy_images_list(
         image_paths: List of Paths of images to copy to a new directory.
         image_dir: Path to the output directory.
         num_downscales: Number of times to downscale the images. Downscales by 2 each time.
+        image_prefix: Prefix for the image filenames.
         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.
+        keep_image_dir: If True, don't delete the output directory if it already exists.
     Returns:
         A list of the copied image Paths.
     """
 
-    # Remove original directory only if we provide a proper image folder path
-    if image_dir.is_dir() and len(image_paths):
+    # Remove original directory and its downscaled versions
+    # only if we provide a proper image folder path and keep_image_dir is False
+    if image_dir.is_dir() and len(image_paths) and not keep_image_dir:
         # check that output directory is not the same as input directory
         if image_dir != image_paths[0].parent:
-            shutil.rmtree(image_dir, ignore_errors=True)
-        image_dir.mkdir(exist_ok=True, parents=True)
+            for i in range(num_downscales + 1):
+                dir_to_remove = image_dir if i == 0 else f"{image_dir}_{2**i}"
+                shutil.rmtree(dir_to_remove, ignore_errors=True)
+    image_dir.mkdir(exist_ok=True, parents=True)
 
     copied_image_paths = []
 
@@ -242,11 +253,11 @@ def copy_images_list(
     for idx, image_path in enumerate(image_paths):
         if verbose:
             CONSOLE.log(f"Copying image {idx + 1} of {len(image_paths)}...")
-        copied_image_path = image_dir / f"frame_{idx + 1:05d}{image_path.suffix}"
+        copied_image_path = image_dir / f"{image_prefix}{idx + 1:05d}{image_path.suffix}"
         try:
             # if CR2 raw, we want to read raw and write RAW_CONVERTED_SUFFIX, and change the file suffix for downstream processing
             if image_path.suffix.lower() in ALLOWED_RAW_EXTS:
-                copied_image_path = image_dir / f"frame_{idx + 1:05d}{RAW_CONVERTED_SUFFIX}"
+                copied_image_path = image_dir / f"{image_prefix}{idx + 1:05d}{RAW_CONVERTED_SUFFIX}"
                 with rawpy.imread(str(image_path)) as raw:
                     rgb = raw.postprocess()
                 imageio.imsave(copied_image_path, rgb)
@@ -283,7 +294,7 @@ def copy_images_list(
     # When this is not the case (e.g. mixed portrait and landscape images), we need to do individually.
     # (Unfortunately, that is much slower.)
     for framenum in range(1, (1 if same_dimensions else num_frames) + 1):
-        framename = "frame_%05d" if same_dimensions else f"frame_{framenum:05d}"
+        framename = f"{image_prefix}%05d" if same_dimensions else f"{image_prefix}{framenum:05d}"
         ffmpeg_cmd = f'ffmpeg -y -noautorotate -i "{image_dir / f"{framename}{copied_image_paths[0].suffix}"}" -q:v 2 '
 
         crop_cmd = ""
@@ -315,7 +326,7 @@ def copy_images_list(
     if num_frames == 0:
         CONSOLE.log("[bold red]:skull: No usable images in the data folder.")
     else:
-        CONSOLE.log("[bold green]:tada: Done copying images.")
+        CONSOLE.log(f"[bold green]:tada: Done copying images with prefix '{image_prefix}'.")
 
     return copied_image_paths
 
@@ -363,7 +374,9 @@ def copy_and_upscale_polycam_depth_maps_list(
 def copy_images(
     data: Path,
     image_dir: Path,
+    image_prefix: str = "frame_",
     verbose: bool = False,
+    keep_image_dir: bool = False,
     crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
     num_downscales: int = 0,
     same_dimensions: bool = True,
@@ -373,8 +386,10 @@ def copy_images(
     Args:
         data: Path to the directory of images.
         image_dir: Path to the output directory.
+        image_prefix: Prefix for the image filenames.
         verbose: If True, print extra logging.
         crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right)
+        keep_image_dir: If True, don't delete the output directory if it already exists.
     Returns:
         The mapping from the original filenames to the new ones.
     """
@@ -390,6 +405,8 @@ def copy_images(
             image_dir=image_dir,
             crop_factor=crop_factor,
             verbose=verbose,
+            image_prefix=image_prefix,
+            keep_image_dir=keep_image_dir,
             num_downscales=num_downscales,
             same_dimensions=same_dimensions,
         )
@@ -406,8 +423,6 @@ def downscale_images(
     """(Now deprecated; much faster integrated into copy_images.)
     Downscales the images in the directory. Uses FFMPEG.
 
-    Assumes images are named frame_00001.png, frame_00002.png, etc.
-
     Args:
         image_dir: Path to the directory containing the images.
         num_downscales: Number of times to downscale the images. Downscales by 2 each time.

nerfstudio/process_data/video_to_nerfstudio_dataset.py

@@ -18,9 +18,7 @@ import shutil
 from dataclasses import dataclass
 
 from nerfstudio.process_data import equirect_utils, process_data_utils
-from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import (
-    ColmapConverterToNerfstudioDataset,
-)
+from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import ColmapConverterToNerfstudioDataset
 from nerfstudio.utils.rich_utils import CONSOLE
 
 
@@ -35,7 +33,7 @@ class VideoToNerfstudioDataset(ColmapConverterToNerfstudioDataset):
     """
 
     num_frames_target: int = 300
-    """Target number of frames to use for the dataset, results may not be exact."""
+    """Target number of frames to use per video, results may not be exact."""
     percent_radius_crop: float = 1.0
     """Create circle crop mask. The radius is the percent of the image diagonal."""
 
@@ -43,6 +41,7 @@ class VideoToNerfstudioDataset(ColmapConverterToNerfstudioDataset):
         """Process video into a nerfstudio dataset."""
 
         summary_log = []
+        summary_log_eval = []
         # Convert video to images
         if self.camera_type == "equirectangular":
             # create temp images folder to store the equirect and perspective images
@@ -65,10 +64,28 @@ class VideoToNerfstudioDataset(ColmapConverterToNerfstudioDataset):
                 num_downscales=self.num_downscales,
                 crop_factor=self.crop_factor,
                 verbose=self.verbose,
+                image_prefix="frame_train_" if self.eval_data is not None else "frame_",
+                keep_image_dir=False,
             )
+            if self.eval_data is not None:
+                summary_log_eval, num_extracted_frames_eval = process_data_utils.convert_video_to_images(
+                    self.eval_data,
+                    image_dir=self.image_dir,
+                    num_frames_target=self.num_frames_target,
+                    num_downscales=self.num_downscales,
+                    crop_factor=self.crop_factor,
+                    verbose=self.verbose,
+                    image_prefix="frame_eval_",
+                    keep_image_dir=True,
+                )
+                summary_log += summary_log_eval
+                num_extracted_frames += num_extracted_frames_eval
 
         # Generate planar projections if equirectangular
         if self.camera_type == "equirectangular":
+            if self.eval_data is not None:
+                raise ValueError("Cannot use eval_data with camera_type equirectangular.")
+
             perspective_image_size = equirect_utils.compute_resolution_from_equirect(
                 self.output_dir / "temp_images", self.images_per_equirect
             )
@@ -97,7 +114,7 @@ class VideoToNerfstudioDataset(ColmapConverterToNerfstudioDataset):
                 process_data_utils.downscale_images(self.image_dir, self.num_downscales, verbose=self.verbose)
             )
 
-        # # Create mask
+        # Create mask
         mask_path = process_data_utils.save_mask(
             image_dir=self.image_dir,
             num_downscales=self.num_downscales,