Build template normalize weights

ants/core/__init__.py

@@ -6,7 +6,8 @@
                             make_image,
                             from_numpy,
                             from_numpy_like,
-                            new_image_like)
+                            new_image_like,
+                            ones_like)
 from .ants_image import (ANTsImage,
                          copy_image_info,
                          set_origin,

ants/core/ants_image_io.py

@@ -527,4 +527,25 @@ def new_image_like(image, data):
         has_components=image.has_components)
 
 def from_numpy_like(data, image):
-    return new_image_like(image, data)
+    return new_image_like(image, data)
+
+
+def ones_like(image):
+    """
+    Return an image of ones with the same shape and info as a given image.
+
+    Arguments
+    ---------
+    image : ANTsImage
+        Image to get image shape and info from.
+
+    Returns
+    -------
+    ANTsImage
+        Image of ones with reference header information
+    """
+    ones = image.clone()
+    ones.view().fill(1)
+    return ones
+
+

ants/registration/build_template.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import os
+import shutil
 from tempfile import mktemp
 
 import ants
@@ -14,6 +15,7 @@ def build_template(
     blending_weight=0.75,
     weights=None,
     useNoRigid=True,
+    output_dir=None,
     **kwargs
 ):
     """
@@ -44,6 +46,10 @@ def build_template(
     useNoRigid : boolean
         equivalent of -y in the script. Template update
         step will not use the rigid component if this is True.
+
+    output_dir : path
+        directory name where intermediate transforms are written
+
     kwargs : keyword args
         extra arguments passed to ants registration
 
@@ -60,6 +66,12 @@ def build_template(
     >>> timage = ants.build_template( image_list = ( image, image2, image3 ) ).resample_image( (45,45))
     >>> timagew = ants.build_template( image_list = ( image, image2, image3 ), weights = (5,1,1) )
     """
+    work_dir = mktemp() if output_dir is None else output_dir
+
+    def make_outprefix(k: int):
+        os.makedirs(os.path.join(work_dir, f"img{k:04d}"), exist_ok=True)
+        return os.path.join(work_dir, f"img{k:04d}", "out")
+
     if "type_of_transform" not in kwargs:
         type_of_transform = "SyN"
     else:
@@ -70,17 +82,22 @@ def build_template(
     weights = [x / sum(weights) for x in weights]
     if initial_template is None:
         initial_template = image_list[0] * 0
+        wimg = initial_template.clone("float")
         for i in range(len(image_list)):
             temp = image_list[i] * weights[i]
             temp = ants.resample_image_to_target(temp, initial_template)
             initial_template = initial_template + temp
+            wtemp = ants.resample_image_to_target(ants.ones_like(image_list[i]), initial_template)
+            wimg = wimg + wtemp * weights[i]
+        nonzero = wimg.view() != 0
+        initial_template.view()[nonzero] = initial_template.view()[nonzero] / wimg.view()[nonzero]
 
     xavg = initial_template.clone()
     for i in range(iterations):
         affinelist = []
         for k in range(len(image_list)):
             w1 = ants.registration(
-                xavg, image_list[k], type_of_transform=type_of_transform, **kwargs
+                xavg, image_list[k], type_of_transform=type_of_transform, outprefix=make_outprefix(k), **kwargs
             )
             L = len(w1["fwdtransforms"])
             # affine is the last one
@@ -90,10 +107,16 @@ def build_template(
                 if L == 2:
                     wavg = ants.image_read(w1["fwdtransforms"][0]) * weights[k]
                 xavgNew = w1["warpedmovout"] * weights[k]
+                wimgNew = ants.apply_transforms(xavg, ants.ones_like(image_list[k]) * weights[k], transformlist=w1["fwdtransforms"])
             else:
                 if L == 2:
                     wavg = wavg + ants.image_read(w1["fwdtransforms"][0]) * weights[k]
                 xavgNew = xavgNew + w1["warpedmovout"] * weights[k]
+                wimgNew = wimgNew + ants.apply_transforms(xavg, ants.ones_like(image_list[k]) * weights[k], transformlist=w1["fwdtransforms"])
+
+        # normalize
+        nonzero = wimgNew.view() != 0
+        xavgNew.view()[nonzero] = xavgNew.view()[nonzero] / wimgNew.view()[nonzero]
 
         if useNoRigid:
             avgaffine = ants.average_affine_transform_no_rigid(affinelist)
@@ -112,6 +135,7 @@ def build_template(
             wavgfn = mktemp(suffix=".nii.gz")
             ants.image_write(wavgA, wavgfn)
             xavg = ants.apply_transforms(fixed=xavgNew, moving=xavgNew, transformlist=[wavgfn, afffn], whichtoinvert=[0, 1])
+            os.remove(wavgfn)
         else:
             xavg = ants.apply_transforms(fixed=xavgNew, moving=xavgNew, transformlist=[afffn], whichtoinvert=[1])
 
@@ -121,4 +145,6 @@ def build_template(
                 1.0 - blending_weight
             )
 
+    if output_dir is None:
+        shutil.rmtree(work_dir)
     return xavg