Implement Fisheye-GS.

examples/benchmarks/fisheye/mcmc_zipnerf.sh

@@ -0,0 +1,25 @@
+SCENE_DIR="data/zipnerf"
+SCENE_LIST="berlin london nyc alameda"
+DATA_FACTOR=4
+RENDER_TRAJ_PATH="ellipse"
+
+RESULT_DIR="results/benchmark_mcmc_2M_zipnerf"
+CAP_MAX=2000000
+
+# RESULT_DIR="results/benchmark_mcmc_4M_zipnerf"
+# CAP_MAX=4000000
+
+for SCENE in $SCENE_LIST;
+do
+    echo "Running $SCENE"
+
+    # train and eval
+    CUDA_VISIBLE_DEVICES=0 python simple_trainer.py mcmc --disable_viewer --data_factor $DATA_FACTOR \
+        --strategy.cap-max $CAP_MAX \
+        --opacity_reg 0.001 \
+        --use_bilateral_grid \
+        --render_traj_path $RENDER_TRAJ_PATH \
+        --camera_model fisheye \
+        --data_dir $SCENE_DIR/$SCENE/ \
+        --result_dir $RESULT_DIR/$SCENE/
+done

examples/benchmarks/fisheye/mcmc_zipnerf_undistorted.sh

@@ -0,0 +1,25 @@
+SCENE_DIR="data/zipnerf_undistorted"
+SCENE_LIST="berlin london nyc alameda"
+DATA_FACTOR=4
+RENDER_TRAJ_PATH="ellipse"
+
+RESULT_DIR="results/benchmark_mcmc_2M_zipnerf_undistorted"
+CAP_MAX=2000000
+
+# RESULT_DIR="results/benchmark_mcmc_4M_zipnerf_undistorted"
+# CAP_MAX=4000000
+
+for SCENE in $SCENE_LIST;
+do
+    echo "Running $SCENE"
+
+    # train and eval
+    CUDA_VISIBLE_DEVICES=0 python simple_trainer.py mcmc --disable_viewer --data_factor $DATA_FACTOR \
+        --strategy.cap-max $CAP_MAX \
+        --opacity_reg 0.001 \
+        --use_bilateral_grid \
+        --render_traj_path $RENDER_TRAJ_PATH \
+        --camera_model pinhole \
+        --data_dir $SCENE_DIR/$SCENE/ \
+        --result_dir $RESULT_DIR/$SCENE/
+done

examples/benchmarks/mcmc.sh

@@ -19,7 +19,7 @@ do
     CUDA_VISIBLE_DEVICES=0 python simple_trainer.py mcmc --eval_steps -1 --disable_viewer --data_factor $DATA_FACTOR \
         --strategy.cap-max $CAP_MAX \
         --render_traj_path $RENDER_TRAJ_PATH \
-        --data_dir data/360_v2/$SCENE/ \
+        --data_dir $SCENE_DIR/$SCENE/ \
         --result_dir $RESULT_DIR/$SCENE/
 
     # run eval and render

examples/datasets/colmap.py

@@ -1,6 +1,7 @@
 import os
 import json
 from typing import Any, Dict, List, Optional
+from typing_extensions import assert_never
 
 import cv2
 import imageio.v2 as imageio
@@ -59,6 +60,7 @@ def __init__(
         Ks_dict = dict()
         params_dict = dict()
         imsize_dict = dict()  # width, height
+        mask_dict = dict()
         bottom = np.array([0, 0, 0, 1]).reshape(1, 4)
         for k in imdata:
             im = imdata[k]
@@ -99,14 +101,12 @@ def __init__(
                 params = np.array([cam.k1, cam.k2, cam.k3, cam.k4], dtype=np.float32)
                 camtype = "fisheye"
             assert (
-                camtype == "perspective"
-            ), f"Only support perspective camera model, got {type_}"
+                camtype == "perspective" or camtype == "fisheye"
+            ), f"Only perspective and fisheye cameras are supported, got {type_}"
 
             params_dict[camera_id] = params
-
-            # image size
             imsize_dict[camera_id] = (cam.width // factor, cam.height // factor)
-
+            mask_dict[camera_id] = None
         print(
             f"[Parser] {len(imdata)} images, taken by {len(set(camera_ids))} cameras."
         )
@@ -203,6 +203,7 @@ def __init__(
         self.Ks_dict = Ks_dict  # Dict of camera_id -> K
         self.params_dict = params_dict  # Dict of camera_id -> params
         self.imsize_dict = imsize_dict  # Dict of camera_id -> (width, height)
+        self.mask_dict = mask_dict  # Dict of camera_id -> mask
         self.points = points  # np.ndarray, (num_points, 3)
         self.points_err = points_err  # np.ndarray, (num_points,)
         self.points_rgb = points_rgb  # np.ndarray, (num_points, 3)
@@ -236,16 +237,60 @@ def __init__(
             ), f"Missing params for camera {camera_id}"
             K = self.Ks_dict[camera_id]
             width, height = self.imsize_dict[camera_id]
-            K_undist, roi_undist = cv2.getOptimalNewCameraMatrix(
-                K, params, (width, height), 0
-            )
-            mapx, mapy = cv2.initUndistortRectifyMap(
-                K, params, None, K_undist, (width, height), cv2.CV_32FC1
-            )
-            self.Ks_dict[camera_id] = K_undist
+
+            if camtype == "perspective":
+                K_undist, roi_undist = cv2.getOptimalNewCameraMatrix(
+                    K, params, (width, height), 0
+                )
+                mapx, mapy = cv2.initUndistortRectifyMap(
+                    K, params, None, K_undist, (width, height), cv2.CV_32FC1
+                )
+                mask = None
+            elif camtype == "fisheye":
+                fx = K[0, 0]
+                fy = K[1, 1]
+                cx = K[0, 2]
+                cy = K[1, 2]
+                grid_x, grid_y = np.meshgrid(
+                    np.arange(width, dtype=np.float32),
+                    np.arange(height, dtype=np.float32),
+                    indexing="xy",
+                )
+                x1 = (grid_x - cx) / fx
+                y1 = (grid_y - cy) / fy
+                theta = np.sqrt(x1**2 + y1**2)
+                r = (
+                    1.0
+                    + params[0] * theta**2
+                    + params[1] * theta**4
+                    + params[2] * theta**6
+                    + params[3] * theta**8
+                )
+                mapx = fx * x1 * r + width // 2
+                mapy = fy * y1 * r + height // 2
+
+                # Use mask to define ROI
+                mask = np.logical_and(
+                    np.logical_and(mapx > 0, mapy > 0),
+                    np.logical_and(mapx < width - 1, mapy < height - 1),
+                )
+                y_indices, x_indices = np.nonzero(mask)
+                y_min, y_max = y_indices.min(), y_indices.max() + 1
+                x_min, x_max = x_indices.min(), x_indices.max() + 1
+                mask = mask[y_min:y_max, x_min:x_max]
+                K_undist = K.copy()
+                K_undist[0, 2] -= x_min
+                K_undist[1, 2] -= y_min
+                roi_undist = [x_min, y_min, x_max - x_min, y_max - y_min]
+            else:
+                assert_never(camtype)
+
             self.mapx_dict[camera_id] = mapx
             self.mapy_dict[camera_id] = mapy
+            self.Ks_dict[camera_id] = K_undist
             self.roi_undist_dict[camera_id] = roi_undist
+            self.imsize_dict[camera_id] = (roi_undist[2], roi_undist[3])
+            self.mask_dict[camera_id] = mask
 
         # size of the scene measured by cameras
         camera_locations = camtoworlds[:, :3, 3]
@@ -284,6 +329,7 @@ def __getitem__(self, item: int) -> Dict[str, Any]:
         K = self.parser.Ks_dict[camera_id].copy()  # undistorted K
         params = self.parser.params_dict[camera_id]
         camtoworlds = self.parser.camtoworlds[index]
+        mask = self.parser.mask_dict[camera_id]
 
         if len(params) > 0:
             # Images are distorted. Undistort them.
@@ -310,6 +356,8 @@ def __getitem__(self, item: int) -> Dict[str, Any]:
             "image": torch.from_numpy(image).float(),
             "image_id": item,  # the index of the image in the dataset
         }
+        if mask is not None:
+            data["mask"] = torch.from_numpy(mask).bool()
 
         if self.load_depths:
             # projected points to image plane to get depths

examples/datasets/download_dataset.py

@@ -13,20 +13,36 @@
     "mipnerf360",
     "mipnerf360_extra",
     "bilarf_data",
+    "zipnerf",
+    "zipnerf_undistorted",
 ]
 
 # dataset urls
 urls = {
     "mipnerf360": "http://storage.googleapis.com/gresearch/refraw360/360_v2.zip",
     "mipnerf360_extra": "https://storage.googleapis.com/gresearch/refraw360/360_extra_scenes.zip",
     "bilarf_data": "https://huggingface.co/datasets/Yuehao/bilarf_data/resolve/main/bilarf_data.zip",
+    "zipnerf": [
+        "https://storage.googleapis.com/gresearch/refraw360/zipnerf/berlin.zip",
+        "https://storage.googleapis.com/gresearch/refraw360/zipnerf/london.zip",
+        "https://storage.googleapis.com/gresearch/refraw360/zipnerf/nyc.zip",
+        "https://storage.googleapis.com/gresearch/refraw360/zipnerf/alameda.zip",
+    ],
+    "zipnerf_undistorted": [
+        "https://storage.googleapis.com/gresearch/refraw360/zipnerf-undistorted/berlin.zip",
+        "https://storage.googleapis.com/gresearch/refraw360/zipnerf-undistorted/london.zip",
+        "https://storage.googleapis.com/gresearch/refraw360/zipnerf-undistorted/nyc.zip",
+        "https://storage.googleapis.com/gresearch/refraw360/zipnerf-undistorted/alameda.zip",
+    ],
 }
 
 # rename maps
 dataset_rename_map = {
     "mipnerf360": "360_v2",
     "mipnerf360_extra": "360_v2",
     "bilarf_data": "bilarf",
+    "zipnerf": "zipnerf",
+    "zipnerf_undistorted": "zipnerf_undistorted",
 }
 
 
@@ -40,57 +56,72 @@ def main(self):
         self.dataset_download(self.dataset)
 
     def dataset_download(self, dataset: dataset_names):
-        (self.save_dir / dataset_rename_map[dataset]).mkdir(parents=True, exist_ok=True)
+        if isinstance(urls[dataset], list):
+            for url in urls[dataset]:
+                url_file_name = Path(url).name
+                extract_path = self.save_dir / dataset_rename_map[dataset]
+                download_path = extract_path / url_file_name
+                download_and_extract(url, download_path, extract_path)
+        else:
+            url = urls[dataset]
+            url_file_name = Path(url).name
+            extract_path = self.save_dir / dataset_rename_map[dataset]
+            download_path = extract_path / url_file_name
+            download_and_extract(url, download_path, extract_path)
 
-        file_name = Path(urls[dataset]).name
 
-        # download
-        download_command = [
-            "curl",
-            "-o",
-            str(self.save_dir / dataset_rename_map[dataset] / file_name),
-            urls[dataset],
-        ]
-        try:
-            subprocess.run(download_command, check=True)
-            print("File file downloaded succesfully.")
-        except subprocess.CalledProcessError as e:
-            print(f"Error downloading file: {e}")
-
-        # if .zip
-        if Path(urls[dataset]).suffix == ".zip":
-            if os.name == "nt":  # Windows doesn't have 'unzip' but 'tar' works
-                extract_command = [
-                    "tar",
-                    "-xvf",
-                    self.save_dir / dataset_rename_map[dataset] / file_name,
-                    "-C",
-                    self.save_dir / dataset_rename_map[dataset],
-                ]
-            else:
-                extract_command = [
-                    "unzip",
-                    self.save_dir / dataset_rename_map[dataset] / file_name,
-                    "-d",
-                    self.save_dir / dataset_rename_map[dataset],
-                ]
-        # if .tar
-        else:
+def download_and_extract(url: str, download_path: Path, extract_path: Path) -> None:
+    download_path.parent.mkdir(parents=True, exist_ok=True)
+    extract_path.mkdir(parents=True, exist_ok=True)
+
+    # download
+    download_command = [
+        "curl",
+        "-L",
+        "-o",
+        str(download_path),
+        url,
+    ]
+    try:
+        subprocess.run(download_command, check=True)
+        print("File file downloaded succesfully.")
+    except subprocess.CalledProcessError as e:
+        print(f"Error downloading file: {e}")
+
+    # if .zip
+    if Path(url).suffix == ".zip":
+        if os.name == "nt":  # Windows doesn't have 'unzip' but 'tar' works
             extract_command = [
                 "tar",
-                "-xvzf",
-                self.save_dir / dataset_rename_map[dataset] / file_name,
+                "-xvf",
+                download_path,
                 "-C",
-                self.save_dir / dataset_rename_map[dataset],
+                extract_path,
             ]
+        else:
+            extract_command = [
+                "unzip",
+                download_path,
+                "-d",
+                extract_path,
+            ]
+    # if .tar
+    else:
+        extract_command = [
+            "tar",
+            "-xvzf",
+            download_path,
+            "-C",
+            extract_path,
+        ]
 
-        # extract
-        try:
-            subprocess.run(extract_command, check=True)
-            os.remove(self.save_dir / dataset_rename_map[dataset] / file_name)
-            print("Extraction complete.")
-        except subprocess.CalledProcessError as e:
-            print(f"Extraction failed: {e}")
+    # extract
+    try:
+        subprocess.run(extract_command, check=True)
+        os.remove(download_path)
+        print("Extraction complete.")
+    except subprocess.CalledProcessError as e:
+        print(f"Extraction failed: {e}")
 
 
 if __name__ == "__main__":