Add SSIM for 3D tensors

[lrlunin]

Fixes #3344

Description:
Implements the SSIM for the 3-dimensional structures.

I was forced to reduce the test data dimensions in a way that the test_ssim_3d does not consume more memory than the regular 2D test_ssim version.

Check list:

• New tests are added (if a new feature is added)
• New doc strings: description and/or example code are in RST format
• Documentation is updated (if required)

[vfdev-5]

Thanks for the PR @lrlunin! Is it possible to extend exisitng SSIM class to cover 3D case instead of creating SSIM3D class. I have an impression that SSIM3D is mostly a copy of SSIM...
For example, we can think to add an argument:

class SSIM(Metric):
    def __init__(
        self,
        data_range: Union[int, float],
        kernel_size: Union[int, Sequence[int]] = 11,
        sigma: Union[float, Sequence[float]] = 1.5,
        k1: float = 0.01,
        k2: float = 0.03,
        gaussian: bool = True,
        output_transform: Callable = lambda x: x,
        device: Union[str, torch.device] = torch.device("cpu"),
        skip_unrolling: bool = False,
        ndims: int = 2,
    ):
        ...
        if ndims not in (2, 3):
            raise ValueError(...) 

What do you think?

[lrlunin]

Thanks for the PR @lrlunin! Is it possible to extend exisitng SSIM class to cover 3D case instead of creating SSIM3D class. I have an impression that SSIM3D is mostly a copy of SSIM... For example, we can think to add an argument:

class SSIM(Metric):
    def __init__(
        self,
        data_range: Union[int, float],
        kernel_size: Union[int, Sequence[int]] = 11,
        sigma: Union[float, Sequence[float]] = 1.5,
        k1: float = 0.01,
        k2: float = 0.03,
        gaussian: bool = True,
        output_transform: Callable = lambda x: x,
        device: Union[str, torch.device] = torch.device("cpu"),
        skip_unrolling: bool = False,
        ndims: int = 2,
    ):
        ...
        if ndims not in (2, 3):
            raise ValueError(...) 

What do you think?

Yeah, this would work. I just wasn't sure if you agree to change the original class. I'll rework this version towards your tipp and mention you after it's done.

[vfdev-5]

Thanks for the updates.
I left few comments.

[vfdev-5]

@lrlunin please fix mypy errors:

Run bash ./tests/run_code_style.sh mypy
+ '[' mypy = lint ']'
+ '[' mypy = fmt ']'
+ '[' mypy = mypy ']'
+ mypy --config-file mypy.ini
ignite/metrics/ssim.py:[15](https://github.com/pytorch/ignite/actions/runs/13957726772/job/39073093453?pr=3345#step:12:16)6: error: Missing return statement  [return]
        def _gaussian_or_uniform_kernel(
        ^
ignite/metrics/ssim.py:224: error: List item 0 has incompatible type
"Optional[int]"; expected "int"  [list-item]
                padding_shape.extend([self.pad_d, self.pad_d])
                                      ^~~~~~~~~~
ignite/metrics/ssim.py:[22](https://github.com/pytorch/ignite/actions/runs/13957726772/job/39073093453?pr=3345#step:12:23)4: error: List item 1 has incompatible type
"Optional[int]"; expected "int"  [list-item]
                padding_shape.extend([self.pad_d, self.pad_d])
                                                  ^~~~~~~~~~
Found 3 errors in 1 file (checked 154 source files)

• https://github.com/pytorch/ignite/actions/runs/13957726772/job/39073093453?pr=3345

Also please rebase the PR on top of the master branch.

[lrlunin]

@lrlunin please fix mypy errors:

Run bash ./tests/run_code_style.sh mypy
+ '[' mypy = lint ']'
+ '[' mypy = fmt ']'
+ '[' mypy = mypy ']'
+ mypy --config-file mypy.ini
ignite/metrics/ssim.py:[15](https://github.com/pytorch/ignite/actions/runs/13957726772/job/39073093453?pr=3345#step:12:16)6: error: Missing return statement  [return]
        def _gaussian_or_uniform_kernel(
        ^
ignite/metrics/ssim.py:224: error: List item 0 has incompatible type
"Optional[int]"; expected "int"  [list-item]
                padding_shape.extend([self.pad_d, self.pad_d])
                                      ^~~~~~~~~~
ignite/metrics/ssim.py:[22](https://github.com/pytorch/ignite/actions/runs/13957726772/job/39073093453?pr=3345#step:12:23)4: error: List item 1 has incompatible type
"Optional[int]"; expected "int"  [list-item]
                padding_shape.extend([self.pad_d, self.pad_d])
                                                  ^~~~~~~~~~
Found 3 errors in 1 file (checked 154 source files)

• https://github.com/pytorch/ignite/actions/runs/13957726772/job/39073093453?pr=3345

Also please rebase the PR on top of the master branch.

Should we add mypy into pre-commit too? I was wondering that there were no errors.

[vfdev-5]

Ultimately, in another PR I would like to remove run_code_style.sh script and use only pre-commit with ruff, black and mypy.

[vfdev-5]

Finally, maybe 1e-4 for all devices is not that a bad idea...

There is however one thing that bothers me: CI time of mps-tests and cpu-tests (macos-latest).
Previously, they complete in about 5-6 minutes:

• https://github.com/pytorch/ignite/actions/runs/13951395188/job/39051591194
• https://github.com/pytorch/ignite/actions/runs/13951395210/job/39051602359

Now they run during:

• 37 minutes mps-tests: https://github.com/pytorch/ignite/actions/runs/13967385618/job/39100936704?pr=3345
• more than 42 minutes cpu-tests (macos-latest): https://github.com/pytorch/ignite/actions/runs/13967385599/job/39100940860?pr=3345

We should do something with that:

• if you have a mac, can you run some benchmarks on how much it takes to run SSIM 3d tests
• try reduce batch size in the 3d SSIM tests-cases, spatial dims as well, maybe less parametrizations?

[vfdev-5]

@lrlunin thanks for working on the PR!
Your last change before further batch size reduction were giving ~15 minutes for macos CI jobs vs 7 minutes on master.
Let's see if we can be <10 minutes with the latest updates.
Otherwise, the PR looks good to me.

[vfdev-5]

Unfortunately, reduced too much some value:

FAILED tests/ignite/metrics/test_ssim.py::test_ssim_uint8[mps-shape2-7-3-False-True] - ValueError: win_size exceeds image extent. Either ensure that your images are at least 7x7; or pass win_size explicitly in the function call, with an odd value less than or equal to the smaller side of your images. If your images are multichannel (with color channels), set channel_axis to the axis number corresponding to the channels.
FAILED tests/ignite/metrics/test_ssim.py::test_ssim_uint8[mps-shape3-11-3-True-False] - ValueError: win_size exceeds image extent. Either ensure that your images are at least 7x7; or pass win_size explicitly in the function call, with an odd value less than or equal to the smaller side of your images. If your images are multichannel (with color channels), set channel_axis to the axis number corresponding to the channels.

Can you check this please ?

[vfdev-5]

Hey @lrlunin, thanks for working on this PR! Is it possible to check the issue: #3345 (comment) and fix it.
I plan to make the next release (v0.5.2) this week/next week and would like to have this metric there.

[lrlunin]

Hey @lrlunin, thanks for working on this PR! Is it possible to check the issue: #3345 (comment) and fix it. I plan to make the next release (v0.5.2) this week/next week and would like to have this metric there.

Sorry, I was busy and ill. I increased the minimal dimensions.

[vfdev-5]

No worries, hope you fill better now. Thanks for the updates!
I may push some changes myself to this PR to satisfy the running time limitation (try to get macosx CI running < 10 minutes)

[vfdev-5]

LGTM, thanks @lrlunin !