Skip to content

Commit

Permalink
Allow skipping copy+downscale of images in ns-process-data, fix bug w…
Browse files Browse the repository at this point in the history
…hen using --skip-colmap on images.bin with non-standard image names
  • Loading branch information
jkulhanek committed Feb 7, 2023
1 parent e749fbc commit e05c355
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 8 deletions.
11 changes: 10 additions & 1 deletion nerfstudio/process_data/colmap_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -604,6 +604,14 @@ def colmap_to_json(
cameras = read_cameras_binary(cameras_path)
images = read_images_binary(images_path)

# Images were renamed to frame_{i:05d}.{ext} and
# the filenames needs to be replaced in the transforms.json as well
original_filenames = [x.name for x in images.values()]
# Sort was used in nerfstudio.process_data.process_data_utils:get_image_filenames
original_filenames.sort()
# Build the map to the new filenames
filename_map = {name: f"frame_{i+1:05d}{os.path.splitext(name)[-1]}" for i, name in enumerate(original_filenames)}

# Only supports one camera
camera_params = cameras[1].params

Expand All @@ -619,7 +627,8 @@ def colmap_to_json(
c2w = c2w[np.array([1, 0, 2, 3]), :]
c2w[2, :] *= -1

name = Path(f"./images/{im_data.name}")
name = filename_map[im_data.name]
name = Path(f"./images/{name}")

frame = {
"file_path": name.as_posix(),
Expand Down
16 changes: 14 additions & 2 deletions nerfstudio/process_data/process_data_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,19 @@ def copy_and_upscale_polycam_depth_maps_list(
return copied_depth_map_paths


def list_images(data: Path) -> List[Path]:
"""Lists all supported images in a directory
Args:
data: Path to the directory of images.
Returns:
Paths to images contained in the directory
"""
allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
image_paths = sorted([p for p in data.glob("[!.]*") if p.suffix.lower() in allowed_exts])
return image_paths


def copy_images(data: Path, image_dir: Path, verbose) -> int:
"""Copy images from a directory to a new directory.
Expand All @@ -240,8 +253,7 @@ def copy_images(data: Path, image_dir: Path, verbose) -> int:
The number of images copied.
"""
with status(msg="[bold yellow]Copying images...", spinner="bouncingBall", verbose=verbose):
allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
image_paths = sorted([p for p in data.glob("[!.]*") if p.suffix.lower() in allowed_exts])
image_paths = list_images(data)

if len(image_paths) == 0:
CONSOLE.log("[bold red]:skull: No usable images in the data folder.")
Expand Down
32 changes: 27 additions & 5 deletions scripts/process_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,10 @@ class ProcessImages:
will downscale the images by 2x, 4x, and 8x."""
skip_colmap: bool = False
"""If True, skips COLMAP and generates transforms.json if possible."""
skip_images: bool = False
"""If True, skips copying and downscaling of images and only runs COLMAP if possible and enabled"""
colmap_model_path: str = "colmap/sparse/0"
"""Optionally sets the path of the colmap model. Used only when --skip-colmap is set to True."""
colmap_cmd: str = "colmap"
"""How to call the COLMAP executable."""
gpu: bool = True
Expand All @@ -83,6 +87,9 @@ class ProcessImages:

def main(self) -> None:
"""Process images into a nerfstudio dataset."""
if not self.skip_colmap and self.colmap_model_path != "colmap/sparse/0":
CONSOLE.log("[bold red]The --colmap-model-path can only be used when --skip-colmap is set.")
sys.exit(1)
install_checks.check_ffmpeg_installed()
install_checks.check_colmap_installed()

Expand All @@ -92,12 +99,22 @@ def main(self) -> None:

summary_log = []

# Copy images to output directory
num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose)
summary_log.append(f"Starting with {num_frames} images")
# Copy and downscale images
if not self.skip_images:
# Copy images to output directory
num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose)
summary_log.append(f"Starting with {num_frames} images")

# Downscale images
summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose))
# Downscale images
summary_log.append(
process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)
)
else:
num_frames = len(process_data_utils.list_images(self.data))
if num_frames == 0:
CONSOLE.log("[bold red]:skull: No usable images in the data folder.")
sys.exit(1)
summary_log.append(f"Starting with {num_frames} images")

# Run COLMAP
colmap_dir = self.output_dir / "colmap"
Expand Down Expand Up @@ -133,6 +150,11 @@ def main(self) -> None:
sys.exit(1)

# Save transforms.json
if self.skip_colmap:
colmap_model_path = self.colmap_model_path
else:
colmap_model_path = colmap_dir

if (colmap_dir / "sparse" / "0" / "cameras.bin").exists():
with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"):
num_matched_frames = colmap_utils.colmap_to_json(
Expand Down

0 comments on commit e05c355

Please sign in to comment.