From b3efcb9bfd08897b86ebf7d6fb07ebdfe9c8b58c Mon Sep 17 00:00:00 2001
From: Kevin Chen <79341521+kevinddchen@users.noreply.github.com>
Date: Tue, 8 Nov 2022 19:35:56 +0000
Subject: [PATCH] Local camera coordinate directions

---
 nerfstudio/cameras/cameras.py   | 23 ++++++++---------------
 nerfstudio/data/datamanagers.py |  4 ++--
 scripts/process_data.py         | 12 ++++++------
 3 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py
index a9a56c6d7..ec9649343 100644
--- a/nerfstudio/cameras/cameras.py
+++ b/nerfstudio/cameras/cameras.py
@@ -46,7 +46,7 @@ CAMERA_MODEL_TO_TYPE = {
     "RADIAL": CameraType.PERSPECTIVE,
     "OPENCV": CameraType.PERSPECTIVE,
     "OPENCV_FISHEYE": CameraType.FISHEYE,
-    "EQUIRECTANGULAR": CameraType.EQUIRECTANGULAR
+    "EQUIRECTANGULAR": CameraType.EQUIRECTANGULAR,
 }
 
 
@@ -315,20 +315,13 @@ class Cameras:
                 dim=-1,
             )
         elif self.camera_type[0] == CameraType.EQUIRECTANGULAR.value:
-            coord = torch.stack([x,y ], -1)
-            coord_x_offset = torch.stack([x+1, y], -1)
-            coord_y_offset = torch.stack([x, y+1], -1)
-
-            coord_stack = torch.stack([coord, coord_x_offset, coord_y_offset], dim=0)
-
-            phi = coord_stack[..., 1] / self.image_height[camera_indices] * torch.pi
-            theta = coord_stack[..., 0] / self.image_width[camera_indices] * 2 * torch.pi
+            # u goes from -w/2 to w/2
+            # v goes from -h/2 to h/2
+            phi = torch.pi * (0.5 - coord_stack[..., 1] / self.image_height[camera_indices])
+            theta = -2 * torch.pi * coord_stack[..., 0] / self.image_width[camera_indices]
+            # use spherical in local camera coordinates (+y up, x=0 and z<0 is theta=0)
             directions_stack = torch.stack(
-                [
-                    torch.cos(theta) * torch.sin(phi), 
-                    torch.sin(theta) * torch.sin(phi), 
-                    torch.cos(phi)
-                ],
+                [-torch.sin(theta) * torch.sin(phi), torch.cos(phi), -torch.cos(theta) * torch.sin(phi)],
                 dim=-1,
             )
         else:
@@ -354,7 +347,7 @@ class Cameras:
         if not isinstance(camera_indices, torch.Tensor):
             ray_bundle_camera_indices = torch.Tensor([camera_indices]).broadcast_to(pixel_area.shape).to(self.device)
         else:
-            ray_bundle_camera_indices =  camera_indices.view(pixel_area.shape)
+            ray_bundle_camera_indices = camera_indices.view(pixel_area.shape)
 
         return RayBundle(
             origins=origins,
diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py
index b98d5d48c..c8b5a40da 100644
--- a/nerfstudio/data/datamanagers.py
+++ b/nerfstudio/data/datamanagers.py
@@ -38,7 +38,7 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
 )
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
 from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig
-from nerfstudio.data.pixel_samplers import PixelSampler, EquirectangularSampler
+from nerfstudio.data.pixel_samplers import EquirectangularSampler, PixelSampler
 from nerfstudio.data.utils.dataloaders import (
     CacheDataloader,
     FixedIndicesEvalDataloader,
@@ -302,7 +302,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             pin_memory=True,
         )
         self.iter_train_image_dataloader = iter(self.train_image_dataloader)
-        self.train_pixel_sampler = EquirectangularSampler(self.config.train_num_rays_per_batch)
+        self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch)
         self.train_camera_optimizer = self.config.camera_optimizer.setup(
             num_cameras=self.train_dataset.dataparser_outputs.cameras.size, device=self.device
         )
diff --git a/scripts/process_data.py b/scripts/process_data.py
index 6a11224ea..3d26dc471 100755
--- a/scripts/process_data.py
+++ b/scripts/process_data.py
@@ -605,12 +605,12 @@ def opensfm_to_json(reconstruction_path: Path, output_dir: Path, camera_model: C
     # For now just assume it's all the same camera
     cam = cams[list(cams.keys())[0]]
     out = {
-        "fl_x": float(cam["focal_x"]) * max(cam["height"], cam["width"]) if "focal_x" in cam else 1.0,
-        "fl_y": float(cam["focal_y"]) * max(cam["height"], cam["width"]) if "focal_y" in cam else 1.0,
-        "cx": cam["width"] / 2.0 + float(cam["c_x"]) * cam["width"] if "c_x" in cam else 1.0,
-        "cy": cam["height"] / 2.0 + float(cam["c_y"]) * cam["height"] if "c_y" in cam else 1.0,
-        "w": cam["width"],
-        "h": cam["height"],
+        "fl_x": float(cam["focal_x"]) if "focal_x" in cam else 1.0,
+        "fl_y": float(cam["focal_y"]) if "focal_y" in cam else 1.0,
+        "cx": float(cam["c_x"]) if "c_x" in cam else 0.5 * cam["width"],
+        "cy": float(cam["c_y"]) if "c_y" in cam else 0.5 * cam["height"],
+        "w": int(cam["width"]),
+        "h": int(cam["height"]),
         "camera_model": camera_model.value,
     }
     if camera_model == CameraModel.OPENCV: