Fast and differentiable MS-SSIM and SSIM for pytorch.
Project description
Pytorch MS-SSIM
Fast and differentiable MS-SSIM and SSIM for pytorch 1.0+
Install
pip install pytorch-msssim
or
python setup.py install
Update
2019.12.10
The negative results or NaN results are caused by neative covariances of input images, which can be avoided by using a larger K2 constant (e.g. 0.4). See 'tests/tests_negative_ssim.py' for more details.
# set K2=0.4 for more stable results
ssim( X, Y, data_range=1, size_average=False, K=(0.01, 0.4))
2019.8.15
Apply to 5D tensor: #6
2019.6.17
Now it is faster than compare_ssim thanks to One-sixth's contribution
Usages
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
# X: (N,3,H,W) a batch of non-negative RGB images (0~255)
# Y: (N,3,H,W)
ssim_val = ssim( X, Y, data_range=255, size_average=False) # return (N,)
ms_ssim_val = ms_ssim( X, Y, data_range=255, size_average=False ) #(N,)
# set 'size_average=True' to get a scalar value as loss.
ssim_loss = 1 - ssim( X, Y, data_range=255, size_average=True) # return a scalar
ms_ssim_loss = 1 - ms_ssim( X, Y, data_range=255, size_average=True )
# reuse windows with SSIM & MS_SSIM.
ssim_module = SSIM(win_size=11, win_sigma=1.5, data_range=255, size_average=True, channel=3)
ms_ssim_module = MS_SSIM(win_size=11, win_sigma=1.5, data_range=255, size_average=True, channel=3)
ssim_loss = 1 - ssim_module(X, Y)
ms_ssim_loss = 1 - ms_ssim_module(X, Y)
Tests
cd tests
1. Compared with skimage.measure.compare_ssim on CPU.
python tests.py
The outputs:
Downloading test image...
====> Single Image
sigma=0.000000 compare_ssim=1.000000 (275.226831 ms) ssim_torch=1.000000 (462.517738 ms)
sigma=10.000000 compare_ssim=0.932497 (389.491558 ms) ssim_torch=0.932494 (63.863516 ms)
sigma=20.000000 compare_ssim=0.785664 (266.695976 ms) ssim_torch=0.785658 (46.617031 ms)
sigma=30.000000 compare_ssim=0.637369 (275.762081 ms) ssim_torch=0.637362 (55.842876 ms)
sigma=40.000000 compare_ssim=0.515707 (236.553907 ms) ssim_torch=0.515700 (45.801163 ms)
sigma=50.000000 compare_ssim=0.422497 (264.705896 ms) ssim_torch=0.422491 (46.895742 ms)
sigma=60.000000 compare_ssim=0.350707 (234.748363 ms) ssim_torch=0.350702 (44.762611 ms)
sigma=70.000000 compare_ssim=0.295998 (210.025072 ms) ssim_torch=0.295993 (45.758247 ms)
sigma=80.000000 compare_ssim=0.253552 (250.259876 ms) ssim_torch=0.253547 (96.461058 ms)
sigma=90.000000 compare_ssim=0.219344 (263.813257 ms) ssim_torch=0.219340 (49.159765 ms)
sigma=100.000000 compare_ssim=0.192421 (258.941889 ms) ssim_torch=0.192418 (47.627449 ms)
Pass
====> Batch
Pass
2. Avoid negative or NaN results
python tests_negative_ssim.py
The outputs:
Negative ssim:
skimage.measure.compare_ssim: -0.96587564223943
pytorch_msssim.ssim: -0.9658759832382202
pytorch_msssim.ms_ssim: nan
Larger K2:
pytorch_msssim.ssim (K2=0.4): 0.0062743788585066795
pytorch_msssim.ms_ssim (K2=0.4): 0.6563504934310913
3. train your autoencoder with MS_SSIM
see 'tests/ae_example'
left: original image, right: reconstructed image
References
https://github.com/jorge-pessoa/pytorch-msssim
https://ece.uwaterloo.ca/~z70wang/research/ssim/
https://ece.uwaterloo.ca/~z70wang/publications/msssim.pdf
Matlab Code
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for pytorch_msssim-0.1.4-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 6e377019b95f67ffe690b382ad84606156e7a0821c0f8dd4d0a2b5bd6c7fcf37 |
|
MD5 | 36f310f28735f35526897fed0ea6ab62 |
|
BLAKE2b-256 | 38c4d42e97aef04c15771714706532edbe61829fd332c0df42e0e105a03a1b47 |