From 19d996ef5e7faffc289310cf51ed199b871f83aa Mon Sep 17 00:00:00 2001
From: Kamyar Salahi <kam.salahi@berkeley.edu>
Date: Fri, 4 Nov 2022 10:13:48 -0700
Subject: [PATCH] fixed bugs in the data processing pipeline

---
 nerfstudio/cameras/cameras.py     | 14 ++++++++++----
 nerfstudio/data/pixel_samplers.py |  3 +--
 scripts/process_data.py           | 10 +++++-----
 3 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py
index 4cfb89c1d..a9a56c6d7 100644
--- a/nerfstudio/cameras/cameras.py
+++ b/nerfstudio/cameras/cameras.py
@@ -315,9 +315,15 @@ class Cameras:
                 dim=-1,
             )
         elif self.camera_type[0] == CameraType.EQUIRECTANGULAR.value:
-            phi = y / self.image_height * torch.pi
-            theta = -x / self.image_width * 2 * torch.pi
-            directions = torch.stack(
+            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
+            directions_stack = torch.stack(
                 [
                     torch.cos(theta) * torch.sin(phi), 
                     torch.sin(theta) * torch.sin(phi), 
@@ -348,7 +354,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/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py
index b2e3122d8..0ec2ccdf9 100644
--- a/nerfstudio/data/pixel_samplers.py
+++ b/nerfstudio/data/pixel_samplers.py
@@ -136,9 +136,8 @@ def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch:
     #     ).long()
 
     number_of_images = torch.rand(num_rays_per_batch, device=device)
-    phi = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / 2 * torch.pi
+    phi = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / torch.pi
     theta = torch.rand(num_rays_per_batch, device=device)
-
     indices = torch.floor(
             torch.stack((number_of_images, phi, theta), dim=-1)
             * torch.tensor([num_images, image_height, image_width], device=device)
diff --git a/scripts/process_data.py b/scripts/process_data.py
index 82774abf5..288cba9dd 100755
--- a/scripts/process_data.py
+++ b/scripts/process_data.py
@@ -602,10 +602,10 @@ 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"]),
-        "fl_y": float(cam["focal_y"]) * max(cam["height"], cam["width"]),
-        "cx": cam["width"] / 2.0 + float(cam["c_x"]) * cam["width"],
-        "cy": cam["height"] / 2.0 + float(cam["c_y"]) * cam["height"],
+        "fl_x": float(cam["focal_x"]) * max(cam["height"], cam["width"]) if "focal_x" in cam else 1.,
+        "fl_y": float(cam["focal_y"]) * max(cam["height"], cam["width"]) if "focal_y" in cam else 1.,
+        "cx": cam["width"] / 2.0 + float(cam["c_x"]) * cam["width"] if "c_x" in cam else 1. ,
+        "cy": cam["height"] / 2.0 + float(cam["c_y"]) * cam["height"] if "c_y" in cam else 1. ,
         "w": cam["width"],
         "h": cam["height"],
         "camera_model": camera_model.value,
@@ -946,7 +946,7 @@ class ProcessInsta360:
                 gpu=self.gpu,
                 verbose=self.verbose,
                 matching_method=self.matching_method,
-                opensfm_install=
+                opensfm_install=Path("/scratch2/ksalahi/OpenSfM")
             )
 
         # Save transforms.json