基于pytorch计算ssim和ms-ssim
使用pytorch计算两组图片的ssim和ms-ssim
首先是SSIM和MS-SSIM类（ssim.py）
import torch
import torch.nn.functional as F
def _fspecial_gauss_1d(size, sigma):
coords = torch.arange(size).to(dtype=torch.float)
coords -= size//2
g = torch.exp(-(coords**2) / (2*sigma**2))
g /= g.sum()
return g.unsqueeze(0).unsqueeze(0)
def gaussian_filter(input, win):
N, C, H, W = input.shape
out = F.conv2d(input, win, stride=1, padding=0, groups=C)
out = F.conv2d(out, win.transpose(2, 3), stride=1, padding=0, groups=C)
return out
def _ssim(X, Y, win, data_range=1023, size_average=True, full=False):
K1 = 0.01
K2 = 0.03
batch, channel, height, width = X.shape
compensation = 1.0
C1 = (K1 * data_range)**2
C2 = (K2 * data_range)**2
win = win.to(X.device, dtype=X.dtype)
mu1 = gaussian_filter(X, win)
mu2 = gaussian_filter(Y, win)
mu1_sq = mu1.pow(2)
mu2_sq = mu2.pow(2)
mu1_mu2 = mu1 * mu2
sigma1_sq = compensation * ( gaussian_filter(X * X, win) - mu1_sq )
sigma2_sq = compensation * ( gaussian_filter(Y * Y, win) - mu2_sq )
sigma12 = compensation * ( gaussian_filter(X * Y, win) - mu1_mu2 )
cs_map = (2 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2)
ssim_map = ((2 * mu1_mu2 + C1) / (mu1_sq + mu2_sq + C1)) * cs_map
if size_average:
ssim_val = ssim_map.mean()
cs = cs_map.mean()
else:
ssim_val = ssim_map.mean(-1).mean(-1).mean(-1) # reduce along CHW
cs = cs_map.mean(-1).mean(-1).mean(-1)
if full:
return ssim_val, cs
else:
return ssim_val
def ssim(X, Y, win_size=11, win_sigma=10, win=None, data_range=1, size_average=True, full=False):
if len(X.shape) != 4:
raise ValueError('Input images must 4-d tensor.')
if not X.type() == Y.type():
raise ValueError('Input images must have the same dtype.')
if not X.shape == Y.shape:
raise ValueError('Input images must have the same dimensions.')
if not (win_size % 2 == 1):
raise ValueError('Window size must be odd.')
win_sigma = win_sigma
if win is None:
win = _fspecial_gauss_1d(win_size, win_sigma)
win = win.repeat(X.shape[1], 1, 1, 1)
else:
win_size = win.shape[-1]
ssim_val, cs = _ssim(X, Y,
win=win,
data_range=data_range,
size_average=False,
full=True)
if size_average:
ssim_val = ssim_val.mean()
cs = cs.mean()
if full:
return ssim_val, cs
else:
return ssim_val
def ms_ssim(X, Y, win_size=11, win_sigma=10, win=None, data_range=1, size_average=True, full=False, weights=None):
if len(X.shape) != 4:
raise ValueError('Input images must 4-d tensor.')
if not X.type() == Y.type():
raise ValueError('Input images must have the same dtype.')
if not X.shape == Y.shape:
raise ValueError('Input images must have the same dimensions.')
if not (win_size % 2 == 1):
raise ValueError('Window size must be odd.')
if weights is None:
weights = torch.FloatTensor(
[0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(X.device, dtype=X.dtype)
win_sigma = win_sigma
if win is None:
win = _fspecial_gauss_1d(win_size, win_sigma)
win = win.repeat(X.shape[1], 1, 1, 1)
else:
win_size = win.shape[-1]
levels = weights.shape[0]
mcs = []
for _ in range(levels):
ssim_val, cs = _ssim(X, Y,
win=win,
data_range=data_range,
size_average=False,
full=True)
mcs.append(cs)
padding = (X.shape[2] % 2, X.shape[3] % 2)
X = F.avg_pool2d(X, kernel_size=2, padding=padding)
Y = F.avg_pool2d(Y, kernel_size=2, padding=padding)
mcs = torch.stack(mcs, dim=0) # mcs, (level, batch)
# weights, (level)
msssim_val = torch.prod((mcs[:-1] ** weights[:-1].unsqueeze(1))
* (ssim_val ** weights[-1]), dim=0) # (batch, )
if size_average:
msssim_val = msssim_val.mean()
return msssim_val
# Classes to re-use window
class SSIM(torch.nn.Module):
def __init__(self, win_size=11, win_sigma=1.5, data_range=255, size_average=True, channel=3):
super(SSIM, self).__init__()
self.win = _fspecial_gauss_1d(
win_size, win_sigma).repeat(channel, 1, 1, 1)
self.size_average = size_average
self.data_range = data_range
def forward(self, X, Y):
return ssim(X, Y, win=self.win, data_range=self.data_range, size_average=self.size_average)
class MS_SSIM(torch.nn.Module):
def __init__(self, win_size=11, win_sigma=1.5, data_range=255, size_average=True, channel=3, weights=None):
super(MS_SSIM, self).__init__()
self.win = _fspecial_gauss_1d(
win_size, win_sigma).repeat(channel, 1, 1, 1)
self.size_average = size_average
self.data_range = data_range
self.weights = weights
def forward(self, X, Y):
return ms_ssim(X, Y, win=self.win, size_average=self.size_average, data_range=self.data_range, weights=self.weights)
上面的工具类我在pytorch中当做损失函数使用
使用
这里还用到几个方法，我在下面给出
import argparse
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
from util import (map_range, cv2torch, random_tone_map,
DirectoryDataset, str2bool)、
#这里我把上面的ssim.py放到了一个文件夹中，所以需要这样导入
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
'--batch_size', type=int, default=1, help='Batch size.')
parser.add_argument(
'-d', '--data_root_path_label', default='D:/project_hdr/hdr-expandnet/test_data_hdr', help='Path to hdr data.')
parser.add_argument(
'-v', '--data_root_path_pre', default='D:/project_hdr/myGan_same_size_wights_noexpand _minD/test_data_ldr/981',
help='Path to hdr data.')
parser.add_argument(
'--num_workers',
type=int,
default=1,
help='Number of data loading workers.')
return parser.parse_args()
def transformh(hdr):
hdr = map_range(hdr)
return cv2torch(hdr)
def train(opt):
# 加载训练集
dataset1 = DirectoryDataset(
data_root_path=opt.data_root_path_label, preprocess=transformh)
loader1 = DataLoader(
dataset1,
batch_size=opt.batch_size,
num_workers=opt.num_workers,)
# 加载训练集
dataset2 = DirectoryDataset(
data_root_path=opt.data_root_path_pre, preprocess=transformh)
loader2 = DataLoader(
dataset2,
batch_size=opt.batch_size,
num_workers=opt.num_workers,)
for (ldr_in, hdr_target) in zip(loader2, loader1):
if torch.cuda.is_available():
ldr_in = ldr_in.cuda()
hdr_target = hdr_target.cuda()
pre = ldr_in
real_B = hdr_target
ssim_val = ssim(real_B, pre, data_range=1, size_average=True, ) # return (N,)
ms_ssim_val = ms_ssim(real_B, pre, data_range=1, size_average=True,) # (N,)
rep = (f'ssim_val: {ssim_val},'
f'ms_ssim_val: {ms_ssim_val},')
tqdm.write(rep)
if __name__ == '__main__':
opt = parse_args()
train(opt)
下面是使用到的几个方法
def map_range(x, low=0, high=1):
return np.interp(x, [x.min(), x.max()], [low, high]).astype(x.dtype)
var foo = 'bar';
def cv2torch(np_img):
rgb = np_img[:, :, (2, 1, 0)]
return torch.from_numpy(rgb.swapaxes(1, 2).swapaxes(0, 1))
下面这个类需要opencv
class DirectoryDataset(Dataset):
def __init__(self,
data_root_path='hdr_data',
data_extensions=['.hdr', '.exr'],
load_fn=None,
preprocess=None):
super(DirectoryDataset, self).__init__()
data_root_path = process_path(data_root_path)
self.file_list = []
for root, _, fnames in sorted(os.walk(data_root_path)):
for fname in fnames:
if any(fname.lower().endswith(extension)
for extension in data_extensions):
self.file_list.append(os.path.join(root, fname))
if len(self.file_list) == 0:
msg = 'Could not find any files with extensions:\n[{0}]\nin\n{1}'
raise RuntimeError(
msg.format(', '.join(data_extensions), data_root_path))
self.preprocess = preprocess
def __getitem__(self, index):
dpoint = cv2.imread(
self.file_list[index],
flags=cv2.IMREAD_ANYDEPTH + cv2.IMREAD_COLOR)
if self.preprocess is not None:
dpoint = self.preprocess(dpoint)
return dpoint
def __len__(self):
return len(self.file_list)
def process_path(directory, create=False):
directory = os.path.expanduser(directory)
directory = os.path.normpath(directory)
directory = os.path.abspath(directory)
if create:
try:
os.makedirs(directory)
except:
pass
return directory
大概就是这些，msssim代码来源于github ： https://github.com/VainF/pytorch-msssim
当然，计算ssim还有一些其他更加简便的方法，比如
from skimage.measure import compare_ssim
(score, diff) = compare_ssim(X, Y, full=True)
diff = (diff * 255).astype("float32")
但是python计算msssim的就没有了，所以我用的这个单独的方法去计算的
