diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py new file mode 100644 index 000000000..596af58fb --- /dev/null +++ b/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.") diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 4919edd93..2d78469eb 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/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, } diff --git a/pyproject.toml b/pyproject.toml index 73bdc21a6..2dae98036 100644 --- a/pyproject.toml +++ b/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", diff --git a/scripts/process_data.py b/scripts/process_data.py index 752e1d15c..7b6bb7922 100644 --- a/scripts/process_data.py +++ b/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(