Function align_surfaces

flowmap/model/projection.py:213–252

(
    surfaces: Float[Tensor, "batch frame height width 3"],
    backward_flows: Float[Tensor, "batch frame-1 height width xy=2"],
    backward_weights: Float[Tensor, "batch frame-1 height width"],
    indices: Int64[Tensor, " pixel_index"],
)

def align_surfaces(
    surfaces: Float[Tensor, "batch frame height width 3"],
    backward_flows: Float[Tensor, "batch frame-1 height width xy=2"],
    backward_weights: Float[Tensor, "batch frame-1 height width"],
    indices: Int64[Tensor, " pixel_index"],
) -> Float[Tensor, "batch frame 4 4"]:
    b, f, h, w, _ = surfaces.shape

    # Convert the depth maps into camera-space 3D surfaces (b, f, h, w, xyz).
    xy, _ = sample_image_grid((h, w), device=surfaces.device)

    # Subsample the surfaces to select points for Procrustes alignment. Select the later
    # points from the surfaces using the provided indices.
    xyz_later = rearrange(later(surfaces), "b f h w xyz -> b f (h w) xyz")
    xyz_later = xyz_later[:, :, indices]

    # Flow the grid of XY locations backwards, then select from the flowed XY locations
    # using the provided indices.
    xy_earlier = rearrange(xy + backward_flows, "b f h w xy -> b f (h w) xy")
    xy_earlier = xy_earlier[:, :, indices]

    # Use the earlier XY locations to select from the earlier 3D surfaces.
    xyz_earlier = F.grid_sample(
        rearrange(earlier(surfaces), "b f h w xyz -> (b f) xyz h w"),
        rearrange(xy_earlier * 2 - 1, "b f p xy -> (b f) p () xy"),
        mode="bilinear",
        padding_mode="border",
        align_corners=False,
    )
    xyz_earlier = rearrange(xyz_earlier, "(b f) xyz p () -> b f p xyz", b=b, f=f - 1)

    # Estimate poses via Procrustes alignment.
    inverse_relative_transformations = align_rigid(
        xyz_later,
        xyz_earlier,
        rearrange(backward_weights, "b f h w -> b f (h w)")[..., indices],
    )
    extrinsics = get_extrinsics(inverse_relative_transformations)

    return extrinsics


def compute_track_flow(