tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Add cropping options to data processing (#1236)

Browse Source
This commit is contained in:
Matthew Tancik 2023-01-16 15:34:20 -08:00 committed by GitHub
parent bd1060ee83
commit 5b50740f45
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 159 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
nerfstudio/data/datasets/base_dataset.py
View File

@ -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

2
nerfstudio/data/utils/data_utils.py
View File

@ -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

17
nerfstudio/process_data/colmap_utils.py
View File

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

114
nerfstudio/process_data/process_data_utils.py
View File

@ -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"

26
scripts/process_data.py
View File

@ -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))