diff --git a/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py index 1c7f8cfbd..a9bc5aadd 100644 --- a/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/base_converter_to_nerfstudio_dataset.py @@ -19,25 +19,24 @@ 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""" + mask_dir: Path = None + """Path to the mask directory.""" verbose: bool = False """If True, print extra logging.""" def __post_init__(self) -> None: self.output_dir.mkdir(parents=True, exist_ok=True) - self.image_dir.mkdir(parents=True, exist_ok=True) + # self.image_dir.mkdir(parents=True, exist_ok=True) @property def image_dir(self) -> Path: diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index ac15d0937..a49a1e2dd 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -138,7 +138,7 @@ def run_colmap( ] if matching_method == "vocab_tree": vocab_tree_filename = get_vocab_tree() - feature_matcher_cmd.append(f'--VocabTreeMatching.vocab_tree_path "{vocab_tree_filename}"') + feature_matcher_cmd.append(f"--VocabTreeMatching.vocab_tree_path {vocab_tree_filename}") feature_matcher_cmd = " ".join(feature_matcher_cmd) with status(msg="[bold yellow]Running COLMAP feature matcher...", spinner="runner", verbose=verbose): run_command(feature_matcher_cmd, verbose=verbose) @@ -438,7 +438,7 @@ def colmap_to_json( "colmap_im_id": im_id, } if camera_mask_path is not None: - frame["mask_path"] = camera_mask_path.relative_to(camera_mask_path.parent.parent).as_posix() + frame["mask_path"] = Path(camera_mask_path.absolute() / im_data.name).as_posix() # camera_mask_path.relative_to(camera_mask_path.parent.parent).as_posix() if image_id_to_depth_path is not None: depth_path = image_id_to_depth_path[im_id] frame["depth_file_path"] = str(depth_path) diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index a04f66363..9d81904ae 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -110,6 +110,7 @@ def generate_planar_projections_from_equirectangular( image_dir: Path, planar_image_size: Tuple[int, int], samples_per_im: int, + mask_dir: Path, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), ) -> Path: """Generate planar projections from an equirectangular image. @@ -170,6 +171,10 @@ def generate_planar_projections_from_equirectangular( frame_dir = image_dir output_dir = image_dir / "planar_projections" output_dir.mkdir(exist_ok=True) + + output_mask_dir = image_dir / "mask_planar_projections" + output_mask_dir.mkdir(exist_ok=True) + num_ims = len(os.listdir(frame_dir)) progress = Progress( TextColumn("[bold blue]Generating Planar Images", justify="right"), @@ -181,21 +186,34 @@ def generate_planar_projections_from_equirectangular( with progress: for i in progress.track(os.listdir(frame_dir), description="", total=num_ims): - if i.lower().endswith((".jpg", ".png", ".jpeg")): + if i.lower().endswith((".jpg", ".png", ".jpeg", ".exr")): im = np.array(cv2.imread(os.path.join(frame_dir, i))) im = torch.tensor(im, dtype=torch.float32, device=device) im = torch.permute(im, (2, 0, 1)) / 255.0 + + mask_addr = os.path.join(mask_dir, i) + if not os.path.isfile(mask_addr): + raise Exception("The corresponding mask for " + i + " does not exist!!!") + mask = np.array(cv2.imread(mask_addr)) + mask = torch.tensor(mask, dtype=torch.float32, device=device) + mask = torch.permute(mask, (2, 0, 1)) / 255.0 + count = 0 for u_deg, v_deg in yaw_pitch_pairs: v_rad = torch.pi * v_deg / 180.0 u_rad = torch.pi * u_deg / 180.0 pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) * 255.0 - assert isinstance(pers_image, torch.Tensor) + pers_mask = equi2pers(mask, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) * 255.0 + + assert isinstance(pers_image, torch.Tensor) and isinstance(pers_mask, torch.Tensor) pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy() + pers_mask = (pers_mask.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy() cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.jpg", pers_image) + cv2.imwrite(f"{output_mask_dir}/{i[:-4]}_{count}.jpg", pers_mask) + count += 1 - return output_dir + return output_dir, output_mask_dir def compute_resolution_from_equirect(image_dir: Path, num_images: int) -> Tuple[int, int]: diff --git a/nerfstudio/process_data/images_to_nerfstudio_dataset.py b/nerfstudio/process_data/images_to_nerfstudio_dataset.py index 7ae7d5bd4..3cce2f092 100644 --- a/nerfstudio/process_data/images_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/images_to_nerfstudio_dataset.py @@ -18,7 +18,9 @@ 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 @@ -30,9 +32,6 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): 2. Calculates the camera poses for each image using `COLMAP `_. """ - percent_radius_crop: float = 1.0 - """Create circle crop mask. The radius is the percent of the image diagonal.""" - def main(self) -> None: """Process images into a nerfstudio dataset.""" @@ -48,15 +47,11 @@ 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.data, perspective_masks = equirect_utils.generate_planar_projections_from_equirectangular( + self.data, pers_size, self.images_per_equirect, self.mask_dir, crop_factor=self.crop_factor ) - self.camera_type = "perspective" summary_log = [] @@ -65,41 +60,25 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): if not self.skip_image_processing: # Copy images to output directory image_rename_map_paths = process_data_utils.copy_images( - 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, + self.data, image_dir=self.image_dir, crop_factor=self.crop_factor, verbose=self.verbose + ) + + (self.output_dir / "masks").mkdir(parents=True, exist_ok=True) + mask_rename_map_paths = process_data_utils.copy_images( + perspective_masks, image_dir=self.output_dir / "masks", crop_factor=self.crop_factor, verbose=self.verbose ) - 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") - # # Create mask - mask_path = process_data_utils.save_mask( - image_dir=self.image_dir, - num_downscales=self.num_downscales, - crop_factor=(0.0, 0.0, 0.0, 0.0), - percent_radius=self.percent_radius_crop, + # Downscale images + summary_log.append( + process_data_utils.downscale_images(self.image_dir, self.num_downscales, verbose=self.verbose) ) - if mask_path is not None: - summary_log.append("Saved mask(s)") + # # Downscale images + # summary_log.append( + # process_data_utils.downscale_images((self.output_dir / "masks"), self.num_downscales, verbose=self.verbose) + # ) else: num_frames = len(process_data_utils.list_images(self.data)) if num_frames == 0: @@ -109,7 +88,7 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): # Run COLMAP if not self.skip_colmap: require_cameras_exist = True - self._run_colmap() + self._run_colmap(mask_path=(self.output_dir / "masks")) # Colmap uses renamed images image_rename_map = None @@ -123,7 +102,7 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): summary_log += self._save_transforms( num_frames, image_id_to_depth_path, - None, + (self.output_dir / "masks"), image_rename_map, )