correlation.py

Perform spatial correlation to generate a cost volume.

This is meant to be a simplified version of the SpatialCorrelationSampler from https://github.com/ClementPinard/Pytorch-Correlation-extension.

This version is implemented purely in PyTorch. However, it only supports correlation with 1x1 kernels. It is also not as efficient as the original SpatialCorrelationSampler.

class ptlflow.utils.correlation.IterSpatialCorrelationSampler(kernel_size: int | Tuple[int, int] = 1, patch_size: int | Tuple[int, int] = 1, stride: int | Tuple[int, int] = 1, padding: int | Tuple[int, int] = 0, dilation: int | Tuple[int, int] = 1, dilation_patch: int | Tuple[int, int] = 1)[source]

Apply spatial correlation sampling from two inputs using iteration in PyTorch.

Initialize IterSpatialCorrelationSampler.

Parameters:

kernel_sizeUnion[int, Tuple[int, int]], default 1: Total size of your correlation kernel, in pixels
patch_sizeUnion[int, Tuple[int, int]], default 1: Total size of your patch, determining how many different shifts will be applied.
strideUnion[int, Tuple[int, int]], default 1: Stride of the spatial sampler, will modify output height and width.
paddingUnion[int, Tuple[int, int]], default 0: Padding applied to input1 and input2 before applying the correlation sampling, will modify output height and width.
dilationUnion[int, Tuple[int, int]], default 1: Similar to dilation in convolution.
dilation_patchUnion[int, Tuple[int, int]], default 1: Step for every shift in patch.

forward(input1: Tensor, input2: Tensor) → Tensor[source]

Compute the correlation sampling from input1 to input2.

Parameters:

input1torch.Tensor: The origin feature map.
input2torch.Tensor: The target feature map.

Returns:

torch.Tensor: Result of correlation sampling.

class ptlflow.utils.correlation.IterTranslatedSpatialCorrelationSampler(kernel_size: int | Tuple[int, int] = 1, patch_size: int | Tuple[int, int] = 1, stride: int | Tuple[int, int] = 1, padding: int | Tuple[int, int] = 0, dilation: int | Tuple[int, int] = 1, dilation_patch: int | Tuple[int, int] = 1)[source]

Apply translated spatial correlation sampling from two inputs using iteration in PyTorch.

This operation is equivalent to first translating the points from input1 using the given flow, and then doing a local correlation sampling around the translated points.

This allows us to do correlation sampling without warping the second input.

Initialize IterTranslatedSpatialCorrelationSampler.

Parameters:

kernel_sizeUnion[int, Tuple[int, int]], default 1: Total size of your correlation kernel, in pixels
patch_sizeUnion[int, Tuple[int, int]], default 1: Total size of your patch, determining how many different shifts will be applied.
strideUnion[int, Tuple[int, int]], default 1: Stride of the spatial sampler, will modify output height and width.
paddingUnion[int, Tuple[int, int]], default 0: Padding applied to input1 and input2 before applying the correlation sampling, will modify output height and width.
dilationUnion[int, Tuple[int, int]], default 1: Similar to dilation in convolution.
dilation_patchUnion[int, Tuple[int, int]], default 1: Step for every shift in patch.

forward(input1: Tensor, input2: Tensor, flow: Tensor) → Tensor[source]

Compute the correlation sampling from input1 to input2.

Parameters:

input1torch.Tensor: The origin feature map.
input2torch.Tensor: The target feature map.
flowtorch.Tensor: The optical flow field to translate the points from input1. The flow values should be represented in number of pixels (do not provide normalized values, e.g. between -1 and 1). It should be a 4D tensor (b, 2, h, w), where flow[:, 0] represent the horizontal flow and flow[:, 1] the vertical ones.

Returns:

torch.Tensor: Result of correlation sampling.

class ptlflow.utils.correlation.IterativeCorrBlock(fmap1: Tensor, fmap2: Tensor, radius: int = 1, num_levels: int = 1)[source]

Another wrapper for iter_translated_spatial_correlation_sample.

This block is designed to mimic the operations of RAFT’s AlternateCorrBlock package (see ptlflow/models/raft/corr.py). This block can be used when alt_cuda_corr has not been compiled (see ptlflow/utils/external/alt_cuda_corr).

IMPORTANT: this implementation is slower than alt_cuda_corr.

__init__(fmap1: Tensor, fmap2: Tensor, radius: int = 1, num_levels: int = 1)[source]

Initialize IterativeCorrBlock.

Parameters:

fmap1torch.Tensor: The origin feature map.
fmap2torch.Tensor: The target feature map.
radiusint, default 1: The radius if the correlation patch. The patch_size will be 2 * radius + 1.
num_levelsint, default 1: Number of correlation pooling levels to use (see ptlflow/models/raft/corr.py).

forward(coords)[source]

Compute the correlation sampling from input1 to input2.

Parameters:

coordstorch.Tensor: The addition (optical flow + coords_grid) to translate the points from input1. The coords values should be represented in number of pixels (do not provide normalized values, e.g. between -1 and 1). It should be a 4D tensor (b, 2, h, w), where coords[:, 0] represent the x axis and flow[:, 1] the y axis.

Returns:

torch.Tensor: Result of correlation sampling.

ptlflow.utils.correlation.iter_spatial_correlation_sample(input1: Tensor, input2: Tensor, kernel_size: int | Tuple[int, int] = 1, patch_size: int | Tuple[int, int] = 1, stride: int | Tuple[int, int] = 1, padding: int | Tuple[int, int] = 0, dilation: int | Tuple[int, int] = 1, dilation_patch: int | Tuple[int, int] = 1) → Tensor[source]

Apply spatial correlation sampling from input1 to input2 using iteration in PyTorch.

This docstring is taken and adapted from the original package.

Every parameter except input1 and input2 can be either single int or a pair of int. For more information about Spatial Correlation Sampling, see this page. https://lmb.informatik.uni-freiburg.de/Publications/2015/DFIB15/

Parameters:

input1torch.Tensor: The origin feature map.
input2torch.Tensor: The target feature map.
kernel_sizeUnion[int, Tuple[int, int]], default 1: Total size of your correlation kernel, in pixels
patch_sizeUnion[int, Tuple[int, int]], default 1: Total size of your patch, determining how many different shifts will be applied.
strideUnion[int, Tuple[int, int]], default 1: Stride of the spatial sampler, will modify output height and width.
paddingUnion[int, Tuple[int, int]], default 0: Padding applied to input1 and input2 before applying the correlation sampling, will modify output height and width.
dilationUnion[int, Tuple[int, int]], default 1: Similar to dilation in convolution.
dilation_patchUnion[int, Tuple[int, int]], default 1: Step for every shift in patch.

Returns:

torch.Tensor: Result of correlation sampling.

Raises:

NotImplementedError: If kernel_size != 1.
NotImplementedError: If dilation != 1.

ptlflow.utils.correlation.iter_translated_spatial_correlation_sample(input1: Tensor, input2: Tensor, flow: Tensor | None = None, coords: Tensor | None = None, kernel_size: int | Tuple[int, int] = 1, patch_size: int | Tuple[int, int] = 1, stride: int | Tuple[int, int] = 1, padding: int | Tuple[int, int] = 0, dilation: int | Tuple[int, int] = 1, dilation_patch: int | Tuple[int, int] = 1, coords_grid: Tensor | None = None) → Tensor[source]

Apply spatial correlation sampling with translation from input1 to input2 using iteration in PyTorch.

This operation is equivalent to first translating the points from input1 using the given flow, and then doing a local correlation sampling around the translated points.

This allows us to do correlation sampling without warping the second input.

Every parameter except input1, input2, and flow can be either single int or a pair of int. For more information about Spatial Correlation Sampling (without translation), see this page: https://lmb.informatik.uni-freiburg.de/Publications/2015/DFIB15/

Parameters:

input1torch.Tensor: The origin feature map.
input2torch.Tensor: The target feature map.
flowOptional[torch.Tensor]: This argument and “coords” are mutually exclusive, only one of them can be not None. The optical flow field to translate the points from input1. The flow values should be represented in number of pixels (do not provide normalized values, e.g. between -1 and 1). It should be a 4D tensor (b, 2, h, w), where flow[:, 0] represent the horizontal flow and flow[:, 1] the vertical ones.
coordstorch.Tensor: This argument and “flow” are mutually exclusive, only one of them can be not None. This value should be equivalent to “flow” + “coords_grid”.
kernel_sizeUnion[int, Tuple[int, int]], default 1: Total size of your correlation kernel, in pixels
patch_sizeUnion[int, Tuple[int, int]], default 1: Total size of your patch, determining how many different shifts will be applied.
strideUnion[int, Tuple[int, int]], default 1: Stride of the spatial sampler, will modify output height and width.
paddingUnion[int, Tuple[int, int]], default 0: Padding applied to input1 and input2 before applying the correlation sampling, will modify output height and width.
dilationUnion[int, Tuple[int, int]], default 1: Similar to dilation in convolution.
dilation_patchUnion[int, Tuple[int, int]], default 1: Step for every shift in patch.
coords_gridOptional[torch.Tensor], default None: A tensor with the same shape as flow containing a grid of 2D coordinates of the pixels. This can be created using torch.meshgrid. This parameter is optional. If not provided, the grid will be created internally. Only useful if the grid can be buffered somewhere to avoid recreating it at every call.

Returns:

torch.Tensor: Result of correlation sampling.

Raises:

NotImplementedError: If kernel_size != 1.
NotImplementedError: If dilation != 1.