File size: 2,518 Bytes
7f43945 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
import time
import numpy as np
from skimage.metrics import peak_signal_noise_ratio, structural_similarity
from skvideo.measure import niqe
class AverageMeter():
""" Computes and stores the average and current value """
def __init__(self):
self.reset()
def reset(self):
""" Reset all statistics """
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
""" Update statistics """
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def accuracy(output, target, topk=(1,)):
""" Computes the precision@k for the specified values of k """
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
# one-hot case
if target.ndimension() > 1:
target = target.max(1)[1]
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0)
res.append(correct_k.mul_(1.0 / batch_size))
return res
def compute_psnr_ssim(recoverd, clean):
assert recoverd.shape == clean.shape
recoverd = np.clip(recoverd.detach().cpu().numpy(), 0, 1)
clean = np.clip(clean.detach().cpu().numpy(), 0, 1)
recoverd = recoverd.transpose(0, 2, 3, 1)
clean = clean.transpose(0, 2, 3, 1)
psnr = 0
ssim = 0
for i in range(recoverd.shape[0]):
print(f"Clean patch size: {clean[i].shape}, Restored size: {recoverd[i].shape}")
# psnr_val += compare_psnr(clean[i], recoverd[i])
# ssim += compare_ssim(clean[i], recoverd[i], multichannel=True)
psnr += peak_signal_noise_ratio(clean[i], recoverd[i], data_range=1)
ssim += structural_similarity(clean[i], recoverd[i], data_range=1, multichannel=True, win_size=3)
return psnr / recoverd.shape[0], ssim / recoverd.shape[0], recoverd.shape[0]
def compute_niqe(image):
image = np.clip(image.detach().cpu().numpy(), 0, 1)
image = image.transpose(0, 2, 3, 1)
niqe_val = niqe(image)
return niqe_val.mean()
class timer():
def __init__(self):
self.acc = 0
self.tic()
def tic(self):
self.t0 = time.time()
def toc(self):
return time.time() - self.t0
def hold(self):
self.acc += self.toc()
def release(self):
ret = self.acc
self.acc = 0
return ret
def reset(self):
self.acc = 0 |