Image Registration

class AbstractRegistration(ABC):
    """Base class for all registration algorithms in FireANTs.

    This abstract class provides the core functionality and interface for image registration,
    handling all the common functionality for all linear and non-linear registration algorithms.
    It handles features like
        - multi-resolution optimization,
        - arbitrary similarity metrics, and
        - convergence monitoring.

    Args:
        scales (List[int]): Downsampling factors for multi-resolution registration.
            Must be in descending order (e.g. [4,2,1]).
        iterations (List[float]): Number of iterations to perform at each scale.
            Must match length of scales.
        fixed_images (BatchedImages): Fixed/reference images to register to.
        moving_images (BatchedImages): Moving images to be registered.
        loss_type (str, optional): Similarity metric to use. A set of predefined loss functions are provided:
            - 'mi': Mutual information
                Global mutual information between the fixed image and the moved image is measured using Parzen windowing using different kernel types.
                Mutual information is implemented using the `GlobalMutualInformationLoss` class.
            - 'cc': Cross correlation (default)
                Local normalized cross correlation between the fixed image and the moved image is measured using a kernel of size `cc_kernel_size`. This is a de-facto standard similarity metric for both linear and non-linear image registration. It also consumes less very little memory than mutual information.
                Cross correlation is implemented using the `LocalNormalizedCrossCorrelationLoss` class.
            - 'mse': Mean squared error
                Mean squared error is implemented using the `MeanSquaredError` class.
                This is the fastest similarity metric, but it is not robust to outliers, multi-modal images, or even intensity inhomogeneities within the same modality. Good for testing registration pipelines.
            - 'custom': Custom loss function (see `custom_loss` for more details)
            - 'noop': No similarity metric  (this is useful for testing regularizations)
        mi_kernel_type (str, optional): Kernel type for mutual information loss. Default: 'gaussian'
        cc_kernel_type (str, optional): Kernel type for cross correlation loss. Default: 'rectangular'
        custom_loss (nn.Module, optional): Custom loss module if loss_type='custom'.
            See [Custom Loss Functions](../advanced/customloss.md) for more details on how to implement custom loss functions.
        loss_params (dict, optional): Additional parameters for loss function. See the documentation of the loss function for more details on what parameters are available. This implementation abstracts away the loss function implementation details from the registration pipeline.
        cc_kernel_size (int, optional): Kernel size for cross correlation loss. Default: 3
        reduction (str, optional): Loss reduction method. Default: 'mean'
        tolerance (float, optional): Convergence tolerance. Default: 1e-6
        max_tolerance_iters (int, optional): Max iterations for convergence check. Default: 10
        progress_bar (bool, optional): Whether to show progress bar. Default: True

    Methods:
        optimize(): Abstract method to perform registration optimization
        get_warped_coordinates(): Abstract method to get transformed coordinates
        evaluate(): Apply learned transformation to new images

    Note:
        The number of scales and iterations must match, and scales must be in descending order.
        The fixed and moving images must be broadcastable in batch dimension.
    """

    def __init__(self,
                scales: List[float], iterations: List[int],
                fixed_images: BatchedImages, moving_images: BatchedImages,
                loss_type: str = "cc",
                mi_kernel_type: str = 'gaussian', cc_kernel_type: str = 'rectangular',
                custom_loss: nn.Module = None,
                loss_params: dict = {},
                cc_kernel_size: int = 3,
                reduction: str = 'mean',
                tolerance: float = 1e-6, max_tolerance_iters: int = 10,
                progress_bar: bool = True,
                dtype: torch.dtype = torch.float32,
                ) -> None:
        '''
        Initialize abstract registration class
        '''
        super().__init__()
        self.scales = scales
        _assert_check_scales_decreasing(self.scales)
        self.iterations = iterations
        assert len(self.iterations) == len(self.scales), "Number of iterations must match number of scales"
        # check for fixed and moving image sizes
        self.fixed_images = fixed_images
        self.moving_images = moving_images
        # assert (self.fixed_images.size() == self.moving_images.size()), "Number of fixed and moving images must match"

        # check if sizes are broadcastable
        fsize, msize = self.fixed_images.size(), self.moving_images.size()
        assert (fsize == msize) or (fsize == 1) or (msize == 1), "Number of fixed and moving images must match or broadcastable"
        self.opt_size = max(fsize, msize)

        self.tolerance = tolerance
        self.max_tolerance_iters = max_tolerance_iters
        self.convergence_monitor = ConvergenceMonitor(self.max_tolerance_iters, self.tolerance)

        self.device = fixed_images.device
        self.dtype = dtype
        if not is_torch_float_type(self.dtype):
            raise ValueError(f"non-float dtype {self.dtype} is not supported for registration")

        self.dims = self.fixed_images.dims
        self.progress_bar = progress_bar        # variable to show or hide progress bar

        # initialize losses
        # track whether we are in masked mode; this is driven by the loss configuration
        if loss_type.startswith("masked_"):
            loss_params['masked'] = True
            self.masked = True
            loss_type = loss_type.replace("masked_", "")
            logger.info(f"Masked mode specified, will use {loss_type} function with masking if it supports it.")
        else:
            loss_params['masked'] = False
            self.masked = False
            logger.info(f"Masked mode not specified, will not use masking.")

        if loss_type == 'mi':
            self.loss_fn = GlobalMutualInformationLoss(kernel_type=mi_kernel_type, reduction=reduction, **loss_params)
        elif loss_type == 'cc':
            self.loss_fn = LocalNormalizedCrossCorrelationLoss(kernel_type=cc_kernel_type, spatial_dims=self.dims,
                                                               kernel_size=cc_kernel_size, reduction=reduction, **loss_params)
        elif loss_type == 'fusedcc':
            try:
                from fireants.losses.fusedcc import FusedLocalNormalizedCrossCorrelationLoss
                self.loss_fn = FusedLocalNormalizedCrossCorrelationLoss(spatial_dims=self.dims,
                                                        kernel_size=cc_kernel_size, reduction=reduction, **loss_params)
            except ImportError:
                logger.warning("fireants_fused_ops not available, falling back to non-fused LocalNormalizedCrossCorrelationLoss")
                self.loss_fn = LocalNormalizedCrossCorrelationLoss(kernel_type=cc_kernel_type, spatial_dims=self.dims,
                                                                   kernel_size=cc_kernel_size, reduction=reduction, **loss_params)
        elif loss_type == 'lite_lncc':
            try:
                from fireants.losses.lite_lncc_backend import LiteLNCCLoss
                self.loss_fn = LiteLNCCLoss(spatial_dims=self.dims,
                                            kernel_size=cc_kernel_size, reduction=reduction, **loss_params)
            except ImportError:
                logger.warning("lite_lncc not available, falling back to non-fused LocalNormalizedCrossCorrelationLoss")
                self.loss_fn = LocalNormalizedCrossCorrelationLoss(kernel_type=cc_kernel_type, spatial_dims=self.dims,
                                                                   kernel_size=cc_kernel_size, reduction=reduction, **loss_params)
        elif loss_type == 'fusedmi':
            try:
                from fireants.losses.fusedmi import FusedGlobalMutualInformationLoss
                self.loss_fn = FusedGlobalMutualInformationLoss(kernel_type=mi_kernel_type, reduction=reduction, **loss_params)
            except ImportError:
                logger.warning("fireants_fused_ops not available, falling back to non-fused GlobalMutualInformationLoss")
                self.loss_fn = GlobalMutualInformationLoss(kernel_type=mi_kernel_type, reduction=reduction, **loss_params)
        elif loss_type == 'custom':
            self.loss_fn = custom_loss
        elif loss_type == 'noop':
            self.loss_fn = NoOp()
        elif loss_type == 'mse':
            self.loss_fn = MeanSquaredError(reduction=reduction, **loss_params)
        else:
            raise ValueError(f"Loss type {loss_type} not supported")

        # see if loss can store the iterations
        if hasattr(self.loss_fn, 'set_iterations'):
            logger.info("Setting iterations for loss function")
            self.loss_fn.set_iterations(self.iterations)
        if hasattr(self.loss_fn, 'set_scales'):
            logger.info("Setting scales for loss function")
            self.loss_fn.set_scales(self.scales)

        self.print_init_msg()

    def print_init_msg(self):
        logger.info(f"Registration of type {self.__class__.__name__} initialized with dtype {self.dtype}")

    def _split_image_and_mask_last_channel(self, arrays: torch.Tensor):
        """Split arrays into (image_channels, mask_channel) if in masked mode.

        Convention: when `self.masked` is True, the last channel is treated as a mask
        and should not be Gaussian-smoothed. All preceding channels are treated as image.

        Returns:
            (img, mask): img is the tensor to be smoothed; mask is the last channel (kept unsmoothed),
            or None if not in masked mode / no explicit mask channel.
        """
        if getattr(self, "masked", False) and arrays.shape[1] > 1:
            return arrays[:, :-1], arrays[:, -1:]
        return arrays, None

    def _concat_image_and_mask_last_channel(self, img: torch.Tensor, mask: Optional[torch.Tensor]):
        """Inverse of `_split_image_and_mask_last_channel`."""
        if mask is None:
            return img
        return torch.cat([img, mask], dim=1)

    def _downsample_image_and_mask(
        self,
        arrays: torch.Tensor,
        *,
        size,
        mode: str,
        gaussians=None,
        align_corners: bool = True,
    ) -> torch.Tensor:
        """Downsample arrays while respecting masked-last-channel convention.

        - If `gaussians` is provided, we Gaussian-smooth + downsample the image channels via `downsample`.
        - The (optional) mask channel is *never* Gaussian-smoothed; it is resized via interpolation only.
        """
        img, mask = self._split_image_and_mask_last_channel(arrays)
        if gaussians is None:
            img_down = F.interpolate(img, size=size, mode=mode, align_corners=align_corners)
        else:
            img_down = downsample(img, size=size, mode=mode, gaussians=gaussians)

        if mask is None:
            return img_down
        mask_down = F.interpolate(mask, size=size, mode=mode, align_corners=align_corners)
        return self._concat_image_and_mask_last_channel(img_down, mask_down)

    def _smooth_image_not_mask(self, arrays: torch.Tensor, gaussians) -> torch.Tensor:
        """Gaussian-smooth image channels but not the (optional) mask channel."""
        img, mask = self._split_image_and_mask_last_channel(arrays)
        img_smooth = separable_filtering(img, gaussians)
        return self._concat_image_and_mask_last_channel(img_smooth, mask)

    @abstractmethod
    def optimize(self):
        '''
        Abstract method to perform registration optimization
        '''
        pass

    @abstractmethod
    def get_warp_parameters(self, fixed_images: Union[BatchedImages, FakeBatchedImages], moving_images: Union[BatchedImages, FakeBatchedImages], shape=None):
        ''' Get dictionary of parameters to pass into fireants_interpolator '''
        raise NotImplementedError("This method must be implemented by the registration class")

    @abstractmethod
    def get_inverse_warp_parameters(self, fixed_images: Union[BatchedImages, FakeBatchedImages], moving_images: Union[BatchedImages, FakeBatchedImages], shape=None):
        ''' Get dictionary of parameters to pass into fireants_interpolator '''
        raise NotImplementedError("This method must be implemented by the registration class")

    def get_warped_coordinates(self, fixed_images: Union[BatchedImages, FakeBatchedImages], moving_images: Union[BatchedImages, FakeBatchedImages], shape=None):
        '''Get the transformed coordinates for warping the moving image.

        This abstract method must be implemented by all registration classes to define how
        coordinates are transformed from the fixed image space to the moving image space.
        The transformed coordinates are used with grid_sample to warp the moving image.

        Args:
            fixed_images (BatchedImages): Fixed/reference images that define the target space.
                The output coordinates will be in this image's coordinate system.
            moving_images (BatchedImages): Moving images to be transformed.
                Used to access the original image space parameters.
            shape (Optional[Tuple[int, ...]]): Optional output shape for the coordinate grid.
                If None, uses the shape of the fixed image.

        Returns:
            torch.Tensor: Transformed coordinates in normalized [-1, 1] space.
                Has shape [N, H, W, [D], dims] where:
                    N = batch size
                    H, W, [D] = spatial dimensions (D only for 3D)
                    dims = number of spatial dimensions (2 or 3)
                These coordinates can be used directly with `torch.nn.functional.grid_sample()`

        Note:
            - Coordinates are returned in the normalized [-1, 1] coordinate system required by grid_sample
            - The transformation maps from fixed image space to moving image space (backward transform)
            - Physical space transformations are handled internally using the image metadata
        '''
        params = self.get_warp_parameters(fixed_images, moving_images, shape)
        if 'affine' in params and 'grid' not in params:
            return F.affine_grid(params['affine'], params['out_shape'], align_corners=True)
        elif 'affine' in params and 'grid' in params:
            # this is just a warp field
            affine = params['affine']
            grid = params['grid']
            grid = fireants_interpolator.affine_warp(affine, grid, align_corners=True)
            return grid
        else:
            raise ValueError(f"Invalid warp parameters with keys {params.keys()}")

    def get_inverse_warped_coordinates(self, fixed_images: Union[BatchedImages, FakeBatchedImages], moving_images: Union[BatchedImages, FakeBatchedImages], shape=None):
        ''' Get inverse warped coordinates for the moving image.

        This method is useful to analyse the effect of how the moving coordinates (fixed images) are transformed
        '''
        params = self.get_inverse_warp_parameters(fixed_images, moving_images, shape)
        if 'affine' in params and 'grid' not in params:
            return F.affine_grid(params['affine'], params['out_shape'], align_corners=True)
        elif 'affine' in params and 'grid' in params:
            # this is just a warp field
            affine = params['affine']
            grid = params['grid']
            shape = [affine.shape[0], 1] + list(grid.shape[1:-1])
            grid = fireants_interpolator.affine_warp(affine, grid, align_corners=True)
            return grid
        else:
            raise ValueError(f"Invalid warp parameters with keys {params.keys()}")

    def save_moved_images(self, moved_images: Union[BatchedImages, FakeBatchedImages, torch.Tensor], filenames: Union[str, List[str]], moving_to_fixed: bool = True, ignore_size_match: bool = False):
        '''
        Save the moved images to disk.

        Args:
            moved_images (Union[BatchedImages, FakeBatchedImages, torch.Tensor]): The moved images to save.
            filenames (Union[str, List[str]]): The filenames to save the moved images to.
            moving_to_fixed (bool, optional): If True, the moving images are saved to the fixed image space. Defaults to True.
                if False, we are dealing with an image that is moved from fixed space to moving space
        '''
        if isinstance(moved_images, BatchedImages):
            moved_images_save = FakeBatchedImages(moved_images(), moved_images, ignore_size_match)   # roundabout way to call the fakebatchedimages
        elif isinstance(moved_images, torch.Tensor):
            moved_images_save = FakeBatchedImages(moved_images, self.fixed_images if moving_to_fixed else self.moving_images, ignore_size_match)
        else:
            # if it is already a fakebatchedimages, we can just use it
            moved_images_save = moved_images
        moved_images_save.write_image(filenames)


    def evaluate_inverse(self, fixed_images: Union[BatchedImages, torch.Tensor], moving_images: Union[BatchedImages, torch.Tensor], shape=None, **kwargs):
        ''' Apply the inverse of the learned transformation to new images.

        This method is useful to analyse the effect of how the moving coordinates (fixed images) are transformed
        '''
        if isinstance(fixed_images, torch.Tensor):
            fixed_images = FakeBatchedImages(fixed_images, self.fixed_images)
        if isinstance(moving_images, torch.Tensor):
            moving_images = FakeBatchedImages(moving_images, self.moving_images)

        fixed_arrays = moving_images()
        fixed_moved_coords = self.get_inverse_warp_parameters(fixed_images, moving_images, shape=shape, **kwargs)
        fixed_moved_image = fireants_interpolator(fixed_arrays, **fixed_moved_coords, mode='bilinear', align_corners=True)  # [N, C, H, W, [D]]
        return fixed_moved_image


    def evaluate(self, fixed_images: Union[BatchedImages, torch.Tensor], moving_images: Union[BatchedImages, torch.Tensor], shape=None):
        '''Apply the learned transformation to new images.

        This method applies the registration transformation learned during optimization
        to a new set of images. It can be used to:
            - Validate registration performance on test images
            - Apply learned transformations to new data
            - Transform auxiliary data (e.g. segmentation masks) using learned parameters

        All registration classes will implement their own `get_warped_coordinates` method,
        which is used to apply the learned transformation to new images.

        Args:
            fixed_images (BatchedImages): Fixed/reference images that define the target space
            moving_images (BatchedImages): Moving images to be transformed
            shape (Optional[Tuple[int, ...]]): Optional output shape for the transformed image.
                If None, uses the shape of the fixed image.

        Returns:
            torch.Tensor: The transformed moving image in the space of the fixed image.
                Has shape [N, C, H, W, [D]] where:
                    N = batch size
                    C = number of channels
                    H, W, D = spatial dimensions (D only for 3D)

        Note:
            The transformation is applied using bilinear interpolation with align_corners=True
            to maintain consistency with the optimization process.
        '''
        if isinstance(fixed_images, torch.Tensor):
            fixed_images = FakeBatchedImages(fixed_images, self.fixed_images)
        if isinstance(moving_images, torch.Tensor):
            moving_images = FakeBatchedImages(moving_images, self.moving_images)

        moving_arrays = moving_images()
        moved_coords = self.get_warp_parameters(fixed_images, moving_images, shape=shape)
        interpolate_mode = moving_images.get_interpolator_type()
        moved_image = fireants_interpolator(moving_arrays, **moved_coords, mode=interpolate_mode, align_corners=True)  # [N, C, H, W, [D]]
        return moved_image

    def evaluate_keypoints(self, fixed_keypoints: BatchedKeypoints, moving_keypoints: BatchedKeypoints):
        '''Apply the learned transformation to keypoints in the fixed image space

        This method applies the registration transformation learned during optimization
        to a set of unordered keypoints in the fixed image space. It can be used to:
            - Validate registration performance on test images 

        All registration classes will implement their own `get_warped_coordinates` method,
        which is used to apply the learned transformation to new images.

        Args:
            fixed_keypoints (BatchedKeypoints): Fixed/reference keypoints that define the target space

        Returns:
            BatchedKeypoints: The transformed keypoints in the space of the fixed image.
        '''
        moved_coords = self.get_warp_parameters(self.fixed_images, self.moving_images)
        kps_tensor = fixed_keypoints.as_torch_coordinates() # tensor if can_collate, list of tensors otherwise
        can_collate = fixed_keypoints.can_collate
        dims = self.dims
        if can_collate:
            # kps_tensor is a [B, N, dims] tensor -> unsqueeze to (B, 1, N, dims) or (B, 1, 1, N, dims)
            kps_tensor_disp = kps_tensor
            for _ in range(dims - 1):
                kps_tensor_disp = kps_tensor_disp.unsqueeze(1)
            # compute displacement first
            final_coord = 0
            if 'grid' in moved_coords:
                # compute u(x_f) and squeeze back
                final_coord = F.grid_sample(moved_coords['grid'].permute(*self.warp.permute_vtoimg), kps_tensor_disp, mode='bilinear', align_corners=True).permute(*self.warp.permute_imgtov)
                for _ in range(dims - 1):
                    final_coord = final_coord.squeeze(1)

            if 'affine' in moved_coords:
                final_coord = final_coord + BatchedKeypoints.transform_keypoints_batch(moved_coords['affine'], kps_tensor)
            else:
                final_coord = final_coord + kps_tensor
            # final_coord is [B, N, dims] tensor
            return BatchedKeypoints.from_tensor_and_metadata(final_coord, moving_keypoints, space='torch')
        else:
            raise NotImplementedError("Not implemented for non-collated keypoints")