From 5b50740f45f573dc9a259252370bd2d3c850de8a Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 16 Jan 2023 15:34:20 -0800 Subject: [PATCH] Add cropping options to data processing (#1236) --- nerfstudio/data/datasets/base_dataset.py | 3 + nerfstudio/data/utils/data_utils.py | 2 + nerfstudio/process_data/colmap_utils.py | 17 ++- nerfstudio/process_data/process_data_utils.py | 114 ++++++++++++++++++ scripts/process_data.py | 26 +++- 5 files changed, 159 insertions(+), 3 deletions(-) diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index c01d26739..45aa188ef 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/nerfstudio/data/datasets/base_dataset.py @@ -99,6 +99,9 @@ class InputDataset(Dataset): if self.has_masks: mask_filepath = self._dataparser_outputs.mask_filenames[image_idx] data["mask"] = get_image_mask_tensor_from_path(filepath=mask_filepath, scale_factor=self.scale_factor) + assert ( + data["mask"].shape[:2] == data["image"].shape[:2] + ), f"Mask and image have different shapes. Got {data['mask'].shape[:2]} and {data['image'].shape[:2]}" metadata = self.get_metadata(data) data.update(metadata) return data diff --git a/nerfstudio/data/utils/data_utils.py b/nerfstudio/data/utils/data_utils.py index 22a3438dd..c4b21febd 100644 --- a/nerfstudio/data/utils/data_utils.py +++ b/nerfstudio/data/utils/data_utils.py @@ -32,6 +32,8 @@ def get_image_mask_tensor_from_path(filepath: Path, scale_factor: float = 1.0) - newsize = (int(width * scale_factor), int(height * scale_factor)) pil_mask = pil_mask.resize(newsize, resample=Image.NEAREST) mask_tensor = torch.from_numpy(np.array(pil_mask)).unsqueeze(-1).bool() + if len(mask_tensor.shape) != 3: + raise ValueError("The mask image should have 1 channel") return mask_tensor diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index b8e91bd29..9db20647f 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -492,6 +492,7 @@ def run_colmap( image_dir: Path, colmap_dir: Path, camera_model: CameraModel, + camera_mask_path: Optional[Path] = None, gpu: bool = True, verbose: bool = False, matching_method: Literal["vocab_tree", "exhaustive", "sequential"] = "vocab_tree", @@ -503,8 +504,11 @@ def run_colmap( image_dir: Path to the directory containing the images. colmap_dir: Path to the output directory. camera_model: Camera model to use. + camera_mask_path: Path to the camera mask. gpu: If True, use GPU. verbose: If True, logs the output of the command. + matching_method: Matching method to use. + colmap_cmd: Path to the COLMAP executable. """ colmap_version = get_colmap_version(colmap_cmd) @@ -523,6 +527,8 @@ def run_colmap( f"--ImageReader.camera_model {camera_model.value}", f"--SiftExtraction.use_gpu {int(gpu)}", ] + if camera_mask_path is not None: + feature_extractor_cmd.append(f"--ImageReader.camera_mask_path {camera_mask_path}") feature_extractor_cmd = " ".join(feature_extractor_cmd) with status(msg="[bold yellow]Running COLMAP feature extractor...", spinner="moon", verbose=verbose): run_command(feature_extractor_cmd, verbose=verbose) @@ -575,13 +581,20 @@ def run_colmap( CONSOLE.log("[bold green]:tada: Done refining intrinsics.") -def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path, camera_model: CameraModel) -> int: +def colmap_to_json( + cameras_path: Path, + images_path: Path, + output_dir: Path, + camera_model: CameraModel, + camera_mask_path: Optional[Path] = None, +) -> int: """Converts COLMAP's cameras.bin and images.bin to a JSON file. Args: cameras_path: Path to the cameras.bin file. images_path: Path to the images.bin file. output_dir: Path to the output directory. + camera_mask_path: Path to the camera mask. camera_model: Camera model used. Returns: @@ -612,6 +625,8 @@ def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path, came "file_path": name.as_posix(), "transform_matrix": c2w.tolist(), } + if camera_mask_path is not None: + frame["mask_path"] = camera_mask_path.relative_to(camera_mask_path.parent.parent).as_posix() frames.append(frame) out = { diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index b03dcd8cb..24797658b 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -20,6 +20,8 @@ from enum import Enum from pathlib import Path from typing import List, Optional, Tuple +import cv2 +import numpy as np from rich.console import Console from typing_extensions import Literal @@ -268,3 +270,115 @@ def find_tool_feature_matcher_combination( return (sfm_tool, feature_type, matcher_type) return (None, None, None) + + +def generate_circle_mask(height: int, width: int, percent_radius) -> Optional[np.ndarray]: + """generate a circle mask of the given size. + + Args: + height: The height of the mask. + width: The width of the mask. + percent_radius: The radius of the circle as a percentage of the image diagonal size. + + Returns: + The mask or None if the radius is too large. + """ + if percent_radius <= 0.0: + CONSOLE.log("[bold red]:skull: The radius of the circle mask must be positive.") + sys.exit(1) + if percent_radius >= 1.0: + return None + mask = np.zeros((height, width), dtype=np.uint8) + center = (width // 2, height // 2) + radius = int(percent_radius * np.sqrt(width**2 + height**2) / 2.0) + cv2.circle(mask, center, radius, 1, -1) + return mask + + +def generate_crop_mask( + height: int, width: int, percent_crop: 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]. + + Returns: + The mask or None if no cropping is performed. + """ + if np.all(np.array(percent_crop) == 0.0): + return None + if np.any(np.array(percent_crop) < 0.0) or np.any(np.array(percent_crop) > 1.0): + CONSOLE.log("[bold red]Invalid crop percentage, must be between 0 and 1.") + sys.exit(1) + top, bottom, left, right = percent_crop + mask = np.zeros((height, width), dtype=np.uint8) + top = int(top * height) + bottom = int(bottom * height) + left = int(left * width) + right = int(right * width) + mask[top : height - bottom, left : width - right] = 1.0 + return mask + + +def generate_mask( + height: int, width: int, percent_crop: 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]. + 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) + circle_mask = generate_circle_mask(height, width, percent_radius) + if crop_mask is None: + return circle_mask + if circle_mask is None: + return crop_mask + return crop_mask * circle_mask + + +def save_mask( + image_dir: Path, + num_downscales: int, + percent_crop: 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. + + 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]. + percent_radius: The radius of the circle as a percentage of the image diagonal size. + + Returns: + The path to the mask file or None if no mask is needed. + """ + 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) + if mask is None: + return None + mask *= 255 + mask_path = image_dir.parent / "masks" + mask_path.mkdir(exist_ok=True) + cv2.imwrite(str(mask_path / "mask.png"), mask) + downscale_factors = [2**i for i in range(num_downscales + 1)[1:]] + for downscale in downscale_factors: + mask_path_i = image_dir.parent / f"masks_{downscale}" + mask_path_i.mkdir(exist_ok=True) + mask_path_i = mask_path_i / "mask.png" + mask_i = cv2.resize(mask, (width // downscale, height // downscale), interpolation=cv2.INTER_NEAREST) + cv2.imwrite(str(mask_path_i), mask_i) + CONSOLE.log(":tada: Generated and saved masks.") + return mask_path / "mask.png" diff --git a/scripts/process_data.py b/scripts/process_data.py index fcea3fb44..483735911 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -7,7 +7,7 @@ import sys import zipfile from dataclasses import dataclass from pathlib import Path -from typing import Union +from typing import Tuple, Union import numpy as np import tyro @@ -201,6 +201,10 @@ class ProcessVideo: """If True, skips COLMAP and generates transforms.json if possible.""" colmap_cmd: str = "colmap" """How to call the COLMAP executable.""" + 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)""" gpu: bool = True """If True, use GPU.""" verbose: bool = False @@ -215,12 +219,23 @@ class ProcessVideo: image_dir = self.output_dir / "images" image_dir.mkdir(parents=True, exist_ok=True) + summary_log = [] # Convert video to images summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose ) - # Downscale images + # Create mask + mask_path = process_data_utils.save_mask( + image_dir=image_dir, + num_downscales=self.num_downscales, + percent_crop=self.percent_crop, + percent_radius=self.percent_radius_crop, + ) + if mask_path is not None: + summary_log.append(f"Saved mask to {mask_path}") + + # # Downscale images summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)) # Run Colmap @@ -237,12 +252,18 @@ class ProcessVideo: image_dir=image_dir, colmap_dir=colmap_dir, camera_model=CAMERA_MODELS[self.camera_type], + camera_mask_path=mask_path, gpu=self.gpu, verbose=self.verbose, matching_method=self.matching_method, colmap_cmd=self.colmap_cmd, ) elif sfm_tool == "hloc": + if mask_path is not None: + CONSOLE.log( + "[bold red]Cannot use a mask with hloc. Please remove the cropping options and try again." + ) + sys.exit(1) hloc_utils.run_hloc( image_dir=image_dir, colmap_dir=colmap_dir, @@ -264,6 +285,7 @@ class ProcessVideo: images_path=colmap_dir / "sparse" / "0" / "images.bin", output_dir=self.output_dir, camera_model=CAMERA_MODELS[self.camera_type], + camera_mask_path=mask_path, ) summary_log.append(f"Colmap matched {num_matched_frames} images") summary_log.append(colmap_utils.get_matching_summary(num_extracted_frames, num_matched_frames))