Added similarity metric option to stage-2 script.

share/example/config/2_block_to_slice.json

@@ -14,7 +14,8 @@
         "verbose": false,
         "outputdir": "/Users/inhuszar/temp/stage2",
         "stages": ["rigid", "affine", "nonlinear"],
-        "warnings": false
+        "warnings": false,
+        "cost": "MIND"
     },
     "block": {
         "file": "/Users/inhuszar/temp/2_tissue_block/tissue_block.tif",

src/tirl/scripts/mnd/block_to_slice.py

@@ -56,6 +56,7 @@ from tirl.chain import Chain
 from tirl.timage import TField
 from tirl.timage import TImage
 from tirl.costs.mi import CostMI
+from tirl.costs.msd import CostMSD
 from tirl.costs.mind import CostMIND
 from tirl.regularisers.diffusion import DiffusionRegulariser
 from tirl.transformations.linear.scale import TxScale
@@ -338,6 +339,7 @@ def find_best_site(fixed, moving, cnf):
     # Find best site
     costvals = np.stack(costvals, axis=0)
     best_site = int(costvals[np.argmin(costvals[:, -1]), 0])
+    logger.debug(f"Cost values at each site:\n {costvals}")
     logger.info(f"Best site: ({best_site}): {sites[best_site]}")
 
     # Update the site offset and rotation parameters to the optima
@@ -376,10 +378,16 @@ def jiggle(fixed, moving, cnf):
                                    endpoint=False, retstep=True)
 
     # Measure grid cost: find N best location + rotation parameters
-    # cost = CostMSD(moving, fixed, normalise=True)
-    cost = CostMIND(moving, fixed, normalise=True, kernel=MK_FULL, sigma=1,
-                    truncate=1.5)
-    # cost = CostMI(moving, fixed, normalise=False)
+    if p.general.cost == "MSD":
+        cost = CostMSD(moving, fixed, normalise=True)
+    elif p.general.cost == "MI":
+        cost = CostMI(moving, fixed, normalise=True)
+    elif p.general.cost == "MIND":
+        cost = CostMIND(moving, fixed, normalise=True, kernel=MK_FULL,
+                        sigma=1, truncate=1.5)
+    else:
+        raise ValueError("Unsupported cost function.")
+
     gridcost = []
     for y in yy:
         for x in xx:
@@ -407,8 +415,7 @@ def jiggle(fixed, moving, cnf):
              1: (x - xstep[1] / 2, x + xstep[1] / 2)})
         tx_rotation.parameters.unlock()
         tx_site.parameters.unlock()
-        # # Set cost function
-        # cost = CostMI(moving, fixed, maskmode="and", normalise=False)
+
         # Start optimisation
         OptNL(og, cost, method="LN_BOBYQA", visualise=q.visualise,
               xtol_abs=q.xtol_abs, xtol_rel=q.xtol_rel, step=q.opt_step,
@@ -458,9 +465,18 @@ def rigid2d(fixed, moving, cnf):
         tx_scale = fixed_smooth.domain.chain["scale"]
         tx_translation = fixed_smooth.domain.chain["translation"]
         og = OptimisationGroup(tx_rotation, tx_scale, tx_translation)
+
         # Set cost function
-        # cost = CostMI(moving_smooth, fixed_smooth, normalise=True)
-        cost = CostMIND(moving_smooth, fixed_smooth, normalise=False)
+        if p.general.cost == "MSD":
+            cost = CostMSD(moving_smooth, fixed_smooth, normalise=True)
+        elif p.general.cost == "MI":
+            cost = CostMI(moving_smooth, fixed_smooth, normalise=True)
+        elif p.general.cost == "MIND":
+            cost = CostMIND(moving_smooth, fixed_smooth, normalise=True,
+                            kernel=MK_FULL, sigma=1, truncate=1.5)
+        else:
+            raise ValueError("Unsupported cost function.")
+
         # Start optimisation
         OptNL(og, cost, method="LN_BOBYQA", visualise=q.visualise,
               xtol_abs=q.xtol_abs, xtol_rel=q.xtol_rel, step=q.opt_step,
@@ -494,9 +510,18 @@ def affine2d(fixed, moving, cnf):
         moving_smooth = moving.smooth(sm, copy=True)
         # Prepare transformation to optimise
         tx_affine = fixed_smooth.domain.chain["affine"]
+
         # Set cost function
-        # cost = CostMI(moving_smooth, fixed_smooth, normalise=True)
-        cost = CostMIND(moving_smooth, fixed_smooth, normalise=False)
+        if p.general.cost == "MSD":
+            cost = CostMSD(moving_smooth, fixed_smooth, normalise=True)
+        elif p.general.cost == "MI":
+            cost = CostMI(moving_smooth, fixed_smooth, normalise=True)
+        elif p.general.cost == "MIND":
+            cost = CostMIND(moving_smooth, fixed_smooth, normalise=True,
+                            kernel=MK_FULL, sigma=1, truncate=1.5)
+        else:
+            raise ValueError("Unsupported cost function.")
+
         # Start optimisation
         OptNL(tx_affine, cost, method="LN_BOBYQA",
               xtol_rel=q.xtol_rel, xtol_abs=q.xtol_abs,

src/tirl/scripts/mnd/config/2_block_to_slice.json

@@ -14,7 +14,8 @@
         "verbose": false,
         "outputdir": "/home/inhuszar/Desktop/example/stage2",
         "stages": ["rigid", "affine", "nonlinear"],
-        "warnings": false
+        "warnings": false,
+        "cost": "MIND"
     },
     "block": {
         "file": "/home/inhuszar/Desktop/example/2_tissue_block/tissue_block.tif",