Giter VIP home page Giter VIP logo

pytorch-msssim's Introduction

Pytorch MS-SSIM

Fast and differentiable MS-SSIM and SSIM for pytorch 1.0+

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

2019.8.15
Apply to 5D tensor #6

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 = 1 - ssim_module(X, Y)
ms_ssim_loss = 1 - 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

results 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

pytorch-msssim's People

Contributors

vainf avatar one-sixth avatar

Watchers

James Cloos avatar

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    ๐Ÿ–– Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. ๐Ÿ“Š๐Ÿ“ˆ๐ŸŽ‰

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google โค๏ธ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.