merge ba_helper into one

Author: zitongzhan

ba_example.py

@@ -1,80 +1,110 @@
 from time import perf_counter
-import torch
+
 import pypose as pp
+import torch
+import torch.nn as nn
 
-from ba_helpers import Reproj, least_square_error
-from datapipes.bal_loader import get_problem, read_bal_data
-from bae.sparse.py_ops import *
+from datapipes.bal_loader import get_problem
+from bae.autograd.function import TrackingTensor, map_transform
 from bae.optim import LM
-from bae.utils.pysolvers import PCG, CuDSS
+from bae.utils.ba import rotate_quat
+from bae.utils.pysolvers import PCG
 
-# TARGET_DATASET = "ladybug"
-# TARGET_PROBLEM = "problem-1723-156502-pre"
-# TARGET_PROBLEM = "problem-49-7776-pre"
-# TARGET_PROBLEM = "problem-1695-155710-pre"  
-# TARGET_PROBLEM = "problem-969-105826-pre"
 TARGET_DATASET = "trafalgar"
 TARGET_PROBLEM = "problem-257-65132-pre"
+# other options:
+# TARGET_DATASET = "ladybug"
+# TARGET_PROBLEM = "problem-1723-156502-pre"
 # TARGET_DATASET = "dubrovnik"
 # TARGET_PROBLEM = "problem-356-226730-pre"
 
-
-
-DEVICE = 'cuda'
+DEVICE = "cuda"
 OPTIMIZE_INTRINSICS = True
-
-USE_QUATERNIONS = True
-
-file_name = f'{TARGET_DATASET}.{TARGET_PROBLEM}'
-dataset = get_problem(TARGET_PROBLEM, TARGET_DATASET, use_quat=USE_QUATERNIONS)
-
-if OPTIMIZE_INTRINSICS:
-    NUM_CAMERA_PARAMS = 10 if USE_QUATERNIONS else 9
-else:
-    NUM_CAMERA_PARAMS = 7 if USE_QUATERNIONS else 6
-
-print(f'Fetched {TARGET_PROBLEM} from {TARGET_DATASET}')
-
-trimmed_dataset = dataset
-trimmed_dataset = {k: v.to(DEVICE) for k, v in trimmed_dataset.items() if type(v) == torch.Tensor}
-
-input = {
-    "points_2d": trimmed_dataset['points_2d'],
-    "camera_indices": trimmed_dataset['camera_index_of_observations'],
-    "point_indices": trimmed_dataset['point_index_of_observations']
-}
-
-model = Reproj(
-    trimmed_dataset['camera_params'][:, :NUM_CAMERA_PARAMS].clone(),
-    trimmed_dataset['points_3d'].clone()
-).to(DEVICE)
-strategy = pp.optim.strategy.TrustRegion(up=2.0, down=0.5**4)
-solver = PCG(tol=1e-4, maxiter=250)  # or CuDSS()
-optimizer = LM(model, strategy=strategy, solver=solver, reject=30)
-
-print('Loss:', least_square_error(
-    model.pose,
-    model.points_3d,
-    trimmed_dataset['camera_index_of_observations'],
-    trimmed_dataset['point_index_of_observations'],
-    trimmed_dataset['points_2d'],
-).item())
-
-print("Initial loss", optimizer.model.loss(input, None).item())
-
-start = perf_counter()
-for idx in range(20):
-    loss = optimizer.step(input)
-    print('Iteration', idx, 'loss', loss.item(), 'time', perf_counter() - start)
-
-torch.cuda.synchronize()
-end = perf_counter()
-print('Time', end - start)
-
-print('Ending loss:', least_square_error(
-    model.pose,
-    model.points_3d,
-    trimmed_dataset['camera_index_of_observations'],
-    trimmed_dataset['point_index_of_observations'],
-    trimmed_dataset['points_2d'],
-).item())
+NUM_CAMERA_PARAMS = 10 if OPTIMIZE_INTRINSICS else 7
+
+
+@map_transform
+def project(points, camera_params):
+    projection = rotate_quat(points, camera_params[..., :7])
+    projection = -projection[..., :2] / projection[..., [2]]
+
+    f = camera_params[..., [-3]]
+    k1 = camera_params[..., [-2]]
+    k2 = camera_params[..., [-1]]
+
+    n = torch.sum(projection**2, axis=-1, keepdim=True)
+    r = 1 + k1 * n + k2 * n**2
+    return projection * r * f
+
+
+class Residual(nn.Module):
+    def __init__(self, camera_params, points):
+        super().__init__()
+        self.pose = nn.Parameter(TrackingTensor(camera_params))
+        self.points = nn.Parameter(TrackingTensor(points))
+        self.pose.trim_SE3_grad = True
+
+    def forward(self, observes, cidx, pidx):
+        points_proj = project(self.points[pidx], self.pose[cidx])
+        return points_proj - observes
+
+
+def least_square_error(camera_params, points, cidx, pidx, observes):
+    model = Residual(camera_params, points)
+    loss = model(observes, cidx, pidx)
+    return torch.sum(loss**2, dim=-1).mean()
+
+
+def main():
+    dataset = get_problem(TARGET_PROBLEM, TARGET_DATASET)
+    print(f"Fetched {TARGET_PROBLEM} from {TARGET_DATASET}")
+
+    dataset = {
+        key: value.to(DEVICE)
+        for key, value in dataset.items()
+        if isinstance(value, torch.Tensor)
+    }
+    input = {
+        "observes": dataset["points_2d"],
+        "cidx": dataset["camera_index_of_observations"],
+        "pidx": dataset["point_index_of_observations"],
+    }
+
+    model = Residual(
+        dataset["camera_params"][:, :NUM_CAMERA_PARAMS].clone(),
+        dataset["points_3d"].clone(),
+    ).to(DEVICE)
+    strategy = pp.optim.strategy.TrustRegion(up=2.0, down=0.5**4)
+    solver = PCG(tol=1e-4, maxiter=250)
+    optimizer = LM(model, strategy=strategy, solver=solver, reject=30)
+
+    print('Loss:', least_square_error(
+        model.pose,
+        model.points,
+        dataset["camera_index_of_observations"],
+        dataset["point_index_of_observations"],
+        dataset["points_2d"],
+    ).item())
+
+    print("Initial loss", optimizer.model.loss(input, None).item())
+
+    start = perf_counter()
+    for idx in range(20):
+        loss = optimizer.step(input)
+        print("Iteration", idx, "loss", loss.item(), "time", perf_counter() - start)
+
+    torch.cuda.synchronize()
+    end = perf_counter()
+    print("Time", end - start)
+
+    print('Ending loss:', least_square_error(
+        model.pose,
+        model.points,
+        dataset["camera_index_of_observations"],
+        dataset["point_index_of_observations"],
+        dataset["points_2d"],
+    ).item())
+
+
+if __name__ == "__main__":
+    main()

ba_helpers.py

@@ -17,7 +17,7 @@ def _rotvec_to_quat_xyzw(rotvec: torch.Tensor) -> torch.Tensor:
     return torch.cat([xyz, cos_half], dim=-1)
 
 
-def read_bal_data(file_name: str, use_quat: bool = False) -> dict:
+def read_bal_data(file_name: str, use_quat: bool = True) -> dict:
     """
     Read a Bundle Adjustment in the Large dataset problem text file.
 
@@ -29,8 +29,8 @@ def read_bal_data(file_name: str, use_quat: bool = False) -> dict:
 
     Each camera has 9 parameters: Rodrigues rotvec (3), translation (3), f, k1, k2.
     This loader outputs either:
-      - use_quat=False: [tx, ty, tz, rx, ry, rz, f, k1, k2]  (9)
       - use_quat=True:  [tx, ty, tz, qx, qy, qz, qw, f, k1, k2] (10)
+      - use_quat=False: [tx, ty, tz, rx, ry, rz, f, k1, k2]  (9)
     """
     with open(file_name, "r") as file:
         n_cameras, n_points, n_observations = map(int, file.readline().split())
@@ -70,4 +70,3 @@ def read_bal_data(file_name: str, use_quat: bool = False) -> dict:
         "camera_index_of_observations": camera_indices,
         "point_index_of_observations": point_indices,
     }
-

datapipes/bal_io.py

@@ -83,7 +83,7 @@ def _ensure_problem_available(dataset: str, problem_name: str, cache_dir: Path)
     os.replace(tmp_path, txt_path)
     return txt_path
 
-def get_problem(problem_name, dataset, cache_dir='bal_data', use_quat=False):
+def get_problem(problem_name, dataset, cache_dir='bal_data', use_quat=True):
     cache_path = Path(cache_dir)
     print(f"Preparing data for {dataset}...")
     _validate_dataset(dataset)

datapipes/bal_loader.py

@@ -19,7 +19,7 @@
 if str(_REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(_REPO_ROOT))
 
-from ba_helpers import Reproj, project, least_square_error  # noqa: E402
+from ba_example import Residual, project, least_square_error  # noqa: E402
 from bae.autograd.function import TrackingTensor, map_transform
 import bae.autograd.graph as autograd_graph  # noqa: E402
 from bae.optim import LM  # noqa: E402
@@ -234,11 +234,11 @@ def _final_bal_per_pixel_error(
     point_idx: torch.Tensor,
 ) -> float:
     input = {
-        "points_2d": points_2d,
-        "camera_indices": camera_idx,
-        "point_indices": point_idx,
+        "observes": points_2d,
+        "cidx": camera_idx,
+        "pidx": point_idx,
     }
-    model = Reproj(camera_params.clone(), points_3d.clone())
+    model = Residual(camera_params.clone(), points_3d.clone())
     strategy = pp.optim.strategy.TrustRegion(up=2.0, down=0.5**4)
     solver = PCG(tol=1e-4, maxiter=250)
     optimizer = LM(model, strategy=strategy, solver=solver, reject=30)
@@ -248,7 +248,7 @@ def _final_bal_per_pixel_error(
 
     return least_square_error(
         model.pose,
-        model.points_3d,
+        model.points,
         camera_idx,
         point_idx,
         points_2d,
@@ -284,16 +284,16 @@ def test_bal_jacobian_structure_no_empty_columns(
     camera_idx = camera_idx.to(device=device)
     point_idx = point_idx.to(device=device)
 
-    model = Reproj(camera_params.clone(), points_3d.clone()).to(device)
+    model = Residual(camera_params.clone(), points_3d.clone()).to(device)
     residual = model(points_2d, camera_idx, point_idx)
     n_obs = int(points_2d.shape[0])
 
-    J_cam, J_pts = autograd_graph.jacobian(residual, [model.pose, model.points_3d])
+    J_cam, J_pts = autograd_graph.jacobian(residual, [model.pose, model.points])
     assert J_cam.layout == torch.sparse_bsr
     assert J_pts.layout == torch.sparse_bsr
 
     n_cams = model.pose.shape[0]
-    n_pts = model.points_3d.shape[0]
+    n_pts = model.points.shape[0]
 
     assert J_cam.shape == (n_obs * 2, n_cams * 9)
     assert J_pts.shape == (n_obs * 2, n_pts * 3)
@@ -414,7 +414,7 @@ def _final_bal_per_pixel_error_fixed_first_camera_cat(
     camera_all = torch.cat([camera_se3_all, model.intrinsics.tensor()], dim=-1)
     return least_square_error(
         camera_all,
-        model.points_3d.tensor(),
+        model.points.tensor(),
         camera_idx,
         point_idx,
         points_2d,
@@ -603,9 +603,9 @@ def test_bal_jacobian_structure_assert_failed_when_missing_observation_appearanc
         n_pts=n_pts,
     )
 
-    model = Reproj(camera_params.clone(), points_3d.clone()).to(device)
+    model = Residual(camera_params.clone(), points_3d.clone()).to(device)
     residual = model(points_2d, camera_idx2, point_idx2)
-    J_cam, J_pts = autograd_graph.jacobian(residual, [model.pose, model.points_3d])
+    J_cam, J_pts = autograd_graph.jacobian(residual, [model.pose, model.points])
 
     with pytest.raises(AssertionError):
         _assert_bal_correctness_criteria(