[ENH] register: option for FOV/COM initialization

Author: balbasty

nitorch/cli/registration/register/cli.py

@@ -97,17 +97,43 @@ def _main(options):
     # ------------------------------------------------------------------
     affine = None
     if options.affine:
+
+        # initialization
+        affine_init = options.affine.init
+        if isinstance(affine_init, str):
+            if affine_init.lower() in ('fov', 'com'):
+                if affine_init.lower() == 'fov':
+                    get_center = 'get_center'
+                else:
+                    get_center = 'get_center_of_mass'
+                loss_list = losses if isinstance(losses, list) else [losses]
+                num = den = 0
+                for loss_sum in loss_list:
+                    for loss in loss_sum:
+                        fix = getattr(loss.fixed, get_center)()
+                        mov = getattr(loss.moving, get_center)()
+                        wgt = 0.5 * (
+                            torch.Size(fix.shape[1:]).numel() +
+                            torch.Size(mov.shape[1:]).numel()
+                        )
+                        num += (mov - fix) * wgt
+                        den += wgt
+                shift = num / den
+                affine_init = torch.eye(4)
+                affine_init[:-1, -1] = shift
+
+        # build affine object
         if options.affine.is2d is not False:
             if isinstance(losses[0], objects.SumSimilarity):
                 affine0 = losses[0][0].fixed.affine
             else:
                 affine0 = losses[0].fixed.affine
             affine = make_affine_2d(options.affine.is2d, affine0,
                                     options.affine.name, options.affine.position,
-                                    init=options.affine.init)
+                                    init=affine_init)
         else:
             affine = make_affine(options.affine.name, options.affine.position,
-                                 init=options.affine.init)
+                                 init=affine_init)
 
     # ------------------------------------------------------------------
     #                           BUILD DENSE

nitorch/cli/registration/register/parser.py

@@ -163,7 +163,7 @@
     Common options:
         -p, --position                  Position of the affine: [sym], mov, fix
         -g, --progressive               Progressive optimization (t -> r -> s -> a) [false]
-        -i, --input                     Path to initial transform
+        -i, --input                     Path to initial transform (fov, com, or file path)
         -o, --output                    Path to the output transform: [{dir}/{name}.lta]
         -2d [AXIS=2]                    Force transform to be 2d about AXIS
 

nitorch/tools/registration/objects.py

@@ -334,6 +334,7 @@ def cpu_(self):
 
     voxel_size = property(lambda self: spatial.voxel_size(self.affine))
     shape = property(lambda self: self.dat.shape[-self.dim:])
+    ndim = property(lambda self: len(self.shape))
     dtype = property(lambda self: self.dat.dtype)
     device = property(lambda self: self.dat.device)
 
@@ -430,13 +431,13 @@ def pull(self, grid, dat=True, mask=False, preview=False):
         if mask:
             msk = None
             if self.masked:
-                msk = self.mask.to(self.dat.dtype)
+                msk = self.mask.to(self.dtype)
                 msk = regutils.smart_pull(msk, grid, bound=self.bound,
                                           extrapolate=self.extrapolate)
             out += [msk]
         if preview:
             if self.previewed:
-                prv = self.preview.to(self.dat.dtype)
+                prv = self.preview.to(self.dtype)
                 prv = regutils.smart_pull(prv, grid, bound=self.bound,
                                           extrapolate=self.extrapolate)
             else:
@@ -481,6 +482,30 @@ def grad(self):
         return spatial.diff(self.dat, dim=list(range(-self.dim, 0)),
                             bound=self.bound)
 
+    def get_center(self):
+        """Compute the RAS coordinate of the center of the field of view."""
+        ndim = self.ndim
+        backend = dict(dtype=self.dtype, device=self.device)
+        affine = self.affine.to(**backend)
+        shape = torch.as_tensor(self.shape, **backend)
+        center = affine[:ndim, :ndim].matmul((shape[:, None] - 1)*0.5)
+        center += affine[:ndim, -1:]
+        return center.squeeze(-1)
+
+    def get_center_of_mass(self, masked=True, **backend):
+        """Compute the RAS coordinate of the center of mass."""
+        backend = dict(dtype=self.dtype, device=self.device)
+        grid = spatial.identity_grid(self.shape, **backend)
+        dat = self.dat
+        if masked and self.masked:
+            msk = self.mask.to(**backend)
+            dat = dat * msk
+            den = msk.sum()
+            if len(msk) == 1:
+                den *= len(dat)
+        num = (dat[..., None] * grid).sum(list(range(dat.ndim)))
+        return num / den
+
     def _prm_as_str(self):
         s = []
         if self.bound != 'dct2':

nitorch/tools/registration/pairwise_makeobj.py

@@ -109,6 +109,10 @@ def make_affine(basis='rigid', position='symmetric', penalty=None, init=None):
         If 'symmetric', both images are rotated by the transformation and
         its inverse, towards a mean space; thereby making the model fully
         symmetric.
+    penalty : None
+        Not implemented
+    init : str | MappedAffine | (4, 4) tensor
+        Initial affine
 
     Returns
     -------