WiNet: Wavelet-based Incremental Learning for Efficient Medical Image Registration

TL;DR

WiNet introduces a wavelet-based incremental learning approach for medical image registration that improves efficiency and explainability by estimating scale-wise wavelet coefficients without cascaded architectures.

Contribution

The paper presents WiNet, a novel wavelet-based model that incrementally estimates deformation fields, reducing memory usage and increasing explainability compared to traditional cascaded methods.

Findings

Achieves high accuracy on 3D medical datasets.

Demonstrates GPU efficiency and reduced memory consumption.

Provides an explainable wavelet-based registration framework.

Abstract

Deep image registration has demonstrated exceptional accuracy and fast inference. Recent advances have adopted either multiple cascades or pyramid architectures to estimate dense deformation fields in a coarse-to-fine manner. However, due to the cascaded nature and repeated composition/warping operations on feature maps, these methods negatively increase memory usage during training and testing. Moreover, such approaches lack explicit constraints on the learning process of small deformations at different scales, thus lacking explainability. In this study, we introduce a model-driven WiNet that incrementally estimates scale-wise wavelet coefficients for the displacement/velocity field across various scales, utilizing the wavelet coefficients derived from the original input image pair. By exploiting the properties of the wavelet transform, these estimated coefficients facilitate the seamless reconstruction of a full-resolution displacement/velocity field via our devised inverse discrete wavelet transform (IDWT) layer. This approach avoids the complexities of cascading networks or composition operations, making our WiNet an explainable and efficient competitor with other coarse-to-fine methods. Extensive experimental results from two 3D datasets show that our WiNet is accurate and GPU efficient. The code is available at https://github.com/x-xc/WiNet .