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.)