Equirectangular input (#1446)

* started adding equirectangular to process_data

* was multiplying frame by 255 twice

* added equi2planar configs for the 3 presets

* fixed presets

* setup process data

* changed variable name and comment

* fixed middle preset

* removed third preset

* added progress bar and fixed video for equirectangular input

* fixed whitespace stuff

* added pyequilib dependencey

* Updated for JPG support for equirectangular image input

Updated process_data_utils for JPG support for equirectangular image input, since even though .jpg images worked, .JPG images (output from 360 camera software) were not working.

* changed file type matching

* added equirect utils

* added license header

* added docstring on equirect utils

* added heuristic for image size

* added heuristic for videos

* fixed linty stuff

* shortened code

* raised error if no ims in directory

* Updated ProcessVideo for fixing directories for equirectangular input

Updated the ProcessVideo for equirectangular input to make the image directory have the correct perspective projections. The original equirectangular images are copied into a temp_images folder which also holds the perspective projections. This folder is deleted once the perspective projections are copied into the process data images folder.

---------

Co-authored-by: cvachha <cyrus.vachha@yahoo.com>
Co-authored-by: cvachha <cvachha549@gmail.com>

nerfstudio/process_data/equirect_utils.py

@@ -0,0 +1,120 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helper utils for processing equirectangular data."""
+
+import os
+from pathlib import Path
+from typing import Tuple
+
+import cv2
+import numpy as np
+import torch
+from equilib import Equi2Pers
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
+
+from nerfstudio.utils.rich_utils import ItersPerSecColumn
+
+
+def generate_planar_projections_from_equirectangular(
+    image_dir: Path,
+    planar_image_size: Tuple[int, int],
+    samples_per_im: int,
+) -> Path:
+    """Generate planar projections from an equirectangular image.
+
+    Args:
+        image_dir: The directory containing the equirectangular image.
+        planar_image_size: The size of the planar projections [width, height].
+        samples_per_im: The number of samples to take per image.
+    returns:
+        The path to the planar projections directory.
+    """
+
+    device = torch.device("cuda")
+
+    fov = 120
+    yaw_pitch_pairs = []
+    if samples_per_im == 8:
+        fov = 120
+        for i in np.arange(-180, 180, 90):
+            yaw_pitch_pairs.append((i, 0))
+        for i in np.arange(-180, 180, 180):
+            yaw_pitch_pairs.append((i, 45))
+        for i in np.arange(-180, 180, 180):
+            yaw_pitch_pairs.append((i, -45))
+    elif samples_per_im == 14:
+        fov = 110
+        for i in np.arange(-180, 180, 60):
+            yaw_pitch_pairs.append((i, 0))
+        for i in np.arange(-180, 180, 90):
+            yaw_pitch_pairs.append((i, 45))
+        for i in np.arange(-180, 180, 90):
+            yaw_pitch_pairs.append((i, -45))
+
+    equi2pers = Equi2Pers(height=planar_image_size[1], width=planar_image_size[0], fov_x=fov, mode="bilinear")
+    frame_dir = image_dir
+    output_dir = image_dir / "planar_projections"
+    output_dir.mkdir(exist_ok=True)
+    num_ims = len(os.listdir(frame_dir))
+    progress = Progress(
+        TextColumn("[bold blue]Generating Planar Images", justify="right"),
+        BarColumn(),
+        TaskProgressColumn(show_speed=True),
+        ItersPerSecColumn(suffix="equirect frames/s"),
+        TimeRemainingColumn(elapsed_when_finished=True, compact=True),
+    )
+
+    curr_im = 0
+    with progress:
+        for i in progress.track(os.listdir(frame_dir), description="", total=num_ims):
+            if i.lower().endswith((".jpg", ".png", ".jpeg")):
+                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
+                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
+                    pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy()
+                    cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.jpg", pers_image)
+                    count += 1
+            curr_im += 1
+
+    return output_dir
+
+
+def compute_resolution_from_equirect(image_dir: Path, num_images: int) -> Tuple[int, int]:
+    """Compute the resolution of the persepctive projections of equirectangular images
+       from the heuristic: num_image * res**2 = orig_height * orig_width.
+
+    Args:
+        image_dir: The directory containing the equirectangular images.
+    returns:
+        The target resolution of the perspective projections.
+    """
+
+    for i in os.listdir(image_dir):
+        if i.lower().endswith((".jpg", ".png", ".jpeg")):
+            im = np.array(cv2.imread(os.path.join(image_dir, i)))
+            res_squared = (im.shape[0] * im.shape[1]) / num_images
+            return (int(np.sqrt(res_squared)), int(np.sqrt(res_squared)))
+    raise ValueError("No images found in the directory.")

nerfstudio/process_data/process_data_utils.py

@@ -43,6 +43,7 @@ class CameraModel(Enum):
 CAMERA_MODELS = {
     "perspective": CameraModel.OPENCV,
     "fisheye": CameraModel.OPENCV_FISHEYE,
+    "equirectangular": CameraModel.OPENCV,
 }
 
 

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "Pillow>=9.3.0",
     "plotly>=5.7.0",
     "protobuf<=3.20.3,!=3.20.0",
+    "pyequilib>=0.5.6",
     "pymeshlab>=2022.2.post2",
     "pyngrok>=5.1.0",
     "python-socketio>=5.7.1",

scripts/process_data.py

@@ -3,6 +3,7 @@
 
 
 import json
+import shutil
 import sys
 import zipfile
 from dataclasses import dataclass
@@ -16,6 +17,7 @@ from typing_extensions import Annotated, Literal
 
 from nerfstudio.process_data import (
     colmap_utils,
+    equirect_utils,
     hloc_utils,
     insta360_utils,
     metashape_utils,
@@ -45,7 +47,7 @@ class ProcessImages:
     """Path the data, either a video file or a directory of images."""
     output_dir: Path
     """Path to the output directory."""
-    camera_type: Literal["perspective", "fisheye"] = "perspective"
+    camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective"
     """Camera model to use."""
     matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
     """Feature matching method to use. Vocab tree is recommended for a balance of speed and
@@ -83,6 +85,10 @@ class ProcessImages:
     """
     colmap_cmd: str = "colmap"
     """How to call the COLMAP executable."""
+    images_per_equirect: Literal[8, 14] = 8
+    """Number of samples per image to take from each equirectangular image.
+       Used only when camera-type is equirectangular.
+    """
     gpu: bool = True
     """If True, use GPU."""
     verbose: bool = False
@@ -110,6 +116,14 @@ class ProcessImages:
         image_dir = self.output_dir / "images"
         image_dir.mkdir(parents=True, exist_ok=True)
 
+        # Generate planar projections if equirectangular
+        if self.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
+            )
+
         summary_log = []
 
         # Copy and downscale images
@@ -208,7 +222,7 @@ class ProcessVideo:
     """Path to the output directory."""
     num_frames_target: int = 300
     """Target number of frames to use for the dataset, results may not be exact."""
-    camera_type: Literal["perspective", "fisheye"] = "perspective"
+    camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective"
     """Camera model to use."""
     matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
     """Feature matching method to use. Vocab tree is recommended for a balance of speed and
@@ -240,6 +254,10 @@ class ProcessVideo:
     """If True, skips COLMAP and generates transforms.json if possible."""
     colmap_cmd: str = "colmap"
     """How to call the COLMAP executable."""
+    images_per_equirect: Literal[8, 14] = 8
+    """Number of samples per image to take from each equirectangular image.
+       Used only when camera-type is equirectangular.
+    """
     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)
@@ -260,9 +278,37 @@ class ProcessVideo:
 
         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
-        )
+        if self.camera_type == "equirectangular":
+            # create temp images folder to store the equirect and perspective images
+            temp_image_dir = self.output_dir / "temp_images"
+            temp_image_dir.mkdir(parents=True, exist_ok=True)
+            summary_log, num_extracted_frames = process_data_utils.convert_video_to_images(
+                self.data, image_dir=temp_image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose
+            )
+        else:
+            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
+            )
+
+        # Generate planar projections if equirectangular
+        if self.camera_type == "equirectangular":
+            perspective_image_size = equirect_utils.compute_resolution_from_equirect(
+                self.output_dir / "temp_images", self.images_per_equirect
+            )
+            image_dir = equirect_utils.generate_planar_projections_from_equirectangular(
+                self.output_dir / "temp_images", perspective_image_size, self.images_per_equirect
+            )
+
+            # copy the perspective images to the image directory
+            process_data_utils.copy_images(
+                self.output_dir / "temp_images" / "planar_projections",
+                image_dir=self.output_dir / "images",
+                verbose=False,
+            )
+            image_dir = self.output_dir / "images"
+
+            # remove the temp_images folder
+            shutil.rmtree(self.output_dir / "temp_images", ignore_errors=True)
 
         # Create mask
         mask_path = process_data_utils.save_mask(