Modifying reprojection cost function

[AlleUndKalle]

Thanks a lot for this great library!

I am trying to implement a small modification over Reprojection cost function used for the bundle adjustment example. I would like to add another SE(3) transformation (we call it point_pose) before camera pose transformation and optimize for both camera_pose and point_pose jointly. And I would like to disable optimizing the world_point. I couldn't figure out how I can calculate the jacobian for this point_pose variable. Here is the modifications I made:

class Reprojection(th.CostFunction):
    def __init__(
        self,
        camera_pose: th.SE3,
        point_pose: th.SE3,
        world_point: th.Point3,
        image_feature_point: th.Point2,
        focal_length: th.Vector,
        weight: Optional[th.CostWeight] = None,
        name: Optional[str] = None,
    ):
        if weight is None:
            weight = th.ScaleCostWeight(torch.tensor(1.0).to(dtype=camera_pose.dtype))
        super().__init__(
            cost_weight=weight,
            name=name,
        )
        self.camera_pose = camera_pose
        self.point_pose = point_pose
        self.focal_length = focal_length
        self.world_point = world_point
        self.image_feature_point = image_feature_point
        
        self.register_optim_vars(["camera_pose", "point_pose"])
        self.register_aux_vars(
            ["focal_length", "image_feature_point", "world_point"]
        )

    def error(self) -> torch.Tensor:
        self.world_point = self.point_pose.transform_from(self.world_point)
        point_cam = self.camera_pose.transform_from(self.world_point)
        proj = -point_cam[:, :2] / point_cam[:, 2:3]
        proj_sqn = (proj * proj).sum(dim=1).unsqueeze(1)
        proj_factor = self.focal_length.tensor * (
            1.0 + proj_sqn * (0 + proj_sqn * 0)
        )
        point_projection = proj * proj_factor

        err = point_projection - self.image_feature_point.tensor
        
        return err

    def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
        ppose_wpt_jacs: List[torch.Tensor] = [] # <-- where to use this list of jacobians??
        self.world_point = self.point_pose.transform_from(self.world_point, ppose_wpt_jacs)
        cpose_wpt_jacs: List[torch.Tensor] = []
        point_cam = self.camera_pose.transform_from(self.world_point, cpose_wpt_jacs)
        J = torch.cat(cpose_wpt_jacs, dim=2)
        
        proj = -point_cam[:, :2] / point_cam[:, 2:3]
        proj_sqn = (proj * proj).sum(dim=1).unsqueeze(1)
        proj_factor = self.focal_length.tensor * (
            1.0 + proj_sqn * (0. + proj_sqn * 0.)
        )
        d_proj_factor = self.focal_length.tensor * (
            0. + 2.0 * proj_sqn * 0.
        )
        point_projection = proj * proj_factor

        # derivative of N/D is (N' - ND'/D) / D
        d_num = J[:, 0:2, :]
        num_dden_den = torch.bmm(
            point_cam[:, :2].unsqueeze(2),
            (J[:, 2, :] / point_cam[:, 2:3]).unsqueeze(1),
        )
        proj_jac = (num_dden_den - d_num) / point_cam[:, 2:].unsqueeze(2)
        proj_sqn_jac = 2.0 * proj.unsqueeze(2) * torch.bmm(proj.unsqueeze(1), proj_jac)
        point_projection_jac = proj_jac * proj_factor.unsqueeze(
            2
        ) + proj_sqn_jac * d_proj_factor.unsqueeze(2)

        err = point_projection - self.image_feature_point.tensor

        return [point_projection_jac[..., :6]], err

    def dim(self) -> int:
        return 2

    def to(self, *args, **kwargs):
        super().to(*args, **kwargs)

    def _copy_impl(self, new_name: Optional[str] = None) -> "Reprojection":
        return Reprojection(
            self.camera_pose.copy(),
            self.point_pose.copy(),
            self.world_point.copy(),
            self.image_feature_point.copy(),
            self.focal_length.copy(),
            weight=self.weight.copy(),
            name=new_name,
        )

I specifically would like to know how I can make use of ppose_wpt_jacs to be able to optimize for point_pose as well.

[maurimo]

Quoting:

self.world_point = self.point_pose.transform_from(self.world_point)
point_cam = self.camera_pose.transform_from(self.world_point)

Probably you don't want to modify self.world_point in place (also to avoid messing with optimizer internals), I would define a temporary:

pose_applied_pt = self.point_pose.transform_from(self.world_point)
point_cam = self.camera_pose.transform_from(pose_applied_pt)

To compute the Jacobian, you need to compute the jacobian of camPt wrt both camera_pose and point_pose:

dPosePt_dPtPose_jacs: List[torch.Tensor] = [] # <-- where to use this list of jacobians??
pose_applied_pt = self.point_pose.transform_from(self.world_point, dPosePt_dPtPose_jacs)
dCamPt_dCamPose_jacs: List[torch.Tensor] = []
point_cam = self.camera_pose.transform_from(pose_applied_pt, dCamPt_dCamPose_jacs)
dPosePt_dPtPose = torch.cat(dPosePt_dPtPose_jacs, dim=2)  # JACOBIAN matrix of ppose
dCamPt_dCamPose = torch.cat(dCamPt_dCamPose_jacs, dim=2)
dCamPt_dPtPose = self.camera_pose.adjoint() @ dPosePt_dPtPose  # check if it's not called `adjoint`... should exist
J = torch.cat([dCamPt_dCamPose, dCamPt_dPtPose], dim=2)         # NOW J contains the jacobian of `camPt` wrt both `camera_pose` and `point_pose`, on sets of colums 0:6 and 6:12
... everything identical ...
# at the end `point_projection_jac` has jacobian of projected point wrt `camera_pose` and `point_pose`, therefore:
return [point_projection_jac[..., :6], point_projection_jac[..., 6:12]], err

[AlleUndKalle]

Thanks a lot for your answer @maurimo!!

[luisenp]

@AlleUndKalle Also, in case you don't know, you can check the correctness of your jacobian implementation using this numerical jacobian utility, which we use in our unit tests. Here you can find an example of how to use it.

[AlleUndKalle]

Thanks a lot @luisenp! I will give it a try