ENH: Add labelmap support to image registration pipeline

.gitignore

@@ -1,4 +1,5 @@
 # Project files
+*_cache
 .coverage
 *.code-workspace
 coverage.xml

docs/API_MAP.md

@@ -168,20 +168,21 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._
   - `def set_metric(self, metric)` (line 107): Set the similarity metric to use for registration.
   - `def itk_affine_transform_to_ants_transform(self, itk_tfm)` (line 317): Convert ITK affine/rigid transform to ANTs affine transform.
   - `def itk_transform_to_antsfile(self, itk_tfm, reference_image, output_filename)` (line 410): Convert ITK transform to ANTs transform file.
-  - `def registration_method(self, moving_image, moving_mask=None, moving_image_pre=None, initial_forward_transform=None)` (line 510): Register moving image to fixed image using ANTs registration algorithm.
+  - `def registration_method(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 510): Register moving image to fixed image using ANTs registration algorithm.
 
 ## src/physiomotion4d/register_images_base.py
 
 - **class RegisterImagesBase** (line 29): Base class for deformable image registration algorithms.
   - `def __init__(self, log_level=logging.INFO)` (line 75): Initialize the base image registration class.
-  - `def set_modality(self, modality)` (line 106): Set the imaging modality for registration optimization.
-  - `def set_fixed_image(self, fixed_image)` (line 122): Set the fixed/target image for registration.
-  - `def set_mask_dilation(self, mask_dilation_mm)` (line 143): Set the dilation of the fixed and moving image masks.
-  - `def set_fixed_mask(self, fixed_mask)` (line 151): Set a binary mask for the fixed image region of interest.
-  - `def preprocess(self, image, modality='ct')` (line 192): Preprocess the image based on modality-specific requirements.
-  - `def registration_method(self, moving_image, moving_mask=None, moving_image_pre=None, initial_forward_transform=None)` (line 213): Main registration method to align moving image to fixed image.
-  - `def register(self, moving_image, moving_mask=None, moving_image_pre=None, initial_forward_transform=None)` (line 246): Register a moving image to the fixed image.
-  - `def get_registered_image(self)` (line 329): Get the registered image.
+  - `def set_modality(self, modality)` (line 108): Set the imaging modality for registration optimization.
+  - `def set_fixed_image(self, fixed_image)` (line 124): Set the fixed/target image for registration.
+  - `def set_mask_dilation(self, mask_dilation_mm)` (line 145): Set the dilation of the fixed and moving image masks.
+  - `def set_fixed_mask(self, fixed_mask)` (line 153): Set a binary mask for the fixed image region of interest.
+  - `def set_fixed_labelmap(self, fixed_labelmap)` (line 194): Set the fixed image labelmap (multi-label segmentation).
+  - `def preprocess(self, image, modality='ct')` (line 207): Preprocess the image based on modality-specific requirements.
+  - `def registration_method(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 228): Main registration method to align moving image to fixed image.
+  - `def register(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 263): Register a moving image to the fixed image.
+  - `def get_registered_image(self)` (line 350): Get the registered image.
 
 ## src/physiomotion4d/register_images_greedy.py
 
@@ -190,7 +191,7 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._
   - `def set_number_of_iterations(self, number_of_iterations)` (line 80): Set the number of iterations per resolution level.
   - `def set_transform_type(self, transform_type)` (line 88): Set the type of transform: Deformable, Affine, or Rigid.
   - `def set_metric(self, metric)` (line 99): Set the similarity metric (CC→NCC, Mattes→NMI, MeanSquares→SSD).
-  - `def registration_method(self, moving_image, moving_mask=None, moving_image_pre=None, initial_forward_transform=None)` (line 260): Register moving image to fixed image using Greedy.
+  - `def registration_method(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 260): Register moving image to fixed image using Greedy.
 
 ## src/physiomotion4d/register_images_icon.py
 
@@ -201,7 +202,7 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._
   - `def set_multi_modality(self, enable)` (line 101): Enable or disable multi-modality registration.
   - `def set_mass_preservation(self, enable)` (line 118): Enable or disable mass preservation constraint.
   - `def preprocess(self, image, modality='ct')` (line 135): Preprocess the image for ICON registration.
-  - `def registration_method(self, moving_image, moving_mask=None, moving_image_pre=None, initial_forward_transform=None)` (line 155): Register moving image to fixed image using ICON registration algorithm.
+  - `def registration_method(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 155): Register moving image to fixed image using ICON registration algorithm.
 
 ## src/physiomotion4d/register_models_distance_maps.py
 
@@ -248,9 +249,9 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._
   - `def set_modality(self, modality)` (line 149): Set the imaging modality for registration optimization.
   - `def set_fixed_image(self, fixed_image)` (line 159): Set the fixed image for registration.
   - `def set_fixed_mask(self, fixed_mask)` (line 170): Set a binary mask for the fixed image region of interest.
-  - `def register_time_series(self, moving_images, moving_masks=None, reference_frame=0, register_reference=True, prior_weight=0.0)` (line 180): Register a time series of images to the fixed image.
-  - `def reconstruct_time_series(self, moving_images, inverse_transforms, upsample_to_fixed_resolution=False)` (line 500): Reconstruct time series images using inverse transforms.
-  - `def registration_method(self, moving_image, moving_mask=None, moving_image_pre=None, initial_forward_transform=None)` (line 630): Registration method required by RegisterImagesBase.
+  - `def register_time_series(self, moving_images, moving_masks=None, moving_labelmaps=None, reference_frame=0, register_reference=True, prior_weight=0.0)` (line 180): Register a time series of images to the fixed image.
+  - `def reconstruct_time_series(self, moving_images, inverse_transforms, upsample_to_fixed_resolution=False)` (line 534): Reconstruct time series images using inverse transforms.
+  - `def registration_method(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 664): Registration method required by RegisterImagesBase.
 
 ## src/physiomotion4d/segment_anatomy_base.py
 

src/physiomotion4d/register_images_ants.py

@@ -511,6 +511,7 @@ def registration_method(
         self,
         moving_image: itk.Image,
         moving_mask: Optional[itk.Image] = None,
+        moving_labelmap: Optional[itk.Image] = None,
         moving_image_pre: Optional[ants.ANTsImage] = None,
         initial_forward_transform: Optional[itk.Transform] = None,
     ) -> dict[str, Union[itk.Transform, float]]:

src/physiomotion4d/register_images_base.py

@@ -16,7 +16,7 @@
 """
 
 import logging
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import itk
 import numpy as np
@@ -91,10 +91,12 @@ def __init__(self, log_level: int | str = logging.INFO) -> None:
         self.fixed_image: Optional[itk.Image] = None
         self.fixed_image_pre: Optional[itk.Image] = None
         self.fixed_mask: Optional[itk.Image] = None
+        self.fixed_labelmap: Optional[itk.Image] = None
 
         self.moving_image: Optional[itk.Image] = None
         self.moving_image_pre: Optional[itk.Image] = None
         self.moving_mask: Optional[itk.Image] = None
+        self.moving_labelmap: Optional[itk.Image] = None
 
         self.mask_dilation_mm: float = 5.0
 
@@ -189,6 +191,19 @@ def set_fixed_mask(self, fixed_mask: Optional[itk.Image]) -> None:
             )
             self.fixed_mask = imMath.GetOutputUChar()
 
+    def set_fixed_labelmap(self, fixed_labelmap: Optional[itk.Image]) -> None:
+        """Set the fixed image labelmap (multi-label segmentation).
+
+        Args:
+            fixed_labelmap (itk.Image, optional): Multi-label segmentation
+                co-registered with the fixed image, or None to clear.
+        """
+        self.fixed_labelmap = fixed_labelmap
+        self.forward_transform = None
+        self.inverse_transform = None
+        self.loss = None
+        self.moving_image_registered = None
+
     def preprocess(self, image: itk.Image, modality: str = "ct") -> itk.Image:
         """Preprocess the image based on modality-specific requirements.
 
@@ -214,9 +229,10 @@ def registration_method(
         self,
         moving_image: itk.Image,
         moving_mask: Optional[itk.Image] = None,
+        moving_labelmap: Optional[itk.Image] = None,
         moving_image_pre: Optional[itk.Image] = None,
         initial_forward_transform: Optional[itk.Transform] = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, Union[itk.Transform, float]]:
         """Main registration method to align moving image to fixed image.
 
         This method serves as the primary interface for performing image
@@ -229,6 +245,7 @@ def registration_method(
         Args:
             moving_image (itk.image): The 3D image to be registered to the fixed image
             moving_mask (itk.image, optional): Binary mask for moving image ROI
+            moving_labelmap (itk.image, optional): Multi-label segmentation for the moving image
             moving_image_pre (itk.image, optional): Preprocessed moving image
             initial_forward_transform (itk.Transform, optional): Initial transformation from moving to fixed
 
@@ -247,9 +264,10 @@ def register(
         self,
         moving_image: itk.Image,
         moving_mask: Optional[itk.Image] = None,
+        moving_labelmap: Optional[itk.Image] = None,
         moving_image_pre: Optional[itk.Image] = None,
         initial_forward_transform: Optional[itk.Transform] = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, Union[itk.Transform, float]]:
         """Register a moving image to the fixed image.
 
         This is the main registration method that must be implemented by
@@ -259,6 +277,7 @@ def register(
         Args:
             moving_image (itk.image): The 3D image to be registered to the fixed image
             moving_mask (itk.image, optional): Binary mask for moving image ROI
+            moving_labelmap (itk.image, optional): Multi-label segmentation for the moving image
             moving_image_pre (itk.image, optional): Preprocessed moving image
             initial_forward_transform (itk.Transform, optional): Initial transformation from moving to fixed
 
@@ -308,10 +327,12 @@ def register(
         self.moving_image = moving_image
         self.moving_image_pre = moving_image_pre
         self.moving_mask = new_moving_mask
+        self.moving_labelmap = moving_labelmap
 
         result = self.registration_method(
             moving_image,
             moving_mask=new_moving_mask,
+            moving_labelmap=moving_labelmap,
             moving_image_pre=moving_image_pre,
             initial_forward_transform=initial_forward_transform,
         )

src/physiomotion4d/register_images_greedy.py

@@ -261,6 +261,7 @@ def registration_method(
         self,
         moving_image: itk.Image,
         moving_mask: Optional[itk.Image] = None,
+        moving_labelmap: Optional[itk.Image] = None,
         moving_image_pre: Optional[itk.Image] = None,
         initial_forward_transform: Optional[itk.Transform] = None,
     ) -> dict[str, Union[itk.Transform, float]]:

src/physiomotion4d/register_images_icon.py

@@ -10,7 +10,7 @@
 """
 
 import logging
-from typing import Optional
+from typing import Optional, Union
 
 import icon_registration as icon
 import icon_registration.itk_wrapper
@@ -156,9 +156,10 @@ def registration_method(
         self,
         moving_image: itk.Image,
         moving_mask: Optional[itk.Image] = None,
+        moving_labelmap: Optional[itk.Image] = None,
         moving_image_pre: Optional[itk.Image] = None,
         initial_forward_transform: Optional[itk.Transform] = None,
-    ) -> dict[str, object]:
+    ) -> dict[str, Union[itk.Transform, float]]:
         """Register moving image to fixed image using ICON registration algorithm.
 
         Implementation of the abstract register() method from RegisterImagesBase.
@@ -220,32 +221,40 @@ def registration_method(
                 self.fixed_image,
             )
 
+        # Prefer labelmap over binary mask when both sides have a labelmap.
+        use_labelmaps = moving_labelmap is not None and self.fixed_labelmap is not None
+        moving_effective_mask = moving_labelmap if use_labelmaps else moving_mask
+        fixed_effective_mask = self.fixed_labelmap if use_labelmaps else self.fixed_mask
+
         if self.net is None:
+            dice_loss_weight = 1.0 if use_labelmaps else 0.0
             if self.use_multi_modality:
                 self.net = get_multigradicon(
                     loss_fn=icon.LNCC(sigma=5),
                     # loss_fn=icon.losses.MINDSSC(radius=2, dilation=2),
                     apply_intensity_conservation_loss=self.use_mass_preservation,
                     weights_location=self.weights_path,
+                    dice_loss_weight=dice_loss_weight,
                 )
             else:
                 self.net = get_unigradicon(
                     loss_fn=icon.LNCC(sigma=5),
                     apply_intensity_conservation_loss=self.use_mass_preservation,
                     weights_location=self.weights_path,
+                    dice_loss_weight=dice_loss_weight,
                 )
 
         inverse_transform = None
         forward_transform = None
         loss_artifacts = None
-        if self.fixed_mask is not None and moving_mask is not None:
+        if fixed_effective_mask is not None and moving_effective_mask is not None:
             inverse_transform, forward_transform, loss_artifacts = (
                 icon_registration.itk_wrapper.register_pair_with_mask(
                     self.net,
                     self.fixed_image_pre,
                     new_moving_image_pre,
-                    self.fixed_mask,
-                    moving_mask,
+                    fixed_effective_mask,
+                    moving_effective_mask,
                     finetune_steps=self.number_of_iterations,
                     return_artifacts=True,
                 )