pytorch_msssim_basic_ssim_and_ms_ssim_functional_module_usage.py

1import torch
2from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
3
4# Create two random images (batch_size, channels, height, width)
5img1 = torch.rand(1, 3, 256, 256)
6img2 = torch.rand(1, 3, 256, 256)
7
8# Functional interface
9ssim_val = ssim(img1, img2, data_range=1.0, size_average=True) # return a scalar
10ms_ssim_val = ms_ssim(img1, img2, data_range=1.0, size_average=True) # return a scalar
11
12# Module interface (can be used as a loss function)
13ssim_module = SSIM(data_range=1.0, size_average=True, channel=3)
14ms_ssim_module = MS_SSIM(data_range=1.0, size_average=True, channel=3)
15
16ssim_loss = 1 - ssim_module(img1, img2)
17ms_ssim_loss = 1 - ms_ssim_module(img1, img2)
18
19print(f"SSIM: {ssim_val.item()}")
20print(f"MS-SSIM: {ms_ssim_val.item()}")