Method forward

monai/networks/blocks/crf.py:71–111

Args: input_tensor: tensor containing initial class logits. reference_tensor: the reference tensor used to guide the message passing. Returns: output (torch.Tensor): output tensor.

(self, input_tensor: torch.Tensor, reference_tensor: torch.Tensor)

        self.compatibility_matrix = compatibility_matrix

    def forward(self, input_tensor: torch.Tensor, reference_tensor: torch.Tensor):
        """
        Args:
            input_tensor: tensor containing initial class logits.
            reference_tensor: the reference tensor used to guide the message passing.

        Returns:
            output (torch.Tensor): output tensor.
        """

        # constructing spatial feature tensor
        spatial_features = _create_coordinate_tensor(reference_tensor)

        # constructing final feature tensors for bilateral and gaussian kernel
        bilateral_features = torch.cat(
            [spatial_features / self.bilateral_spatial_sigma, reference_tensor / self.bilateral_color_sigma], dim=1
        )
        gaussian_features = spatial_features / self.gaussian_spatial_sigma

        # setting up output tensor
        output_tensor = softmax(input_tensor, dim=1)

        # mean field loop
        for _ in range(self.iterations):
            # message passing step for both kernels
            bilateral_output = PHLFilter.apply(output_tensor, bilateral_features)
            gaussian_output = PHLFilter.apply(output_tensor, gaussian_features)

            # combining filter outputs
            combined_output = self.bilateral_weight * bilateral_output + self.gaussian_weight * gaussian_output

            # optionally running a compatibility transform
            if self.compatibility_matrix is not None:
                flat = combined_output.flatten(start_dim=2).permute(0, 2, 1)
                flat = torch.matmul(flat, self.compatibility_matrix)
                combined_output = flat.permute(0, 2, 1).reshape(combined_output.shape)

            # update and normalize
            output_tensor = softmax(input_tensor + self.update_factor * combined_output, dim=1)

        return output_tensor


# helper methods