Metrics for registration of labelmaps

With the availability of machine learning it is now often possible to obtain reasonable guesses for segmentations. However, ML-based inference usually does not provide topological/connectivity guarantees for the resulting segmentation.

On the other hand, registration has often been very difficult to make robust - especially the initialization of the transform often fails (in my experience). I believe, registration of labelmaps can be much more robust as a strategy to initialize the registration (maybe even close to the final result).

Currently, there are few metrics in ITK to perform labelmap to labelmap registration. The MatchCardinalityImageToImageMetric has never been ported to the v4 registration framework. The mean square metric is suboptimal for labelmaps with many labels (difference between labels 1 and 2 vs difference between labels 1 and 500). An alternative to the MatchCardinalityImageToImageMetric could be to use the mean square metric on a one-hot encoded labelmap, i.e. a vector image.

I attempted to code a MatchCardinalityImageToImageMetricv4, i.e. a v4 version, but stopped because I am unsure what to do about the moving image gradient (I based it on the mean square metric). Maybe the metric should one-hot encode the labels on the fly, and compute the gradient of the one-hot encoded vector pixels?

I also looked into adding vector image support to the ImageRegistrationMethod in SimpleITK (which is based on the v4 registration framework), but realized the VectorToImageMetricTraits currently is templated over the number of components (or “channels”). This is too restrictive, especially for exporting/wrapping to Python, since the main use-case would be many labels, e.g. up to 30 or more.

If anybody has some thoughts or recommendations that would be great! I think this feature would also be interesting for the MONAI community.

More thoughts (originally I posted this on github):