The registration module is responsible for spatially aligning feature maps extracted from two views of the same 3D scene before the change-detection decoder compares them. Depending on what geometry data is available in a batch, the module routes each sample through one of three strategies — 3D camera-pose-based warping, 2D linear warping, or identity (no warping) — and stores the resulting coordinate-transform callables back into the batch dictionary so that downstream bbox matching can use them.Documentation Index
Fetch the complete documentation index at: https://mintlify.com/ragavsachdeva/CYWS-3D/llms.txt
Use this file to discover all available pages before exploring further.
FeatureRegisterationModule
FeatureRegisterationModule is a torch.nn.Module that wraps the full registration pipeline. It is constructed once during model initialisation and called inside Model.forward() after the shared DINOv2 encoder has produced feature maps for both images.
Constructor
Parsed configuration namespace. No fields on
args are read directly in __init__; the object is stored for potential sub-module construction.forward()
Dispatches each sample in the batch to the correct registration strategy, then merges the results back into a single aligned batch.Batch dictionary produced by the dataloader. Must contain
"registration_strategy" — a list of strings, one per sample, where each entry is "3d", "2d", or "identity". Camera parameters and correspondences are read from the same dict when present. Modified in place: after forward() returns, batch will contain two new keys:batch["transform_points_1_to_2"]— callable(points) → pointsthat maps (B, N, 2) pixel coords in image 1 to image 2.batch["transform_points_2_to_1"]— the inverse callable.
Feature map for the first image, shape
(B, C, H, W).Feature map for the second image, shape
(B, C, H, W).(image1_registered, image2_registered) — a 2-tuple of Tensor objects, each (B, C, H, W). Both tensors have undergone visibility-masked feature subtraction:
Mixed batches — where some samples use
"3d", others "2d", and others "identity" — are handled transparently. The module uses slice_batch_given_bool_array to split the batch by strategy, processes each subset independently, and recombines the results before returning.Code example
Registration methods
register_3d_features()
Handles samples where full 3D geometry is available. Uses rigid-body camera transformations to project each image’s feature map onto the other’s view plane via depth-based unprojection and reprojection.Batch slice containing only the
"3d" samples. If "intrinsics" keys are present the rotation matrices and camera positions are used via estimate_Rt_using_camera_parameters; otherwise 2D–3D point correspondences are used via estimate_Rt_using_points.Feature map for image 1, shape
(B, C, H, W).Feature map for image 2, shape
(B, C, H, W).Features from image 1 rendered into the coordinate frame of image 2.
Features from image 2 rendered into the coordinate frame of image 1.
Closure that applies the estimated rigid transform to pixel coordinates.
Inverse closure.
register_2d_features()
Handles samples where only 2D correspondences or a pre-computed 2D transformation matrix are available. No depth information is required.Batch slice for
"2d" samples. If batch["transfm2d"] is present it is used directly; otherwise estimate_linear_warp is called on the provided 2D point correspondences.Feature map for image 1, shape
(B, C, H, W).Feature map for image 2, shape
(B, C, H, W).register_3d_features(). Warping is performed by DifferentiableFeatureWarper.render_features_from_points() rather than the full depth-based warp() path.
register_identity_features()
A no-op registration strategy for samples where the two images are already aligned (e.g. synthetic pairs or when no geometry metadata is available).Batch slice for
"identity" samples. No geometry keys are read.Feature map for image 1, shape
(B, C, H, W).Feature map for image 2, shape
(B, C, H, W).transform_points callables are identity functions that return their input unchanged.
DifferentiableFeatureWarper
DifferentiableFeatureWarper is a torch.nn.Module that performs differentiable rendering of feature maps using PyTorch3D’s PointsRenderer with an AlphaCompositor. It supports two entry points depending on whether depth maps are available.
warp()
Full 3D warp: unprojects image pixels to world space using depth, re-projects into the target camera, and composites the features.Source feature map, shape
(B, C, H, W).Per-pixel depth for the source view, shape
(B, H, W).Inverse intrinsic matrix for the source camera, shape
(B, 3, 3).Inverse intrinsic matrix for the destination camera, shape
(B, 3, 3).Rigid-body transform from source to destination camera frame, shape
(B, 4, 4).radius = 1.5 / min(image_H, image_W) * 2.0points_per_pixel = 8
render_features_from_points()
Lightweight renderer used by the 2D registration path. Accepts already-computed 3D point positions and renders their associated features without requiring a depth map.Point positions in the PyTorch3D coordinate system, shape
(B, N, 3).Per-point feature vectors, shape
(B, N, C).Tensor of shape (B, C, H, W). Rendering uses PerspectiveCameras with an identity transform (points are assumed to already be in the target camera’s coordinate system).
Helper functions
estimate_Rt_using_camera_parameters()
Computes inverse intrinsics and relative rigid-body transforms from known camera calibration and pose data.Intrinsic matrix for camera 1, shape
(B, 3, 3).Intrinsic matrix for camera 2, shape
(B, 3, 3).Rotation matrix for camera 1 in world frame, shape
(B, 3, 3).Rotation matrix for camera 2 in world frame, shape
(B, 3, 3).World-space position of camera 1, shape
(B, 3).World-space position of camera 2, shape
(B, 3).(K_inv_1, K_inv_2, Rt_1_to_2, Rt_2_to_1) — all Tensor objects. K_inv_* are (B, 3, 3); Rt_* are (B, 4, 4) homogeneous matrices.
estimate_Rt_using_points()
Estimates camera transforms from 2D–3D point correspondences when intrinsic/extrinsic calibration is not available. Assumes a canonical (identity) camera for unprojection.2D keypoint locations in image 1, shape
(B, N, 2), normalised to [0, 1].Corresponding 2D keypoint locations in image 2, shape
(B, N, 2), normalised to [0, 1].Depth values sampled at
points1, shape (B, N).Depth values sampled at
points2, shape (B, N).(K_inv, K_inv, Rt_1_to_2, Rt_2_to_1). Both K_inv entries are identical canonical-camera inverse intrinsics; Rt_* are (B, 4, 4).