pytorch ssim

pytorch-ssim（不保留此存储库）

该代码不起作用，因为它在超级旧的pytorch上。

可区分的结构相似性（SSIM）指数。

安装

1. 克隆这个仓库。
2. 复制项目中的“ pytorch_ssim”文件夹。

例子

基本用法

 import pytorch_ssim
import torch
from torch . autograd import Variable

img1 = Variable ( torch . rand ( 1 , 1 , 256 , 256 ))
img2 = Variable ( torch . rand ( 1 , 1 , 256 , 256 ))

if torch . cuda . is_available ():
    img1 = img1 . cuda ()
    img2 = img2 . cuda ()

print ( pytorch_ssim . ssim ( img1 , img2 ))

ssim_loss = pytorch_ssim . SSIM ( window_size = 11 )

print ( ssim_loss ( img1 , img2 ))

最大化SSIM

 import pytorch_ssim
import torch
from torch . autograd import Variable
from torch import optim
import cv2
import numpy as np

npImg1 = cv2 . imread ( "einstein.png" )

img1 = torch . from_numpy ( np . rollaxis ( npImg1 , 2 )). float (). unsqueeze ( 0 ) / 255.0
img2 = torch . rand ( img1 . size ())

if torch . cuda . is_available ():
    img1 = img1 . cuda ()
    img2 = img2 . cuda ()


img1 = Variable ( img1 ,  requires_grad = False )
img2 = Variable ( img2 , requires_grad = True )


# Functional: pytorch_ssim.ssim(img1, img2, window_size = 11, size_average = True)
ssim_value = pytorch_ssim . ssim ( img1 , img2 ). data [ 0 ]
print ( "Initial ssim:" , ssim_value )

# Module: pytorch_ssim.SSIM(window_size = 11, size_average = True)
ssim_loss = pytorch_ssim . SSIM ()

optimizer = optim . Adam ([ img2 ], lr = 0.01 )

while ssim_value < 0.95 :
    optimizer . zero_grad ()
    ssim_out = - ssim_loss ( img1 , img2 )
    ssim_value = - ssim_out . data [ 0 ]
    print ( ssim_value )
    ssim_out . backward ()
    optimizer . step ()

参考

https://ece.uwaterloo.ca/~z70wang/research/ssim/