def test_vif_simmular_to_matlab_implementation():
# Greyscale images
goldhill = torch.tensor(imread('tests/assets/goldhill.gif'))[None, None,
...]
goldhill_jpeg = torch.tensor(
imread('tests/assets/goldhill_jpeg.gif'))[None, None, ...]
score = vif_p(goldhill_jpeg, goldhill, data_range=255, reduction='none')
score_baseline = torch.tensor(0.2665)
assert torch.isclose(score, score_baseline, atol=1e-4), \
f'Expected PyTorch score to be equal to MATLAB prediction. Got {score} and {score_baseline}'
# RGB images
I01 = torch.tensor(imread('tests/assets/I01.BMP')).permute(2, 0, 1)[None,
...]
i1_01_5 = torch.tensor(imread('tests/assets/i01_01_5.bmp')).permute(
2, 0, 1)[None, ...]
score = vif_p(i1_01_5, I01, data_range=255, reduction='none')
# RGB images are not supported by MATLAB code. Here is result for luminance channel taken from YIQ colour space
score_baseline = torch.tensor(0.3147)
assert torch.isclose(score, score_baseline, atol=1e-4), \
f'Expected PyTorch score to be equal to MATLAB prediction. Got {score} and {score_baseline}'
def test_vif_fails_for_incorrect_data_range(prediction: torch.Tensor,
target: torch.Tensor,
device: str) -> None:
# Scale to [0, 255]
prediction_scaled = (prediction * 255).type(torch.uint8)
target_scaled = (target * 255).type(torch.uint8)
with pytest.raises(AssertionError):
vif_p(prediction_scaled.to(device),
target_scaled.to(device),
data_range=1.0)
def test_vif_supports_different_data_ranges(x, y, data_range, device: str) -> None:
x_scaled = (x * data_range).type(torch.uint8)
y_scaled = (y * data_range).type(torch.uint8)
measure_scaled = vif_p(x_scaled.to(device), y_scaled.to(device), data_range=data_range)
measure = vif_p(
x_scaled.to(device) / float(data_range),
y_scaled.to(device) / float(data_range),
data_range=1.0
)
diff = torch.abs(measure_scaled - measure)
assert diff <= 1e-5, f'Result for same tensor with different data_range should be the same, got {diff}'
def test_vif_supports_different_data_ranges(prediction: torch.Tensor,
target: torch.Tensor, data_range,
device: str) -> None:
prediction_scaled = (prediction * data_range).type(torch.uint8)
target_scaled = (target * data_range).type(torch.uint8)
measure_scaled = vif_p(prediction_scaled.to(device),
target_scaled.to(device),
data_range=data_range)
measure = vif_p(prediction_scaled.to(device) / float(data_range),
target_scaled.to(device) / float(data_range),
data_range=1.0)
diff = torch.abs(measure_scaled - measure)
assert diff <= 1e-5, f'Result for same tensor with different data_range should be the same, got {diff}'
def test_vif_p_works_for_zeros_tensors() -> None:
x = torch.zeros(4, 3, 256, 256)
y = torch.zeros(4, 3, 256, 256)
measure = vif_p(x, y, data_range=1.)
assert torch.isclose(measure, torch.tensor(
1.0)), f'VIF for 2 zero tensors shouls be 1.0, got {measure}.'
def test_vif_p_one_for_equal_tensors(x) -> None:
y = x.clone()
measure = vif_p(x, y)
assert torch.isclose(measure, torch.tensor(
1.0)), f'VIF for equal tensors shouls be 1.0, got {measure}.'
def test_vif_p(input_tensors: Tuple[torch.Tensor, torch.Tensor],
device: str) -> None:
x, y = input_tensors
vif_p(x.to(device), y.to(device), data_range=1.)
def main():
# Read RGB image and it's noisy version
x = torch.tensor(imread('tests/assets/i01_01_5.bmp')).permute(2, 0,
1) / 255.
y = torch.tensor(imread('tests/assets/I01.BMP')).permute(2, 0, 1) / 255.
if torch.cuda.is_available():
# Move to GPU to make computaions faster
x = x.cuda()
y = y.cuda()
# To compute BRISQUE score as a measure, use lower case function from the library
brisque_index: torch.Tensor = piq.brisque(x,
data_range=1.,
reduction='none')
# In order to use BRISQUE as a loss function, use corresponding PyTorch module.
# Note: the back propagation is not available using torch==1.5.0.
# Update the environment with latest torch and torchvision.
brisque_loss: torch.Tensor = piq.BRISQUELoss(data_range=1.,
reduction='none')(x)
print(
f"BRISQUE index: {brisque_index.item():0.4f}, loss: {brisque_loss.item():0.4f}"
)
# To compute Content score as a loss function, use corresponding PyTorch module
# By default VGG16 model is used, but any feature extractor model is supported.
# Don't forget to adjust layers names accordingly. Features from different layers can be weighted differently.
# Use weights parameter. See other options in class docstring.
content_loss = piq.ContentLoss(feature_extractor="vgg16",
layers=("relu3_3", ),
reduction='none')(x, y)
print(f"ContentLoss: {content_loss.item():0.4f}")
# To compute DISTS as a loss function, use corresponding PyTorch module
# By default input images are normalized with ImageNet statistics before forwarding through VGG16 model.
# If there is no need to normalize the data, use mean=[0.0, 0.0, 0.0] and std=[1.0, 1.0, 1.0].
dists_loss = piq.DISTS(reduction='none')(x, y)
print(f"DISTS: {dists_loss.item():0.4f}")
# To compute FSIM as a measure, use lower case function from the library
fsim_index: torch.Tensor = piq.fsim(x, y, data_range=1., reduction='none')
# In order to use FSIM as a loss function, use corresponding PyTorch module
fsim_loss = piq.FSIMLoss(data_range=1., reduction='none')(x, y)
print(
f"FSIM index: {fsim_index.item():0.4f}, loss: {fsim_loss.item():0.4f}")
# To compute GMSD as a measure, use lower case function from the library
# This is port of MATLAB version from the authors of original paper.
# In any case it should me minimized. Usually values of GMSD lie in [0, 0.35] interval.
gmsd_index: torch.Tensor = piq.gmsd(x, y, data_range=1., reduction='none')
# In order to use GMSD as a loss function, use corresponding PyTorch module:
gmsd_loss: torch.Tensor = piq.GMSDLoss(data_range=1., reduction='none')(x,
y)
print(
f"GMSD index: {gmsd_index.item():0.4f}, loss: {gmsd_loss.item():0.4f}")
# To compute HaarPSI as a measure, use lower case function from the library
# This is port of MATLAB version from the authors of original paper.
haarpsi_index: torch.Tensor = piq.haarpsi(x,
y,
data_range=1.,
reduction='none')
# In order to use HaarPSI as a loss function, use corresponding PyTorch module
haarpsi_loss: torch.Tensor = piq.HaarPSILoss(data_range=1.,
reduction='none')(x, y)
print(
f"HaarPSI index: {haarpsi_index.item():0.4f}, loss: {haarpsi_loss.item():0.4f}"
)
# To compute LPIPS as a loss function, use corresponding PyTorch module
lpips_loss: torch.Tensor = piq.LPIPS(reduction='none')(x, y)
print(f"LPIPS: {lpips_loss.item():0.4f}")
# To compute MDSI as a measure, use lower case function from the library
mdsi_index: torch.Tensor = piq.mdsi(x, y, data_range=1., reduction='none')
# In order to use MDSI as a loss function, use corresponding PyTorch module
mdsi_loss: torch.Tensor = piq.MDSILoss(data_range=1., reduction='none')(x,
y)
print(
f"MDSI index: {mdsi_index.item():0.4f}, loss: {mdsi_loss.item():0.4f}")
# To compute MS-SSIM index as a measure, use lower case function from the library:
ms_ssim_index: torch.Tensor = piq.multi_scale_ssim(x, y, data_range=1.)
# In order to use MS-SSIM as a loss function, use corresponding PyTorch module:
ms_ssim_loss = piq.MultiScaleSSIMLoss(data_range=1., reduction='none')(x,
y)
print(
f"MS-SSIM index: {ms_ssim_index.item():0.4f}, loss: {ms_ssim_loss.item():0.4f}"
)
# To compute Multi-Scale GMSD as a measure, use lower case function from the library
# It can be used both as a measure and as a loss function. In any case it should me minimized.
# By defualt scale weights are initialized with values from the paper.
# You can change them by passing a list of 4 variables to scale_weights argument during initialization
# Note that input tensors should contain images with height and width equal 2 ** number_of_scales + 1 at least.
ms_gmsd_index: torch.Tensor = piq.multi_scale_gmsd(x,
y,
data_range=1.,
chromatic=True,
reduction='none')
# In order to use Multi-Scale GMSD as a loss function, use corresponding PyTorch module
ms_gmsd_loss: torch.Tensor = piq.MultiScaleGMSDLoss(chromatic=True,
data_range=1.,
reduction='none')(x, y)
print(
f"MS-GMSDc index: {ms_gmsd_index.item():0.4f}, loss: {ms_gmsd_loss.item():0.4f}"
)
# To compute PSNR as a measure, use lower case function from the library.
psnr_index = piq.psnr(x, y, data_range=1., reduction='none')
print(f"PSNR index: {psnr_index.item():0.4f}")
# To compute PieAPP as a loss function, use corresponding PyTorch module:
pieapp_loss: torch.Tensor = piq.PieAPP(reduction='none', stride=32)(x, y)
print(f"PieAPP loss: {pieapp_loss.item():0.4f}")
# To compute SSIM index as a measure, use lower case function from the library:
ssim_index = piq.ssim(x, y, data_range=1.)
# In order to use SSIM as a loss function, use corresponding PyTorch module:
ssim_loss: torch.Tensor = piq.SSIMLoss(data_range=1.)(x, y)
print(
f"SSIM index: {ssim_index.item():0.4f}, loss: {ssim_loss.item():0.4f}")
# To compute Style score as a loss function, use corresponding PyTorch module:
# By default VGG16 model is used, but any feature extractor model is supported.
# Don't forget to adjust layers names accordingly. Features from different layers can be weighted differently.
# Use weights parameter. See other options in class docstring.
style_loss = piq.StyleLoss(feature_extractor="vgg16",
layers=("relu3_3", ))(x, y)
print(f"Style: {style_loss.item():0.4f}")
# To compute TV as a measure, use lower case function from the library:
tv_index: torch.Tensor = piq.total_variation(x)
# In order to use TV as a loss function, use corresponding PyTorch module:
tv_loss: torch.Tensor = piq.TVLoss(reduction='none')(x)
print(f"TV index: {tv_index.item():0.4f}, loss: {tv_loss.item():0.4f}")
# To compute VIF as a measure, use lower case function from the library:
vif_index: torch.Tensor = piq.vif_p(x, y, data_range=1.)
# In order to use VIF as a loss function, use corresponding PyTorch class:
vif_loss: torch.Tensor = piq.VIFLoss(sigma_n_sq=2.0, data_range=1.)(x, y)
print(f"VIFp index: {vif_index.item():0.4f}, loss: {vif_loss.item():0.4f}")
# To compute VSI score as a measure, use lower case function from the library:
vsi_index: torch.Tensor = piq.vsi(x, y, data_range=1.)
# In order to use VSI as a loss function, use corresponding PyTorch module:
vsi_loss: torch.Tensor = piq.VSILoss(data_range=1.)(x, y)
print(f"VSI index: {vsi_index.item():0.4f}, loss: {vsi_loss.item():0.4f}")
def test_vif_p_works_for_1_channel(prediction_1d: torch.Tensor,
target_1d: torch.Tensor) -> None:
vif_p(prediction_1d, target_1d, data_range=1.)
def test_vif_p_works_for_zeros_tensors() -> None:
prediction = torch.zeros(4, 3, 256, 256)
target = torch.zeros(4, 3, 256, 256)
measure = vif_p(prediction, target, data_range=1.)
assert torch.isclose(measure, torch.tensor(
1.0)), f'VIF for 2 zero tensors shouls be 1.0, got {measure}.'
def test_vif_p_one_for_equal_tensors(prediction: torch.Tensor) -> None:
target = prediction.clone()
measure = vif_p(prediction, target)
assert torch.isclose(measure, torch.tensor(
1.0)), f'VIF for equal tensors shouls be 1.0, got {measure}.'
def test_vif_p(input_tensors: Tuple[torch.Tensor, torch.Tensor],
device: str) -> None:
prediction, target = input_tensors
vif_p(prediction.to(device), target.to(device), data_range=1.)
def test_vif_preserves_dtype(x, y, dtype, device: str) -> None:
output = vif_p(x.to(device=device, dtype=dtype), y.to(device=device, dtype=dtype))
assert output.dtype == dtype
np2torch(gt_mag[-1],
"cuda").permute(2, 1, 0).unsqueeze(0)[:, 0:3, :, :],
data_range=1,
size_average=False))
bicub_ms_ssim.append(
ms_ssim(np2torch(bicub_mag[-1],
"cuda").permute(2, 1, 0).unsqueeze(0)[:, 0:3, :, :],
np2torch(gt_mag[-1],
"cuda").permute(2, 1, 0).unsqueeze(0)[:, 0:3, :, :],
data_range=1,
size_average=False))
singan_vif.append(
vif_p(np2torch(singan_mag[-1],
"cuda").permute(2, 1, 0).unsqueeze(0)[:, 0:3, :, :],
np2torch(gt_mag[-1],
"cuda").permute(2, 1, 0).unsqueeze(0)[:, 0:3, :, :],
data_range=1))
bicub_vif.append(
vif_p(np2torch(bicub_mag[-1],
"cuda").permute(2, 1, 0).unsqueeze(0)[:, 0:3, :, :],
np2torch(gt_mag[-1],
"cuda").permute(2, 1, 0).unsqueeze(0)[:, 0:3, :, :],
data_range=1))
imageio.mimwrite("singan.gif", singan_frames)
imageio.mimwrite("singan_mag.gif", singan_mag)
imageio.mimwrite("GT.gif", GT_frames)
imageio.mimwrite("gt_mag.gif", gt_mag)
imageio.mimwrite("bicub.gif", bicub_frames)
imageio.mimwrite("bicub_mag.gif", bicub_mag)
def test_vif_p_works_for_different_data_range(prediction: torch.Tensor,
target: torch.Tensor) -> None:
prediction_255 = (prediction * 255).type(torch.uint8)
target_255 = (target * 255).type(torch.uint8)
vif_p(prediction_255, target_255, data_range=255)
def test_vif_p_fails_for_small_images() -> None:
x = torch.rand(2, 3, 32, 32)
y = torch.rand(2, 3, 32, 32)
with pytest.raises(ValueError):
vif_p(x, y)
def test_vif_fails_for_incorrect_data_range(x, y, device: str) -> None:
# Scale to [0, 255]
x_scaled = (x * 255).type(torch.uint8)
y_scaled = (y * 255).type(torch.uint8)
with pytest.raises(AssertionError):
vif_p(x_scaled.to(device), y_scaled.to(device), data_range=1.0)
def test_vif_p_fails_for_small_images() -> None:
prediction = torch.rand(2, 3, 32, 32)
target = torch.rand(2, 3, 32, 32)
with pytest.raises(ValueError):
vif_p(prediction, target)
# plt.show()
psnr_n = psnr(img_sharp, img_deblu, data_range=255)
ssim_n = ssim(img_deblu / 255, img_sharp / 255, gaussian_weights=True, multichannel=True,
use_sample_covariance=False, sigma=1.5)
if name_sharp[-7:-4] == "001":
print(name_sharp, (psnr_n, ssim_n))
sharp = Image.fromarray(np.uint8(img_sharp))
deblu = Image.fromarray(np.uint8(img_deblu))
sharp_ts = TF.to_tensor(sharp).unsqueeze(0)
deblu_ts = TF.to_tensor(deblu).unsqueeze(0)
sharp_ts, deblu_ts = sharp_ts/255.0, deblu_ts/255.0
vif_n = piq.vif_p(deblu_ts, sharp_ts)
vsi_n = piq.vsi(deblu_ts, sharp_ts)
haar_n = piq.haarpsi(deblu_ts, sharp_ts)
#
# if count_k < 198:
# psnr_k.append(psnr_n)
# ssim_k.append(ssim_n)
# count_k += 1
# elif count_k == 198:
# psnr_k.append(psnr_n)
# ssim_k.append(ssim_n)
# kernel_p.append(max(psnr_k))
# kernel_s.append(max(ssim_k))
# psnr_k = []
# ssim_k = []
def test_vif_p_works_for_3_channels(prediction: torch.Tensor,
target: torch.Tensor) -> None:
vif_p(prediction, target, data_range=1.)
