Add cropping options to data processing (#1236)

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 = {

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"