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+
For faster calculation speed, the 2D convolution (Gaussian Blur) is replaced by two 1D convolutions.
see Gaussian_blur wiki.
All calculations will be on the same device as inputs.
update
2019.6.17
Now it is faster than compare_ssim thanks to One-sixth's contribution
Install
python setup.py install
or
pip install pytorch-msssim
Example
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
# X: (N,3,H,W) a batch of RGB images with values ranging from 0 to 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,)
# or set 'size_average=True' to get a scalar value as loss.
ssim_loss = ssim( X, Y, data_range=255, size_average=True) # return a scalar value
ms_ssim_loss = ms_ssim( X, Y, data_range=255, size_average=True )
# or 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 = ssim_module(X, Y)
ms_ssim_loss = ms_ssim_module(X, Y)
Tests
Compared with skimage.measure.compare_ssim on CPU.
The outputs:
Downloading test image...
====> Single Image
sigma=0.000000 compare_ssim=1.000000 (417.248964 ms) ssim_torch=1.000000 (257.593870 ms)
sigma=1.000000 compare_ssim=0.991320 (326.905012 ms) ssim_torch=0.991320 (135.488033 ms)
sigma=2.000000 compare_ssim=0.966521 (485.862017 ms) ssim_torch=0.966520 (237.199068 ms)
sigma=3.000000 compare_ssim=0.928799 (323.492050 ms) ssim_torch=0.928797 (148.905993 ms)
sigma=4.000000 compare_ssim=0.882271 (290.801048 ms) ssim_torch=0.882267 (146.914005 ms)
sigma=5.000000 compare_ssim=0.831310 (282.787085 ms) ssim_torch=0.831306 (148.653984 ms)
sigma=6.000000 compare_ssim=0.778222 (308.619022 ms) ssim_torch=0.778217 (147.915840 ms)
sigma=7.000000 compare_ssim=0.726444 (290.637970 ms) ssim_torch=0.726438 (133.754253 ms)
sigma=8.000000 compare_ssim=0.676345 (294.582129 ms) ssim_torch=0.676339 (144.154072 ms)
sigma=9.000000 compare_ssim=0.629922 (300.610065 ms) ssim_torch=0.629916 (141.150951 ms)
Pass
====> Batch
Pass
An autoencoder trained with MS_SSIM
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 Distributions
Hashes for pytorch_msssim-0.1.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | ffb2697244a2bdd0341df4391d3f8de02ef2cb6591910aa05d1c1f4e764e4264 |
|
MD5 | f8adf8aef43452c1c2003c32990a2269 |
|
BLAKE2b-256 | 1431fcd8ca660dbb10ab592b828da96ac31963d3261b4382393ccff995baa33a |