"shape_avg: true" causes squished brain template

[nimzodisaster]

On the left, when I leave shape_avg off, on the right, shape_avg on (becomes squished).

Building a template:
My yaml:
`debug: True
verbose: True

input params

image_list_file: /mnt/rohit_data2/neurite-OASIS/paths.txt
image_prefix: null
image_suffix: null
num_subjects: 8 # set to null to use all subjects

template params

init_template_path: /mnt/rohit_data2/neurite-OASIS/OASIS_OAS1_0270_MR1/aligned_norm.nii.gz

init_template_path: null

template_iterations: 6
normalize_images: True
save_as_uint8: False

laplacian params

num_laplacian: 2
laplace_params:
learning_rate: 0.5
itk_scale: true

shape averaging

shape_avg: true

save outputs

save_dir: ./saved_templates
save_every: 1 # number of epochs to save the template (last epoch will always be saved)
save_init_template: true
save_moved_images: false # set to true if you want to save the moved images at last step (used for template creation)

no num_subjects: will be the same as the number of subjects used for creation

save_additional:
image_file: []
image_prefix: []
image_suffix: []
is_segmentation: []

registration params

batch_size: 8

do_moments: False
moments:
scale: 4
moments: 1

rigid transform

do_rigid: False
rigid:
iterations: [200, 100, 25]
scales: [4, 2, 1]
progress_bar: ${verbose}
loss_type: cc

affine transform

do_affine: False
affine:
iterations: [200, 100, 25]
scales: [4, 2, 1]
progress_bar: ${verbose}
loss_type: cc

deformation transform

do_deform: True
deform_algo: greedy
deform:
iterations: [200, 100, 50]
scales: [4, 2, 1]
optimizer_lr: 0.25 # for template creation, use lower learning rate than pairwise
smooth_warp_sigma: 0.5
smooth_grad_sigma: 1.0
progress_bar: ${verbose}
cc_kernel_size: 5
loss_type: fusedcc

additional params

tmpdir: null
last_reg: null`

[nimzodisaster]

I thout that maybe it was expecting to get a displacement field, but was getting a coordinate grid....but changing the call to fireants interpolator did not work out well.

[nimzodisaster]

What I think is happening:
When shape averaging is enabled, the template becomes severely distorted because the averaging code mixes voxel-space displacement fields with the normalized displacement fields required by the FireANTs interpolator. The quantity accumulated in avg_warp represents full coordinate fields rather than displacements, as far as I can tell. Later, shape_averaging_invwarp takes these coordinate fields and produces a displacement field in voxel units through compositive_warp_inverse(..., displacement=True). However, the 3D interpolation backend always interprets its grid argument as a displacement field expressed in normalized coordinates ranging from –1 to 1, scaled relative to the image dimensions. Passing voxel-scale displacements directly into the sampler results in extreme geometric compression and warping, which explains the squishing observed in the output templates.

To correct this behavior, I guessed that the displacement returned by shape_averaging_invwarp must be converted from voxel units into normalized units before it is passed to fireants_interpolator. Normalization should scale each displacement component by the corresponding factor 2/(size-1) for the appropriate spatial dimension. Once this conversion is applied, the interpolator behaves correctly and the averaged template geometry is preserved.
`
if avg_warp is not None:
# avg_warp is an average of FULL COORDINATE FIELDS (not displacements)
avg_warp = avg_warp / total_dataset_size
parallel_state.all_reduce_across_dp_ranks(avg_warp, op=torch.distributed.ReduceOp.SUM)

take inverse: returns voxel-space displacement field

init_template_batch.broadcast(1)
inverse_disp_vox = shape_averaging_invwarp(init_template_batch, avg_warp)

# convert voxel displacement → normalized displacement in [-1,1]
_, H, W, D, _ = inverse_disp_vox.shape
norm_factors = torch.tensor([
    2.0 / (W - 1),    # x normalized by width
    2.0 / (H - 1),    # y normalized by height
    2.0 / (D - 1),    # z normalized by depth
], device=inverse_disp_vox.device)

inverse_disp_norm = inverse_disp_vox * norm_factors.view(1, 1, 1, 1, 3)

# laplacian smoothing BEFORE applying the averaged warp
updated_template_arr = laplace(updated_template_arr, itk_scale=True, learning_rate=1)

# apply normalized displacement field
updated_template_arr = fireants_interpolator(
    input=updated_template_arr,
    affine=None,
    grid=inverse_disp_norm,
    mode='bilinear',
    align_corners=True,
    is_displacement=True,   # must remain True
)

`

[rohitrango]

When shape averaging is enabled, the template becomes severely distorted because the averaging code mixes voxel-space displacement fields with the normalized displacement fields required by the FireANTs interpolator.

This is not correct. FireANTs uses PyTorch-space displacement fields everywhere only except calls to save_as_ants_transforms and save_as_scipy_transforms (to be added soon).

I'm not able to reproduce the behavior with your config. Can you share a minimal set of images that I can use to reproduce this behavior?