Skip to main content

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.

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.

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

args
Namespace
required
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
dict
required
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) → points that maps (B, N, 2) pixel coords in image 1 to image 2.
  • batch["transform_points_2_to_1"] — the inverse callable.
image1
Tensor
required
Feature map for the first image, shape (B, C, H, W).
image2
Tensor
required
Feature map for the second image, shape (B, C, H, W).
Returns (image1_registered, image2_registered) — a 2-tuple of Tensor objects, each (B, C, H, W). Both tensors have undergone visibility-masked feature subtraction:
image1_registered = visibility2 * (image1 - image2_warped_onto_image1)
Regions not visible in the warped view receive a visibility weight of zero, suppressing false-positive change signals at occlusion boundaries.
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

# Inside Model.forward()
registration_module = FeatureRegisterationModule(args)

# features1, features2: (B, C, H, W) tensors from the shared encoder
image1_registered, image2_registered = registration_module(
    batch,        # batch["registration_strategy"] drives routing
    features1,
    features2,
)

# batch now contains transform callables for bbox matching
projected = batch["transform_points_1_to_2"](left_bbox_centers)

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
dict
required
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.
features1
Tensor
required
Feature map for image 1, shape (B, C, H, W).
features2
Tensor
required
Feature map for image 2, shape (B, C, H, W).
Returns a 4-tuple:
image1_warped_onto_image2
Tensor
Features from image 1 rendered into the coordinate frame of image 2.
image2_warped_onto_image1
Tensor
Features from image 2 rendered into the coordinate frame of image 1.
transform_points_1_to_2
callable
Closure that applies the estimated rigid transform to pixel coordinates.
transform_points_2_to_1
callable
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
dict
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.
features1
Tensor
required
Feature map for image 1, shape (B, C, H, W).
features2
Tensor
required
Feature map for image 2, shape (B, C, H, W).
Returns the same 4-tuple as 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
dict
required
Batch slice for "identity" samples. No geometry keys are read.
features1
Tensor
required
Feature map for image 1, shape (B, C, H, W).
features2
Tensor
required
Feature map for image 2, shape (B, C, H, W).
Returns the same 4-tuple. Visibility is set to 1 everywhere (no occlusion masking), and both 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.
features_src
Tensor
required
Source feature map, shape (B, C, H, W).
depth_src
Tensor
required
Per-pixel depth for the source view, shape (B, H, W).
src_camera_K_inv
Tensor
required
Inverse intrinsic matrix for the source camera, shape (B, 3, 3).
dst_camera_K_inv
Tensor
required
Inverse intrinsic matrix for the destination camera, shape (B, 3, 3).
Rt_src_to_dst
Tensor
required
Rigid-body transform from source to destination camera frame, shape (B, 4, 4).
Rendering parameters (fixed):
  • radius = 1.5 / min(image_H, image_W) * 2.0
  • points_per_pixel = 8
The rasterisation radius is normalised by the shorter image side so that the point splat size scales correctly regardless of image resolution.

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.
points_in_3d
Tensor
required
Point positions in the PyTorch3D coordinate system, shape (B, N, 3).
features
Tensor
required
Per-point feature vectors, shape (B, N, C).
Returns a rendered feature image 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.
intrinsics1
Tensor
required
Intrinsic matrix for camera 1, shape (B, 3, 3).
intrinsics2
Tensor
required
Intrinsic matrix for camera 2, shape (B, 3, 3).
rotation1
Tensor
required
Rotation matrix for camera 1 in world frame, shape (B, 3, 3).
rotation2
Tensor
required
Rotation matrix for camera 2 in world frame, shape (B, 3, 3).
position1
Tensor
required
World-space position of camera 1, shape (B, 3).
position2
Tensor
required
World-space position of camera 2, shape (B, 3).
Returns (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.
points1
Tensor
required
2D keypoint locations in image 1, shape (B, N, 2), normalised to [0, 1].
points2
Tensor
required
Corresponding 2D keypoint locations in image 2, shape (B, N, 2), normalised to [0, 1].
depth1
Tensor
required
Depth values sampled at points1, shape (B, N).
depth2
Tensor
required
Depth values sampled at points2, shape (B, N).
Returns (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).
When using estimate_Rt_using_points, the quality of the estimated transform depends on the accuracy and spread of the input correspondences. Sparse or near-planar point sets may produce degenerate solutions.

Build docs developers (and LLMs) love