hausdorff_distance() throws IndexError for certain input shapes

[pyscorcher]

🐛 Bug

For certain shapes of inputs the torchmetrics.functional.segmentation.hausdorff_distance() throws an index error. It works e.g. well for 10x10 and 10x11 images, but not for 11x10 images.

To Reproduce

Code sample

import torchmetrics.functional.segmentation
import torch
import sys
print(torch.__version__)
print(torchmetrics.__version__)
print(sys.version)

def main(i):
    if i == 0:
        shape = (1, 1, 10, 10) # works
    elif i == 1:
        shape = (1, 1, 10, 11) # works
    else:
        shape = (1, 1, 11, 10) # throws error
    a = torch.randint(0, 2, size=shape)
    b = torch.randint(0, 2, size=shape)
    return torchmetrics.functional.segmentation.hausdorff_distance(a, b, num_classes=1)

print(f'{main(0)=}')
print(f'{main(1)=}')
print(f'{main(2)=}')

Error

  File "/home/user/script.py", line 21, in <module>
    print(f'{main(2)=}')
             ~~~~^^^
  File "/home/user/script.py", line 17, in main
    return torchmetrics.functional.segmentation.hausdorff_distance(a, b, num_classes=1)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/general/lib/python3.13/site-packages/torchmetrics/functional/segmentation/hausdorff_distance.py", line 99, in hausdorff_distance
    dist = edge_surface_distance(
        preds=preds[b, c],
    ...<3 lines>...
        symmetric=not directed,
    )
  File "/home/user/miniforge3/envs/general/lib/python3.13/site-packages/torchmetrics/functional/segmentation/utils.py", line 499, in edge_surface_distance
    surface_distance(edges_preds, edges_target, distance_metric=distance_metric, spacing=spacing),
    ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/general/lib/python3.13/site-packages/torchmetrics/functional/segmentation/utils.py", line 469, in surface_distance
    dis = distance_transform(~target, sampling=spacing, metric=distance_metric)
  File "/home/user/miniforge3/envs/general/lib/python3.13/site-packages/torchmetrics/functional/segmentation/utils.py", line 351, in distance_transform
    z[i1 * h + j1] = mindis
    ~^^^^^^^^^^^^^
IndexError: index 156 is out of bounds for dimension 0 with size 156

shell returned 1

Environment

• TorchMetrics version: 1.8.2
• Python & PyTorch Version: 3.13.3 (conda-forge), 2.6.0
• Any other relevant information such as OS (e.g., Linux): Ubuntu

[VijayVignesh1]

Hi @pyscorcher
I just confirmed the bug. It basically happens in line 351 in /torchmetrics/functional/segmentation/utils.py.
The correct calculation should be z[i1 * w + j1] = mindis instead of z[i1 * h + j1] = mindis where h, w = x.shape.
So the function will work without any errors as long as the height <= width in the input shape (samples, num_classes, height, width). The calculation will however be wrong since the inherent mapping is incorrect.
@justusschock can you confirm this?

[justusschock]

yes, it looks like this is the source. @pyscorcher would you be willing to submit a fix?

[pyscorcher]

Thank you for looking into it and reproducing it. At the moment I don't feel confident submitting a fix as I don't completely understand the underlying implementation yet, so I'm definitely not gonna be mad if someone more knowledgeable is gonna do it!

Maybe to understand it better, a follow up question: Is it correct that the implementation with the pytorch engine (line 331 and following) is hardcoded for 2D images only? (As opposed to the scipy.ndimage implementation, that seems to work for any number of dimensions?) If this is the case, is this documented somewhere or would it maybe make sense to implement an algorithm that can deal with arbitrary number of dimensions also for the pytorch engine to make the behaviour consistent?

[VijayVignesh1]

Yes, the distance_transform function is right now hardcoded to 2D images. I don't see this in the documentation anywhere.
I guess, extending it to multiple dimensions is a good feature request.