Pythorch的快速可微MS-SSIM和SSIM。
pytorch-msssim的Python项目详细描述
Pythorch MS-SSIM
pytorch 1.0+的快速可微ms-ssim和ssim
所有计算将与输入在同一设备上。
更新
2019.6.17
现在,由于One-sixth's contribution,它比比较ssim快
安装
python setup.py install
或
pip install pytorch-msssim
示例
frompytorch_msssimimportssim,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 valuems_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)
测试
与CPU上的skimage.measure.compare_ssim相比。
输出:
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
使用MS SSIM训练的自动编码器
左:原始图像,右:重建图像
参考资料
https://github.com/jorge-pessoa/pytorch-msssim
https://ece.uwaterloo.ca/~z70wang/research/ssim/
https://ece.uwaterloo.ca/~z70wang/publications/msssim.pdf
Matlab Code