help for developemnt of nn module for structural similarity index measure (SSIM) ?

[sri phani krishna karri]

Hi all,

I am planning to construct a cost function based on SSIM for image filtering purposed. Conventional MSE based CNN improves PSNR but creates smudged images.

A typical SSIM look slike

n=8
h = torch.ones(n,n:div(n*n))
       

ssim = function (x,y)
       mu_x = im.convolve(x,h)
       mu_y = im.convolve(y,h)
       mu_x_sq = torch.cmul(mu_x,mu_x)
       mu_y_sq = torch.cmul(mu_y,mu_y)
       mu_xy = torch.cmul(mu_x,mu_y)
       sigma_x_sq = im.convolve(torch.cmul(x,x),h)-mu_x_sq
       sigma_y_sq = im.convolve(torch.cmul(y,y),h)-mu_y_sq
       sigma_xy = im.convolve(torch.cmul(x,y),h)-torch.cmul(mu_x,mu_y)
       A1 = mu_xy:mul(2):add(5)
       A2 = sigma_xy:mul(2):add(45.5)
       B1 = mu_x_sq+mu_y_sq:add(5)
       B2 = sigma_x_sq+sigma_y_sq:add(45.5)
       --v = torch.cdiv(torch.cmul(B1,B2),torch.cmul(A1,A2)):mean()
       v = torch.ones(1)-torch.cdiv(torch.cmul(A1,A2),torch.cmul(B1,B2)):mean() --(converting to cost)
       return v
end

How should the nn module representation look like

local ssim, Parent = torch.class('nn.ssim', 'nn.Module') local im = require('image') function ssim:__init() parent.__init(self)

self.n = n or 8

self.mu_x = torch.Tensor()

self.mu_y = torch.Tensor()

self.mu_x_sq = torch.Tensor()

self.mu_y_sq  = torch.Tensor()

self.mu_xy = torch.Tensor()

self.sigma_x_sq = torch.Tensor()

self.sigma_y_sq = torch.Tensor()

self.sigma_xy = torch.Tensor()

self.mu_A1 = torch.Tensor()

self.mu_A2 = torch.Tensor()

self.mu_B1 = torch.Tensor()

self.mu_B2 = torch.Tensor()

end

function ssim:updateOutput(input,target)

       self.output:resizeAs(input):copy(input)
       self.target:resizeAs(target):copy(target)
       self.h = torch.ones(self.n,self.n:div(self.n*self.n))

       self.mu_x = im.convolve(self.output,self.h)

       self.mu_y = im.convolve(self.target,h)

       self.mu_x_sq = torch.cmul(self.mu_x,self.mu_x)

       self.mu_y_sq = torch.cmul(self.mu_y,self.mu_y)
       self.mu_xy = torch.cmul(self.mu_x,self.mu_y)

       self.sigma_x_sq = im.convolve(torch.cmul(self.output,self.self.output),self.h)-self.mu_x_sq

       self.sigma_y_sq = im.convolve(torch.cmul(self.target,self.target),self.h)-self.mu_y_sq
       self.sigma_xy = im.convolve(torch.cmul(self.output,self.target),selfh)-torch.cmul(self.mu_x,self.mu_y)

       self.A1 = self.mu_xy:mul(2):add(5)
       self.A2 = self.sigma_xy:mul(2):add(45.5)
       self.B1 = self.mu_x_sq+self.mu_y_sq:add(5)
       self.B2 = self.sigma_x_sq+self.sigma_y_sq:add(45.5)
       v = torch.ones(1)-torch.cdiv(torch.cmul(self.A1,self.A2),torch.cmul(self.B1,self.B2)):mean()                    --(converting to cost)
       return v
end function ssim:updateGradInput(input, gradOutput)

-- to be done end

Am I going in the right path

[Ahmed Abo Bakr]

Hi sri phani, 

have you figured this out, as I am gonna do the same thing?

[Mika S]

Are you sure you would be able to train such a NN. Cost functions are tradeoffs between what is ideal and what gives us nicer gradients

[sri phani krishna karri]

https://github.com/ultrai/SRCNN/blob/master/SRCNN_graph_ssim4#L15-L60

[aMY]

can anyone help with the code for multi scale ssim? Thanks.