Hi.
I’m trying to implement the MEFSSIM metric proposed in 《Perceptual Quality Assessment for
Multi-Exposure Image Fusion》.
However, I find it hard to implement it in an efficient and highly-vectorized way. Does anyone have a better solution? (There are several pytorch implementations of mefssim available on github, but they are all somewhat inconsistent with the matlab implementation provided by the author)
The matlab implementation is as follow:
function [Q, qMap] = mef_ssim(imgSeq, fI,  K, window) 

if (nargin < 2 || nargin > 4) % restrict # of params to be in range [2, 4]
   Q = -Inf;
   qMap = Inf;
   return;
end

if (~exist('K', 'var'))
   K = 0.03;
end

if (~exist('window', 'var'))
   window = fspecial('gaussian', 11, 1.5); % fspecial('gaussian',hsize,sigma) returns a rotationally symmetric Gaussian lowpass filter of size hsize with standard deviation sigma.
end


imgSeq = double(imgSeq);
fI = double(fI);
[s1, s2, s3] = size(imgSeq); % s1, s2 is height and width of imgseq, s3 is number of imgs
wSize = size(window,1);
sWindow = ones(wSize) / wSize^2; % square window used to calculate the distance, ones(size) return an all-one matirx of size (size,size)
bd = floor(wSize/2);
mu = zeros(s1-2*bd, s2-2*bd, s3);
ed = zeros(s1-2*bd, s2-2*bd, s3);
for i = 1:s3 % for all input imgs in input sequence, apply mean filter
    img = squeeze(imgSeq(:,:,i));
    % mu is mean value of the patch u_xk
    mu(:,:,i) = filter2(sWindow, img, 'valid'); % valid means no zero-padding, thus the size changes
    muSq = mu(:,:,i) .* mu(:,:,i);
    sigmaSq = filter2(sWindow, img.*img, 'valid') - muSq;
    ed(:,:,i) =  sqrt( max( wSize^2 * sigmaSq, 0 ) ) + 0.001; % add a small constant to avoid instability
end

R = zeros(s1-2*bd,s2-2*bd); % consistency map which could be used as an output if necessary
for i = bd+1:s1-bd
    for j = bd+1:s2-bd
        vecs = reshape(imgSeq(i-bd:i+bd,j-bd:j+bd,:),[wSize*wSize, s3]);
        denominator = 0;
        for k = 1:s3
            denominator = denominator + norm(vecs(:,k) - mu(i-bd,j-bd,k));
        end
        numerator = norm(sum(vecs,2) - mean(sum(vecs,2)));
        R(i-bd,j-bd) = (numerator + eps) / (denominator + eps); % eq (6)
    end
end

R(R > 1) = 1 - eps; % get rid of numerical instability
R(R < 0) = 0 + eps;


p = tan(pi/2 * R); % eq(7)
p( p >  10 ) = 10; % to avoid blow up (large number such as 10 is equivalent to taking maximum)
p = repmat(p,[1,1,s3]);


wMap = (ed / wSize).^p + eps; % to avoid blowing up eq(5)
normalizer = sum(wMap,3);
wMap = wMap ./ repmat(normalizer,[1,1,s3]);
maxEd = max(ed,[],3); % eq(3)

C = (K * 255)^2;
qMap = zeros(s1-2*bd, s2-2*bd); 
for i = bd+1:s1-bd
    for j = bd+1:s2-bd
        blocks = imgSeq(i-bd:i+bd,j-bd:j+bd,:);
        rBlock = zeros(wSize,wSize);
        for k = 1 : s3
            rBlock = rBlock  + wMap(i-bd,j-bd,k) * ( blocks(:,:,k) - mu(i-bd,j-bd,k) ) / ed(i-bd,j-bd,k); % eq(4)
        end  
        if norm(rBlock(:)) > 0
            rBlock = rBlock / norm(rBlock(:)) * maxEd(i-bd,j-bd);
        end
        fBlock = fI(i-bd:i+bd,j-bd:j+bd);
        rVec = rBlock(:);
        fVec = fBlock(:);
        mu1 = sum( window(:) .* rVec );
        mu2 = sum( window(:) .* fVec );
        sigma1Sq = sum( window(:) .* (rVec - mu1).^2 );
        sigma2Sq = sum( window(:) .* (fVec - mu2).^2 );
        sigma12 = sum(  window(:) .* (rVec - mu1) .* (fVec - mu2)  );
        qMap(i-bd,j-bd) = ( 2 * sigma12 + C ) ./ ( sigma1Sq + sigma2Sq + C ); 
    end
end

Q = mean2(qMap); % Average or mean of matrix elements

My pytorch implementation:
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
import math


def gaussian(window_size, sigma):
    gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
    return gauss / gauss.sum()


def create_window(window_size, channel):
    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
    return window


def mef_ssim(imgSeq, fI, K=0.03):
    fI = fI.squeeze()  # (batch_size, h, w)
    eps = 2.2204e-16
    window = create_window(11, 1).cuda()

    (batch_size, imgSeq_channel, s1, s2) = imgSeq[0].shape
    s3 = len(imgSeq)
    imgSeq_stack = torch.stack(imgSeq).squeeze().unsqueeze(1).cuda().float()  # (img_num, batch_size, h, w)

    wSize = window.shape[2]
    sWindow = torch.ones(window.size()).cuda().float() / wSize ** 2
    bd = wSize // 2
    mu = []
    ed = []
    for i in range(s3):
        img = imgSeq[i].cuda().float()  # (batch_size, channel, h, w)
        mu.append(F.conv2d(img, sWindow, padding = wSize//2, groups = imgSeq_channel))
        muSq = mu[i] * mu[i]
        sigmaSq = F.conv2d(img * img, sWindow, padding = wSize//2, groups = imgSeq_channel) - muSq
        ed.append(torch.sqrt(wSize ** 2 * sigmaSq))
        ed[i][ed[i]<0.001] = 0.001
    ed = torch.stack(ed).squeeze().unsqueeze(1)  # (img_num, batch_size, h, w)
    mu = torch.stack(mu).squeeze().unsqueeze(1)  # (img_num, batch_size, h, w)
    x_hat = (imgSeq_stack - mu).float()

    SumFilter = torch.ones(window.size()).cuda().float()
    numerator = torch.zeros((x_hat.size())).cuda()
    denominator = torch.zeros((x_hat.size())).cuda()
    for i in range(s3):
        denominator = denominator + torch.sqrt(F.conv2d(x_hat**2, SumFilter, padding = wSize//2, groups = imgSeq_channel))
    numerator = torch.sqrt(F.conv2d((torch.sum(x_hat,0)**2).unsqueeze(0), SumFilter, padding = wSize//2, groups = imgSeq_channel))
    R = (numerator + eps) / (denominator+eps)

    # R = torch.zeros((batch_size, s1 - 2 * bd, s2 - 2 * bd)).cuda()  # (batch_size, h, w)
    # for batch in range(batch_size):
    #     for i in range(bd, s1 - bd):
    #         for j in range(bd, s2 - bd):
    #             vecs = torch.reshape(imgSeq_stack[:, batch, i - bd: i + bd + 1, j - bd : j + bd + 1], (wSize * wSize, s3))
    #             denominator = 0
    #             for k in range(s3):
    #                 denominator = denominator + torch.norm(vecs[:, k] - mu[k, batch, i - bd - 1, j - bd - 1])
    #             numerator = torch.norm(torch.sum(vecs, 1) - torch.mean(torch.sum(vecs, 1)))
    #             R[batch, i - bd - 1, j - bd - 1] = (numerator + eps) / (denominator + eps)

    R[R > 1] = 1 - eps
    R[R < 0] = 0 + eps

    p = torch.tan( 1.57 * R)
    p[p > 10] = 10
    p_rep = []
    for i in range(s3):
        p_rep.append(p)
    p = torch.stack(p_rep)  # (img_num, batch_size, h, w)

    wMap = (ed / wSize) ** p + eps  # (img_num, batch_size, h, w)
    normalizer = torch.sum(wMap, 0)
    normalizer_rep = []
    for i in range(s3):
        normalizer_rep.append(normalizer)
    normalizer_rep = torch.stack(normalizer_rep)
    wMap = wMap / normalizer_rep
    maxEd, maxEd_idx = torch.max(ed, dim=0)

    C = (K * 1) ** 2
    qMap = torch.zeros((x_hat.size())).cuda()

    # UnitFliter = torch.zeros(window.size()).cuda().float()
    # UnitFliter[:,:,5,5] = 1
    # rBlock = torch.sum(F.conv2d(wMap, UnitFliter, padding = wSize//2, groups = imgSeq_channel) *  F.conv2d(x_hat, UnitFliter, padding = wSize//2, groups = imgSeq_channel) /  F.conv2d(ed, UnitFliter, padding = wSize//2, groups = imgSeq_channel), dim=0)
    # if ()
    for batch in range(batch_size):
        for i in range(bd, s1 - bd):
            for j in range(bd, s2 - bd):
                blocks = imgSeq_stack[:, batch, i - bd: i + bd, j-bd : j + bd]
                rBlock = torch.zeros((wSize, wSize))
                for k in range(s3):
                    rBlock = rBlock + wMap[k, batch, i - bd, j - bd] * (
                                blocks[k, :, :, :] - mu[k, batch, i - bd, j - bd]) / ed[k, batch, i - bd, j - bd]

                if (torch.norm(rBlock) > 0):
                    rBlock = rBlock / torch.norm(rBlock) * maxEd[batch, i - bd, j - bd]

                fBlock = fI[batch, i - bd:i + bd, j - bd:j + bd]
                rVec = rBlock.view(-1, 1)
                fVec = fBlock.view(-1, 1)
                windowVec = window.view(-1, 1)
                mu1 = torch.sum(windowVec * rVec)
                mu2 = torch.sum(windowVec * fVec)
                sigma1Sq = torch.sum(windowVec * (rVec - mu1) ** 2)
                sigma2Sq = torch.sum(windowVec * (fVec - mu2) ** 2)
                sigma12 = torch.sum(windowVec * (rVec - mu1) * (fVec - mu2))
                qMap[batch, i - bd, j - bd] = (2 * sigma12 + C) / (sigma1Sq + sigma2Sq + C)

    Q = torch.mean(qMap)