Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added compatibility with uint8 to SSIM metric #3045

Merged
merged 6 commits into from
Sep 13, 2023
Merged

Added compatibility with uint8 to SSIM metric #3045

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
237d8de
feat: add compatibility with uint8
MarcBresson Aug 24, 2023
e04cdbb
style: format using the run_code_style script
MarcBresson Aug 25, 2023
fc52540
Merge branch 'master' into feat-uint8-ssim
vfdev-5 Aug 25, 2023
87d25ea
Merge branch 'master' into feat-uint8-ssim
MarcBresson Sep 13, 2023
67edd84
refactor: delete warning and independantly convert y and y_pred to fp
MarcBresson Sep 13, 2023
2cf4fef
feat: remove uint8 warning test
MarcBresson Sep 13, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
11 changes: 11 additions & 0 deletions ignite/metrics/ssim.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import warnings
from typing import Callable, Sequence, Union

import torch
Expand Down Expand Up @@ -98,6 +99,7 @@ def __init__(

super(SSIM, self).__init__(output_transform=output_transform, device=device)
self.gaussian = gaussian
self.data_range = data_range
self.c1 = (k1 * data_range) ** 2
self.c2 = (k2 * data_range) ** 2
self.pad_h = (self.kernel_size[0] - 1) // 2
Expand Down Expand Up @@ -157,6 +159,15 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
f"Expected y_pred and y to have BxCxHxW shape. Got y_pred: {y_pred.shape} and y: {y.shape}."
)

if y_pred.dtype == torch.uint8 or y.dtype == torch.uint8:
if self.data_range != 255:
warnings.warn(
"dtypes of the input tensors are torch.uint8 but data range is not set to 255.", RuntimeWarning
)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not a big fan of showing a warning in this case. I would follow the same strategy as skimage where they just transform to floating point. So, I would do the following:

if not y.is_floating_point():
    y = y.float()
if not y_pred.is_floating_point():
    y_pred = y_pred.float()

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right, as data_range is a required argument in ignite (while it has a default value in skiimg), the warning is not required.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we should put a default value to the data_range argument too? Most of the time we are dealing with [0; 1] images.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do not have a strong opinion on this. We have to figure out if there wont be any negative impact for the users in this case.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the skiimg approach is quite good: default to data_range = 1.0, and warn if the detected type is of integer and that data range is still set to 1


y_pred = y_pred.to(dtype=torch.float32)
y = y.to(dtype=torch.float32)

channel = y_pred.size(1)
if len(self._kernel.shape) < 4:
self._kernel = self._kernel.expand(channel, 1, -1, -1).to(device=y_pred.device)
Expand Down
46 changes: 46 additions & 0 deletions tests/ignite/metrics/test_ssim.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,52 @@ def test_cuda_ssim_dtypes(available_device, dtype, precision):
test_ssim(available_device, (12, 3, 28, 28), 11, True, False, dtype=dtype, precision=precision)


@pytest.mark.parametrize(
"shape, kernel_size, gaussian, use_sample_covariance",
[[(8, 3, 224, 224), 7, False, True], [(12, 3, 28, 28), 11, True, False]],
)
def test_ssim_uint8(available_device, shape, kernel_size, gaussian, use_sample_covariance):
y_pred = torch.randint(0, 255, shape, device=available_device, dtype=torch.uint8)
y = (y_pred * 0.8).to(dtype=torch.uint8)

sigma = 1.5
data_range = 255
ssim = SSIM(data_range=data_range, sigma=sigma, device=available_device)
ssim.update((y_pred, y))
ignite_ssim = ssim.compute()

skimg_pred = y_pred.cpu().numpy()
skimg_y = (skimg_pred * 0.8).astype(np.uint8)
skimg_ssim = ski_ssim(
skimg_pred,
skimg_y,
win_size=kernel_size,
sigma=sigma,
channel_axis=1,
gaussian_weights=gaussian,
data_range=data_range,
use_sample_covariance=use_sample_covariance,
)

assert isinstance(ignite_ssim, float)
assert np.allclose(ignite_ssim, skimg_ssim, atol=1e-5)


def test_ssim_uint8_warning(available_device):
shape = (7, 3, 256, 256)
y_pred = torch.randint(0, 255, shape, device=available_device, dtype=torch.uint8)
y = (y_pred * 0.8).to(dtype=torch.uint8)

sigma = 1.5
data_range = 1.0
ssim = SSIM(data_range=data_range, sigma=sigma, device=available_device)

with pytest.warns(
RuntimeWarning, match=r"dtypes of the input tensors are torch.uint8 but data range is not set to 255."
):
ssim.update((y_pred, y))


@pytest.mark.parametrize("metric_device", ["cpu", "process_device"])
def test_distrib_integration(distributed, metric_device):
from ignite.engine import Engine
Expand Down