3D extension for SSIM

CHANGELOG.md

@@ -17,6 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added new image metric `UniversalImageQualityIndex` ([#824](https://github.com/PyTorchLightning/metrics/pull/824))
 
 
+- Added support for 3D image and uniform kernel in `StructuralSimilarityIndexMeasure` ([#818](https://github.com/PyTorchLightning/metrics/pull/818))
+
+
 - Added smart update of `MetricCollection` ([#709](https://github.com/PyTorchLightning/metrics/pull/709))
 
 

tests/image/test_ms_ssim.py

@@ -29,7 +29,7 @@
 BATCH_SIZE = 1
 
 _inputs = []
-for size, coef in [(128, 0.9), (128, 0.7)]:
+for size, coef in [(182, 0.9), (182, 0.7)]:
  preds = torch.rand(NUM_BATCHES, BATCH_SIZE, 1, size, size)
  _inputs.append(
  Input(
@@ -40,40 +40,41 @@
 
 
 def pytorch_ms_ssim(preds, target, data_range, kernel_size):
- return ms_ssim(preds, target, data_range=data_range, win_size=kernel_size)
+ return ms_ssim(preds, target, data_range=data_range, win_size=kernel_size, size_average=False)
 
 
 @pytest.mark.parametrize(
  "preds, target",
  [(i.preds, i.target) for i in _inputs],
 )
-@pytest.mark.parametrize("kernel_size", [5, 7])
 class TestMultiScaleStructuralSimilarityIndexMeasure(MetricTester):
  atol = 6e-3
 
+ # in the pytorch-msssim package, sigma is hardcoded to 1.5. We can thus only test this value, which corresponds
+ # to a kernel size of 11
+
  @pytest.mark.parametrize("ddp", [False, True])
  @pytest.mark.parametrize("dist_sync_on_step", [False, True])
- def test_ms_ssim(self, preds, target, kernel_size, ddp, dist_sync_on_step):
+ def test_ms_ssim(self, preds, target, ddp, dist_sync_on_step):
  self.run_class_metric_test(
  ddp,
  preds,
  target,
  MultiScaleStructuralSimilarityIndexMeasure,
- partial(pytorch_ms_ssim, data_range=1.0, kernel_size=kernel_size),
- metric_args={"data_range": 1.0, "kernel_size": (kernel_size, kernel_size)},
+ partial(pytorch_ms_ssim, data_range=1.0, kernel_size=11),
  dist_sync_on_step=dist_sync_on_step,
  )
 
- def test_ms_ssim_functional(self, preds, target, kernel_size):
+ def test_ms_ssim_functional(self, preds, target):
  self.run_functional_metric_test(
  preds,
  target,
  multiscale_structural_similarity_index_measure,
- partial(pytorch_ms_ssim, data_range=1.0, kernel_size=kernel_size),
- metric_args={"data_range": 1.0, "kernel_size": (kernel_size, kernel_size)},
+ partial(pytorch_ms_ssim, data_range=1.0, kernel_size=11),
+ metric_args={"data_range": 1.0, "kernel_size": 11},
  )
 
- def test_ms_ssim_differentiability(self, preds, target, kernel_size):
+ def test_ms_ssim_differentiability(self, preds, target):
  # We need to minimize this example to make the test tractable
  single_beta = (1.0,)
  _preds = preds[:, :, :, :16, :16]
@@ -86,7 +87,7 @@ def test_ms_ssim_differentiability(self, preds, target, kernel_size):
  metric_module=MultiScaleStructuralSimilarityIndexMeasure,
  metric_args={
  "data_range": 1.0,
- "kernel_size": (kernel_size, kernel_size),
+ "kernel_size": 11,
  "betas": single_beta,
  },
  )

tests/image/test_ssim.py

@@ -25,84 +25,122 @@
 
 seed_all(42)
 
-Input = namedtuple("Input", ["preds", "target", "multichannel"])
+Input = namedtuple("Input", ["preds", "target"])
 
 _inputs = []
-for size, channel, coef, multichannel, dtype in [
- (12, 3, 0.9, True, torch.float),
- (13, 1, 0.8, False, torch.float32),
- (14, 1, 0.7, False, torch.double),
- (15, 3, 0.6, True, torch.float64),
+for size, channel, coef, dtype in [
+ (12, 3, 0.9, torch.float),
+ (13, 1, 0.8, torch.float32),
+ (14, 1, 0.7, torch.double),
+ (13, 3, 0.6, torch.float32),
 ]:
- preds = torch.rand(NUM_BATCHES, BATCH_SIZE, channel, size, size, dtype=dtype)
+ preds2d = torch.rand(NUM_BATCHES, BATCH_SIZE, channel, size, size, dtype=dtype)
  _inputs.append(
  Input(
- preds=preds,
- target=preds * coef,
- multichannel=multichannel,
+ preds=preds2d,
+ target=preds2d * coef,
+ )
+ )
+ preds3d = torch.rand(NUM_BATCHES, BATCH_SIZE, channel, size, size, size, dtype=dtype)
+ _inputs.append(
+ Input(
+ preds=preds3d,
+ target=preds3d * coef,
  )
  )
 
 
-def _sk_ssim(preds, target, data_range, multichannel, kernel_size):
- c, h, w = preds.shape[-3:]
- sk_preds = preds.view(-1, c, h, w).permute(0, 2, 3, 1).numpy()
- sk_target = target.view(-1, c, h, w).permute(0, 2, 3, 1).numpy()
- if not multichannel:
- sk_preds = sk_preds[:, :, :, 0]
- sk_target = sk_target[:, :, :, 0]
-
- return structural_similarity(
- sk_target,
- sk_preds,
- data_range=data_range,
- multichannel=multichannel,
- gaussian_weights=True,
- win_size=kernel_size,
- sigma=1.5,
- use_sample_covariance=False,
- )
+def _sk_ssim(preds, target, data_range, sigma, kernel_size=None, return_ssim_image=False, gaussian_weights=True):
+ if len(preds.shape) == 4:
+ c, h, w = preds.shape[-3:]
+ sk_preds = preds.view(-1, c, h, w).permute(0, 2, 3, 1).numpy()
+ sk_target = target.view(-1, c, h, w).permute(0, 2, 3, 1).numpy()
+ elif len(preds.shape) == 5:
+ c, d, h, w = preds.shape[-4:]
+ sk_preds = preds.view(-1, c, d, h, w).permute(0, 2, 3, 4, 1).numpy()
+ sk_target = target.view(-1, c, d, h, w).permute(0, 2, 3, 4, 1).numpy()
+
+ results = torch.zeros(sk_preds.shape[0], dtype=target.dtype)
+ if not return_ssim_image:
+ for i in range(sk_preds.shape[0]):
+ res = structural_similarity(
+ sk_target[i],
+ sk_preds[i],
+ data_range=data_range,
+ multichannel=True,
+ gaussian_weights=gaussian_weights,
+ win_size=kernel_size,
+ sigma=sigma,
+ use_sample_covariance=False,
+ full=return_ssim_image,
+ )
+ results[i] = torch.from_numpy(res).type(preds.dtype)
+ return results
+ else:
+ fullimages = torch.zeros(target.shape, dtype=target.dtype)
+ for i in range(sk_preds.shape[0]):
+ res, fullimage = structural_similarity(
+ sk_target[i],
+ sk_preds[i],
+ data_range=data_range,
+ multichannel=True,
+ gaussian_weights=gaussian_weights,
+ win_size=kernel_size,
+ sigma=sigma,
+ use_sample_covariance=False,
+ full=return_ssim_image,
+ )
+ results[i] = torch.from_numpy(res).type(preds.dtype)
+ fullimage = torch.from_numpy(fullimage).type(preds.dtype)
+ if len(preds.shape) == 4:
+ fullimages[i] = fullimage.permute(2, 0, 1)
+ elif len(preds.shape) == 5:
+ fullimages[i] = fullimage.permute(3, 0, 1, 2)
+ return results, fullimages
 
 
 @pytest.mark.parametrize(
- "preds, target, multichannel",
- [(i.preds, i.target, i.multichannel) for i in _inputs],
+ "preds, target",
+ [(i.preds, i.target) for i in _inputs],
 )
-@pytest.mark.parametrize("kernel_size", [5, 11])
+@pytest.mark.parametrize("sigma", [1.5, 0.5])
 class TestSSIM(MetricTester):
  atol = 6e-3
 
  @pytest.mark.parametrize("ddp", [True, False])
  @pytest.mark.parametrize("dist_sync_on_step", [True, False])
- def test_ssim(self, preds, target, multichannel, kernel_size, ddp, dist_sync_on_step):
+ def test_ssim(self, preds, target, sigma, ddp, dist_sync_on_step):
  self.run_class_metric_test(
  ddp,
  preds,
  target,
  StructuralSimilarityIndexMeasure,
- partial(_sk_ssim, data_range=1.0, multichannel=multichannel, kernel_size=kernel_size),
- metric_args={"data_range": 1.0, "kernel_size": (kernel_size, kernel_size)},
+ partial(_sk_ssim, data_range=1.0, sigma=sigma, kernel_size=None),
+ metric_args={
+ "data_range": 1.0,
+ "sigma": sigma,
+ },
  dist_sync_on_step=dist_sync_on_step,
  )
 
- def test_ssim_functional(self, preds, target, multichannel, kernel_size):
+ def test_ssim_functional(self, preds, target, sigma):
  self.run_functional_metric_test(
  preds,
  target,
  structural_similarity_index_measure,
- partial(_sk_ssim, data_range=1.0, multichannel=multichannel, kernel_size=kernel_size),
- metric_args={"data_range": 1.0, "kernel_size": (kernel_size, kernel_size)},
+ partial(_sk_ssim, data_range=1.0, sigma=sigma, kernel_size=None),
+ metric_args={"data_range": 1.0, "sigma": sigma},
  )
 
  # SSIM half + cpu does not work due to missing support in torch.log
  @pytest.mark.xfail(reason="SSIM metric does not support cpu + half precision")
- def test_ssim_half_cpu(self, preds, target, multichannel, kernel_size):
+ def test_ssim_half_cpu(self, preds, target, sigma):
  self.run_precision_test_cpu(
  preds, target, StructuralSimilarityIndexMeasure, structural_similarity_index_measure, {"data_range": 1.0}
  )
 
  @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires cuda")
- def test_ssim_half_gpu(self, preds, target, multichannel, kernel_size):
+ def test_ssim_half_gpu(self, preds, target, sigma):
  self.run_precision_test_gpu(
  preds, target, StructuralSimilarityIndexMeasure, structural_similarity_index_measure, {"data_range": 1.0}
  )
@@ -111,8 +149,8 @@ def test_ssim_half_gpu(self, preds, target, multichannel, kernel_size):
 @pytest.mark.parametrize(
  ["pred", "target", "kernel", "sigma"],
  [
- ([1, 16, 16], [1, 16, 16], [11, 11], [1.5, 1.5]), # len(shape)
  ([1, 1, 16, 16], [1, 1, 16, 16], [11, 11], [1.5]), # len(kernel), len(sigma)
+ ([1, 16, 16], [1, 16, 16], [11, 11], [1.5, 1.5]), # len(shape)
  ([1, 1, 16, 16], [1, 1, 16, 16], [11], [1.5, 1.5]), # len(kernel), len(sigma)
  ([1, 1, 16, 16], [1, 1, 16, 16], [11], [1.5]), # len(kernel), len(sigma)
  ([1, 1, 16, 16], [1, 1, 16, 16], [11, 0], [1.5, 1.5]), # invalid kernel input
@@ -123,15 +161,15 @@ def test_ssim_half_gpu(self, preds, target, multichannel, kernel_size):
  ],
 )
 def test_ssim_invalid_inputs(pred, target, kernel, sigma):
- pred_t = torch.rand(pred)
+ pred_t = torch.rand(pred, dtype=torch.float32)
  target_t = torch.rand(target, dtype=torch.float64)
  with pytest.raises(TypeError):
  structural_similarity_index_measure(pred_t, target_t)
 
  pred = torch.rand(pred)
  target = torch.rand(target)
  with pytest.raises(ValueError):
- structural_similarity_index_measure(pred, target, kernel, sigma)
+ structural_similarity_index_measure(pred, target, kernel_size=kernel, sigma=sigma)
 
 
 def test_ssim_unequal_kernel_size():
@@ -167,5 +205,20 @@ def test_ssim_unequal_kernel_size():
  ]
  )
  # kernel order matters
- assert structural_similarity_index_measure(preds, target, kernel_size=(3, 5)) == torch.tensor(0.10814697)
- assert structural_similarity_index_measure(preds, target, kernel_size=(5, 3)) != torch.tensor(0.10814697)
+ assert torch.isclose(
+ structural_similarity_index_measure(preds, target, gaussian_kernel=True, sigma=(0.25, 0.5)),
+ torch.tensor(0.08869550),
+ )
+ assert not torch.isclose(
+ structural_similarity_index_measure(preds, target, gaussian_kernel=True, sigma=(0.5, 0.25)),
+ torch.tensor(0.08869550),
+ )
+
+ assert torch.isclose(
+ structural_similarity_index_measure(preds, target, gaussian_kernel=False, kernel_size=(3, 5)),
+ torch.tensor(0.05131844),
+ )
+ assert not torch.isclose(
+ structural_similarity_index_measure(preds, target, gaussian_kernel=False, kernel_size=(5, 3)),
+ torch.tensor(0.05131844),
+ )

torchmetrics/functional/image/helper.py

@@ -1,8 +1,12 @@
-from typing import Sequence
+from typing import Sequence, Union
 
 import torch
+import torch.nn.functional as F
 from torch import Tensor
 
+from torchmetrics.utilities import rank_zero_warn
+from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_10
+
 
 def _gaussian(kernel_size: int, sigma: float, dtype: torch.dtype, device: torch.device) -> Tensor:
  """Computes 1D gaussian kernel.
@@ -22,8 +26,12 @@ def _gaussian(kernel_size: int, sigma: float, dtype: torch.dtype, device: torch.
  return (gauss / gauss.sum()).unsqueeze(dim=0) # (1, kernel_size)
 
 
-def _gaussian_kernel(
- channel: int, kernel_size: Sequence[int], sigma: Sequence[float], dtype: torch.dtype, device: torch.device
+def _gaussian_kernel_2d(
+ channel: int,
+ kernel_size: Sequence[int],
+ sigma: Sequence[float],
+ dtype: torch.dtype,
+ device: Union[torch.device, str],
 ) -> Tensor:
  """Computes 2D gaussian kernel.
 
@@ -35,7 +43,7 @@ def _gaussian_kernel(
  device: device of the output tensor
 
  Example:
- >>> _gaussian_kernel(1, (5,5), (1,1), torch.float, "cpu")
+ >>> _gaussian_kernel_2d(1, (5,5), (1,1), torch.float, "cpu")
  tensor([[[[0.0030, 0.0133, 0.0219, 0.0133, 0.0030],
  [0.0133, 0.0596, 0.0983, 0.0596, 0.0133],
  [0.0219, 0.0983, 0.1621, 0.0983, 0.0219],
@@ -48,3 +56,46 @@ def _gaussian_kernel(
  kernel = torch.matmul(gaussian_kernel_x.t(), gaussian_kernel_y) # (kernel_size, 1) * (1, kernel_size)
 
  return kernel.expand(channel, 1, kernel_size[0], kernel_size[1])
+
+
+def _gaussian_kernel_3d(
+ channel: int, kernel_size: Sequence[int], sigma: Sequence[float], dtype: torch.dtype, device: torch.device
+) -> Tensor:
+ """Computes 3D gaussian kernel.
+
+ Args:
+ channel: number of channels in the image
+ kernel_size: size of the gaussian kernel as a tuple (h, w, d)
+ sigma: Standard deviation of the gaussian kernel
+ dtype: data type of the output tensor
+ device: device of the output tensor
+ """
+
+ gaussian_kernel_x = _gaussian(kernel_size[0], sigma[0], dtype, device)
+ gaussian_kernel_y = _gaussian(kernel_size[1], sigma[1], dtype, device)
+ gaussian_kernel_z = _gaussian(kernel_size[2], sigma[2], dtype, device)
+ kernel_xy = torch.matmul(gaussian_kernel_x.t(), gaussian_kernel_y) # (kernel_size, 1) * (1, kernel_size)
+ kernel = torch.mul(
+ kernel_xy.unsqueeze(-1).repeat(1, 1, kernel_size[2]),
+ gaussian_kernel_z.expand(kernel_size[0], kernel_size[1], kernel_size[2]),
+ )
+ return kernel.expand(channel, 1, kernel_size[0], kernel_size[1], kernel_size[2])
+
+
+def _single_dimension_pad(inputs: Tensor, dim: int, pad: int) -> Tensor:
+ _max = inputs.shape[dim] - 2
+ x = torch.index_select(inputs, dim, torch.arange(pad, 0, -1))
+ y = torch.index_select(inputs, dim, torch.arange(_max, _max - pad, -1))
+ return torch.cat((x, inputs, y), dim)
+
+
+def _reflection_pad_3d(inputs: Tensor, pad_h: int, pad_w: int, pad_d: int) -> Tensor:
+ if _TORCH_GREATER_EQUAL_1_10:
+ inputs = F.pad(inputs, (pad_h, pad_h, pad_w, pad_w, pad_d, pad_d), mode="reflect")
+ else:
+ rank_zero_warn(
+ "An older version of pyTorch is used. For optimal speed, please upgrade to at least pyTorch 1.10."
+ )
+ for dim, pad in enumerate([pad_h, pad_w, pad_d]):
+ inputs = _single_dimension_pad(inputs, dim + 2, pad)
+ return inputs