mirror of
https://github.com/AIGC-Audio/AudioGPT.git
synced 2026-09-02 04:00:43 +02:00
392 lines
14 KiB
Python
392 lines
14 KiB
Python
# '''
|
||
# https://github.com/One-sixth/ms_ssim_pytorch/blob/master/ssim.py
|
||
# '''
|
||
#
|
||
# import torch
|
||
# import torch.jit
|
||
# import torch.nn.functional as F
|
||
#
|
||
#
|
||
# @torch.jit.script
|
||
# def create_window(window_size: int, sigma: float, channel: int):
|
||
# '''
|
||
# Create 1-D gauss kernel
|
||
# :param window_size: the size of gauss kernel
|
||
# :param sigma: sigma of normal distribution
|
||
# :param channel: input channel
|
||
# :return: 1D kernel
|
||
# '''
|
||
# coords = torch.arange(window_size, dtype=torch.float)
|
||
# coords -= window_size // 2
|
||
#
|
||
# g = torch.exp(-(coords ** 2) / (2 * sigma ** 2))
|
||
# g /= g.sum()
|
||
#
|
||
# g = g.reshape(1, 1, 1, -1).repeat(channel, 1, 1, 1)
|
||
# return g
|
||
#
|
||
#
|
||
# @torch.jit.script
|
||
# def _gaussian_filter(x, window_1d, use_padding: bool):
|
||
# '''
|
||
# Blur input with 1-D kernel
|
||
# :param x: batch of tensors to be blured
|
||
# :param window_1d: 1-D gauss kernel
|
||
# :param use_padding: padding image before conv
|
||
# :return: blured tensors
|
||
# '''
|
||
# C = x.shape[1]
|
||
# padding = 0
|
||
# if use_padding:
|
||
# window_size = window_1d.shape[3]
|
||
# padding = window_size // 2
|
||
# out = F.conv2d(x, window_1d, stride=1, padding=(0, padding), groups=C)
|
||
# out = F.conv2d(out, window_1d.transpose(2, 3), stride=1, padding=(padding, 0), groups=C)
|
||
# return out
|
||
#
|
||
#
|
||
# @torch.jit.script
|
||
# def ssim(X, Y, window, data_range: float, use_padding: bool = False):
|
||
# '''
|
||
# Calculate ssim index for X and Y
|
||
# :param X: images [B, C, H, N_bins]
|
||
# :param Y: images [B, C, H, N_bins]
|
||
# :param window: 1-D gauss kernel
|
||
# :param data_range: value range of input images. (usually 1.0 or 255)
|
||
# :param use_padding: padding image before conv
|
||
# :return:
|
||
# '''
|
||
#
|
||
# K1 = 0.01
|
||
# K2 = 0.03
|
||
# compensation = 1.0
|
||
#
|
||
# C1 = (K1 * data_range) ** 2
|
||
# C2 = (K2 * data_range) ** 2
|
||
#
|
||
# mu1 = _gaussian_filter(X, window, use_padding)
|
||
# mu2 = _gaussian_filter(Y, window, use_padding)
|
||
# sigma1_sq = _gaussian_filter(X * X, window, use_padding)
|
||
# sigma2_sq = _gaussian_filter(Y * Y, window, use_padding)
|
||
# sigma12 = _gaussian_filter(X * Y, window, use_padding)
|
||
#
|
||
# mu1_sq = mu1.pow(2)
|
||
# mu2_sq = mu2.pow(2)
|
||
# mu1_mu2 = mu1 * mu2
|
||
#
|
||
# sigma1_sq = compensation * (sigma1_sq - mu1_sq)
|
||
# sigma2_sq = compensation * (sigma2_sq - mu2_sq)
|
||
# sigma12 = compensation * (sigma12 - mu1_mu2)
|
||
#
|
||
# cs_map = (2 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2)
|
||
# # Fixed the issue that the negative value of cs_map caused ms_ssim to output Nan.
|
||
# cs_map = cs_map.clamp_min(0.)
|
||
# ssim_map = ((2 * mu1_mu2 + C1) / (mu1_sq + mu2_sq + C1)) * cs_map
|
||
#
|
||
# ssim_val = ssim_map.mean(dim=(1, 2, 3)) # reduce along CHW
|
||
# cs = cs_map.mean(dim=(1, 2, 3))
|
||
#
|
||
# return ssim_val, cs
|
||
#
|
||
#
|
||
# @torch.jit.script
|
||
# def ms_ssim(X, Y, window, data_range: float, weights, use_padding: bool = False, eps: float = 1e-8):
|
||
# '''
|
||
# interface of ms-ssim
|
||
# :param X: a batch of images, (N,C,H,W)
|
||
# :param Y: a batch of images, (N,C,H,W)
|
||
# :param window: 1-D gauss kernel
|
||
# :param data_range: value range of input images. (usually 1.0 or 255)
|
||
# :param weights: weights for different levels
|
||
# :param use_padding: padding image before conv
|
||
# :param eps: use for avoid grad nan.
|
||
# :return:
|
||
# '''
|
||
# levels = weights.shape[0]
|
||
# cs_vals = []
|
||
# ssim_vals = []
|
||
# for _ in range(levels):
|
||
# ssim_val, cs = ssim(X, Y, window=window, data_range=data_range, use_padding=use_padding)
|
||
# # Use for fix a issue. When c = a ** b and a is 0, c.backward() will cause the a.grad become inf.
|
||
# ssim_val = ssim_val.clamp_min(eps)
|
||
# cs = cs.clamp_min(eps)
|
||
# cs_vals.append(cs)
|
||
#
|
||
# ssim_vals.append(ssim_val)
|
||
# padding = (X.shape[2] % 2, X.shape[3] % 2)
|
||
# X = F.avg_pool2d(X, kernel_size=2, stride=2, padding=padding)
|
||
# Y = F.avg_pool2d(Y, kernel_size=2, stride=2, padding=padding)
|
||
#
|
||
# cs_vals = torch.stack(cs_vals, dim=0)
|
||
# ms_ssim_val = torch.prod((cs_vals[:-1] ** weights[:-1].unsqueeze(1)) * (ssim_vals[-1] ** weights[-1]), dim=0)
|
||
# return ms_ssim_val
|
||
#
|
||
#
|
||
# class SSIM(torch.jit.ScriptModule):
|
||
# __constants__ = ['data_range', 'use_padding']
|
||
#
|
||
# def __init__(self, window_size=11, window_sigma=1.5, data_range=255., channel=3, use_padding=False):
|
||
# '''
|
||
# :param window_size: the size of gauss kernel
|
||
# :param window_sigma: sigma of normal distribution
|
||
# :param data_range: value range of input images. (usually 1.0 or 255)
|
||
# :param channel: input channels (default: 3)
|
||
# :param use_padding: padding image before conv
|
||
# '''
|
||
# super().__init__()
|
||
# assert window_size % 2 == 1, 'Window size must be odd.'
|
||
# window = create_window(window_size, window_sigma, channel)
|
||
# self.register_buffer('window', window)
|
||
# self.data_range = data_range
|
||
# self.use_padding = use_padding
|
||
#
|
||
# @torch.jit.script_method
|
||
# def forward(self, X, Y):
|
||
# r = ssim(X, Y, window=self.window, data_range=self.data_range, use_padding=self.use_padding)
|
||
# return r[0]
|
||
#
|
||
#
|
||
# class MS_SSIM(torch.jit.ScriptModule):
|
||
# __constants__ = ['data_range', 'use_padding', 'eps']
|
||
#
|
||
# def __init__(self, window_size=11, window_sigma=1.5, data_range=255., channel=3, use_padding=False, weights=None,
|
||
# levels=None, eps=1e-8):
|
||
# '''
|
||
# class for ms-ssim
|
||
# :param window_size: the size of gauss kernel
|
||
# :param window_sigma: sigma of normal distribution
|
||
# :param data_range: value range of input images. (usually 1.0 or 255)
|
||
# :param channel: input channels
|
||
# :param use_padding: padding image before conv
|
||
# :param weights: weights for different levels. (default [0.0448, 0.2856, 0.3001, 0.2363, 0.1333])
|
||
# :param levels: number of downsampling
|
||
# :param eps: Use for fix a issue. When c = a ** b and a is 0, c.backward() will cause the a.grad become inf.
|
||
# '''
|
||
# super().__init__()
|
||
# assert window_size % 2 == 1, 'Window size must be odd.'
|
||
# self.data_range = data_range
|
||
# self.use_padding = use_padding
|
||
# self.eps = eps
|
||
#
|
||
# window = create_window(window_size, window_sigma, channel)
|
||
# self.register_buffer('window', window)
|
||
#
|
||
# if weights is None:
|
||
# weights = [0.0448, 0.2856, 0.3001, 0.2363, 0.1333]
|
||
# weights = torch.tensor(weights, dtype=torch.float)
|
||
#
|
||
# if levels is not None:
|
||
# weights = weights[:levels]
|
||
# weights = weights / weights.sum()
|
||
#
|
||
# self.register_buffer('weights', weights)
|
||
#
|
||
# @torch.jit.script_method
|
||
# def forward(self, X, Y):
|
||
# return ms_ssim(X, Y, window=self.window, data_range=self.data_range, weights=self.weights,
|
||
# use_padding=self.use_padding, eps=self.eps)
|
||
#
|
||
#
|
||
# if __name__ == '__main__':
|
||
# print('Simple Test')
|
||
# im = torch.randint(0, 255, (5, 3, 256, 256), dtype=torch.float, device='cuda')
|
||
# img1 = im / 255
|
||
# img2 = img1 * 0.5
|
||
#
|
||
# losser = SSIM(data_range=1.).cuda()
|
||
# loss = losser(img1, img2).mean()
|
||
#
|
||
# losser2 = MS_SSIM(data_range=1.).cuda()
|
||
# loss2 = losser2(img1, img2).mean()
|
||
#
|
||
# print(loss.item())
|
||
# print(loss2.item())
|
||
#
|
||
# if __name__ == '__main__':
|
||
# print('Training Test')
|
||
# import cv2
|
||
# import torch.optim
|
||
# import numpy as np
|
||
# import imageio
|
||
# import time
|
||
#
|
||
# out_test_video = False
|
||
# # 最好不要直接输出gif图，会非常大，最好先输出mkv文件后用ffmpeg转换到GIF
|
||
# video_use_gif = False
|
||
#
|
||
# im = cv2.imread('test_img1.jpg', 1)
|
||
# t_im = torch.from_numpy(im).cuda().permute(2, 0, 1).float()[None] / 255.
|
||
#
|
||
# if out_test_video:
|
||
# if video_use_gif:
|
||
# fps = 0.5
|
||
# out_wh = (im.shape[1] // 2, im.shape[0] // 2)
|
||
# suffix = '.gif'
|
||
# else:
|
||
# fps = 5
|
||
# out_wh = (im.shape[1], im.shape[0])
|
||
# suffix = '.mkv'
|
||
# video_last_time = time.perf_counter()
|
||
# video = imageio.get_writer('ssim_test' + suffix, fps=fps)
|
||
#
|
||
# # 测试ssim
|
||
# print('Training SSIM')
|
||
# rand_im = torch.randint_like(t_im, 0, 255, dtype=torch.float32) / 255.
|
||
# rand_im.requires_grad = True
|
||
# optim = torch.optim.Adam([rand_im], 0.003, eps=1e-8)
|
||
# losser = SSIM(data_range=1., channel=t_im.shape[1]).cuda()
|
||
# ssim_score = 0
|
||
# while ssim_score < 0.999:
|
||
# optim.zero_grad()
|
||
# loss = losser(rand_im, t_im)
|
||
# (-loss).sum().backward()
|
||
# ssim_score = loss.item()
|
||
# optim.step()
|
||
# r_im = np.transpose(rand_im.detach().cpu().numpy().clip(0, 1) * 255, [0, 2, 3, 1]).astype(np.uint8)[0]
|
||
# r_im = cv2.putText(r_im, 'ssim %f' % ssim_score, (10, 30), cv2.FONT_HERSHEY_PLAIN, 2, (255, 0, 0), 2)
|
||
#
|
||
# if out_test_video:
|
||
# if time.perf_counter() - video_last_time > 1. / fps:
|
||
# video_last_time = time.perf_counter()
|
||
# out_frame = cv2.cvtColor(r_im, cv2.COLOR_BGR2RGB)
|
||
# out_frame = cv2.resize(out_frame, out_wh, interpolation=cv2.INTER_AREA)
|
||
# if isinstance(out_frame, cv2.UMat):
|
||
# out_frame = out_frame.get()
|
||
# video.append_data(out_frame)
|
||
#
|
||
# cv2.imshow('ssim', r_im)
|
||
# cv2.setWindowTitle('ssim', 'ssim %f' % ssim_score)
|
||
# cv2.waitKey(1)
|
||
#
|
||
# if out_test_video:
|
||
# video.close()
|
||
#
|
||
# # 测试ms_ssim
|
||
# if out_test_video:
|
||
# if video_use_gif:
|
||
# fps = 0.5
|
||
# out_wh = (im.shape[1] // 2, im.shape[0] // 2)
|
||
# suffix = '.gif'
|
||
# else:
|
||
# fps = 5
|
||
# out_wh = (im.shape[1], im.shape[0])
|
||
# suffix = '.mkv'
|
||
# video_last_time = time.perf_counter()
|
||
# video = imageio.get_writer('ms_ssim_test' + suffix, fps=fps)
|
||
#
|
||
# print('Training MS_SSIM')
|
||
# rand_im = torch.randint_like(t_im, 0, 255, dtype=torch.float32) / 255.
|
||
# rand_im.requires_grad = True
|
||
# optim = torch.optim.Adam([rand_im], 0.003, eps=1e-8)
|
||
# losser = MS_SSIM(data_range=1., channel=t_im.shape[1]).cuda()
|
||
# ssim_score = 0
|
||
# while ssim_score < 0.999:
|
||
# optim.zero_grad()
|
||
# loss = losser(rand_im, t_im)
|
||
# (-loss).sum().backward()
|
||
# ssim_score = loss.item()
|
||
# optim.step()
|
||
# r_im = np.transpose(rand_im.detach().cpu().numpy().clip(0, 1) * 255, [0, 2, 3, 1]).astype(np.uint8)[0]
|
||
# r_im = cv2.putText(r_im, 'ms_ssim %f' % ssim_score, (10, 30), cv2.FONT_HERSHEY_PLAIN, 2, (255, 0, 0), 2)
|
||
#
|
||
# if out_test_video:
|
||
# if time.perf_counter() - video_last_time > 1. / fps:
|
||
# video_last_time = time.perf_counter()
|
||
# out_frame = cv2.cvtColor(r_im, cv2.COLOR_BGR2RGB)
|
||
# out_frame = cv2.resize(out_frame, out_wh, interpolation=cv2.INTER_AREA)
|
||
# if isinstance(out_frame, cv2.UMat):
|
||
# out_frame = out_frame.get()
|
||
# video.append_data(out_frame)
|
||
#
|
||
# cv2.imshow('ms_ssim', r_im)
|
||
# cv2.setWindowTitle('ms_ssim', 'ms_ssim %f' % ssim_score)
|
||
# cv2.waitKey(1)
|
||
#
|
||
# if out_test_video:
|
||
# video.close()
|
||
|
||
"""
|
||
Adapted from https://github.com/Po-Hsun-Su/pytorch-ssim
|
||
"""
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch.autograd import Variable
|
||
import numpy as np
|
||
from math import exp
|
||
|
||
|
||
def gaussian(window_size, sigma):
|
||
gauss = torch.Tensor([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 _ssim(img1, img2, window, window_size, channel, size_average=True):
|
||
mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
|
||
mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
|
||
|
||
mu1_sq = mu1.pow(2)
|
||
mu2_sq = mu2.pow(2)
|
||
mu1_mu2 = mu1 * mu2
|
||
|
||
sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
|
||
sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
|
||
sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
|
||
|
||
C1 = 0.01 ** 2
|
||
C2 = 0.03 ** 2
|
||
|
||
ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
|
||
|
||
if size_average:
|
||
return ssim_map.mean()
|
||
else:
|
||
return ssim_map.mean(1)
|
||
|
||
|
||
class SSIM(torch.nn.Module):
|
||
def __init__(self, window_size=11, size_average=True):
|
||
super(SSIM, self).__init__()
|
||
self.window_size = window_size
|
||
self.size_average = size_average
|
||
self.channel = 1
|
||
self.window = create_window(window_size, self.channel)
|
||
|
||
def forward(self, img1, img2):
|
||
(_, channel, _, _) = img1.size()
|
||
|
||
if channel == self.channel and self.window.data.type() == img1.data.type():
|
||
window = self.window
|
||
else:
|
||
window = create_window(self.window_size, channel)
|
||
|
||
if img1.is_cuda:
|
||
window = window.cuda(img1.get_device())
|
||
window = window.type_as(img1)
|
||
|
||
self.window = window
|
||
self.channel = channel
|
||
|
||
return _ssim(img1, img2, window, self.window_size, channel, self.size_average)
|
||
|
||
|
||
window = None
|
||
|
||
|
||
def ssim(img1, img2, window_size=11, size_average=True):
|
||
(_, channel, _, _) = img1.size()
|
||
global window
|
||
if window is None:
|
||
window = create_window(window_size, channel)
|
||
if img1.is_cuda:
|
||
window = window.cuda(img1.get_device())
|
||
window = window.type_as(img1)
|
||
return _ssim(img1, img2, window, window_size, channel, size_average)
|