Add mask to the euirectangular processing

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:

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)

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]:

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 <https://colmap.github.io/>`_.
     """
 
-    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,
         )