IIR-Lab/utils.py

import torch
import numpy as np
import torch.nn.functional as F
import cv2
def padding_4x(seq_noise):
    sh_im = seq_noise.size()
    expanded_h = sh_im[-2]%16

    if expanded_h:
        expanded_h = 16-expanded_h
    expanded_w = sh_im[-1]%16
    if expanded_w:
        expanded_w = 16-expanded_w

    padexp = (0, expanded_w, 0, expanded_h)
    seq_noise = F.pad(input=seq_noise, pad=padexp, mode='reflect')
    return seq_noise, expanded_h, expanded_w

def depadding(seq_denoise,expanded_h, expanded_w):
    if expanded_h:
        seq_denoise = seq_denoise[:, :, :-expanded_h, :]
    if expanded_w:
        seq_denoise = seq_denoise[:, :, :, :-expanded_w]
    return seq_denoise
def chunkV3(net, input_data, option, patch_h = 516, patch_w = 516, patch_h_overlap = 16, patch_w_overlap = 16):
    #input_data  [1,6,4,1500, 2000]

    # H = input_data.shape[3]
    # W = input_data.shape[4]

    shape_list = input_data.shape

    if option == 'image':
        B, C, H, W = shape_list[0], shape_list[1], shape_list[2], shape_list[3]                   # 1，4，1500，2000
    if option == 'RViDeformer':
        B, F, C, H, W = shape_list[0], shape_list[1], shape_list[2], shape_list[3], shape_list[4] # 1，6， 4，1500，2000
    if option == 'three2one':
        B, FC , H, W = shape_list[0], shape_list[1], shape_list[2], shape_list[3]                # 1，12，1500，2000
    
    if option == 'image':
        test_result = torch.zeros_like(input_data).cpu()  # 和input的shape一样
    if option == 'RViDeformer':
        test_result = torch.zeros_like(input_data).cpu()  # 和input的shape一样
    if option == 'three2one':
        test_result = torch.zeros((B, 4 , H, W)).cpu()  # 和input的shape一样


    # t0 = time.perf_counter()
    h_index = 1
    while (patch_h*h_index-patch_h_overlap*(h_index-1)) < H:
        if option == 'image':
            test_horizontal_result = torch.zeros((B,C,patch_h,W)).cpu()  #和input的shape一样 patch_h不一样
        if option == 'RViDeformer':
            test_horizontal_result = torch.zeros((B, F, C, patch_h, W)).cpu()
        if option == 'three2one':
            test_horizontal_result = torch.zeros((B, 4, patch_h, W)).cpu()

        h_begin = patch_h*(h_index-1)-patch_h_overlap*(h_index-1)
        h_end = patch_h*h_index-patch_h_overlap*(h_index-1) 
        w_index = 1
        while (patch_w*w_index-patch_w_overlap*(w_index-1)) < W:
            w_begin = patch_w*(w_index-1)-patch_w_overlap*(w_index-1)
            w_end = patch_w*w_index-patch_w_overlap*(w_index-1)
            test_patch = input_data[...,h_begin:h_end,w_begin:w_end]        

            with torch.no_grad():
                test_patch_result = net(test_patch).detach().cpu()

            if w_index == 1:
                test_horizontal_result[...,w_begin:w_end] = test_patch_result
            else:
                for i in range(patch_w_overlap):
                    test_horizontal_result[...,w_begin+i] = test_horizontal_result[...,w_begin+i]*(patch_w_overlap-1-i)/(patch_w_overlap-1)+test_patch_result[...,i]*i/(patch_w_overlap-1)
                test_horizontal_result[...,w_begin+patch_w_overlap:w_end] = test_patch_result[...,patch_w_overlap:]
            w_index += 1                   
    
        test_patch = input_data[...,h_begin:h_end,-patch_w:]        

        with torch.no_grad():
            test_patch_result = net(test_patch).detach().cpu()
        last_range = w_end-(W-patch_w)       

        for i in range(last_range):
            test_horizontal_result[...,W-patch_w+i] = test_horizontal_result[...,W-patch_w+i]*(last_range-1-i)/(last_range-1)+test_patch_result[...,i]*i/(last_range-1)
        test_horizontal_result[...,w_end:] = test_patch_result[...,last_range:]       

        if h_index == 1:
            test_result[...,h_begin:h_end,:] = test_horizontal_result
        else:
            for i in range(patch_h_overlap):
                test_result[...,h_begin+i,:] = test_result[...,h_begin+i,:]*(patch_h_overlap-1-i)/(patch_h_overlap-1)+test_horizontal_result[...,i,:]*i/(patch_h_overlap-1)
            test_result[...,h_begin+patch_h_overlap:h_end,:] = test_horizontal_result[...,patch_h_overlap:,:] 
        h_index += 1

    if option == 'image':
        test_horizontal_result = torch.zeros((B,C,patch_h,W)).cpu()  #和input的shape一样 patch_h不一样
    if option == 'RViDeformer':
        test_horizontal_result = torch.zeros((B, F, C, patch_h, W)).cpu()
    if option == 'three2one':
        test_horizontal_result = torch.zeros((B, 4, patch_h, W)).cpu()
        
    w_index = 1
    while (patch_w*w_index-patch_w_overlap*(w_index-1)) < W:
        w_begin = patch_w*(w_index-1)-patch_w_overlap*(w_index-1)
        w_end = patch_w*w_index-patch_w_overlap*(w_index-1)
        test_patch = input_data[...,-patch_h:,w_begin:w_end]            
              
        with torch.no_grad():
            test_patch_result = net(test_patch).detach().cpu()

        if w_index == 1:
            test_horizontal_result[...,w_begin:w_end] = test_patch_result
        else:
            for i in range(patch_w_overlap):
                test_horizontal_result[...,w_begin+i] = test_horizontal_result[...,w_begin+i]*(patch_w_overlap-1-i)/(patch_w_overlap-1)+test_patch_result[...,i]*i/(patch_w_overlap-1)
            test_horizontal_result[...,w_begin+patch_w_overlap:w_end] = test_patch_result[...,patch_w_overlap:]   
        w_index += 1

    test_patch = input_data[...,-patch_h:,-patch_w:]         

    with torch.no_grad():
        test_patch_result = net(test_patch).detach().cpu()
    last_range = w_end-(W-patch_w)       
    for i in range(last_range):
        test_horizontal_result[...,W-patch_w+i] = test_horizontal_result[...,W-patch_w+i]*(last_range-1-i)/(last_range-1)+test_patch_result[...,i]*i/(last_range-1) 
    test_horizontal_result[...,w_end:] = test_patch_result[...,last_range:] 

    last_last_range = h_end-(H-patch_h)
    for i in range(last_last_range):
        test_result[...,H-patch_w+i,:] = test_result[...,H-patch_w+i,:]*(last_last_range-1-i)/(last_last_range-1)+test_horizontal_result[...,i,:]*i/(last_last_range-1)
    test_result[...,h_end:,:] = test_horizontal_result[...,last_last_range:,:]
   
    # t1 = time.perf_counter()
    # print('Total running time: %s s' % (str(t1 - t0)))

    return test_result


def calculate_psnr(img, img2, input_order='HWC'):


    assert img.shape == img2.shape, (f'Image shapes are different: {img.shape}, {img2.shape}.')
    if input_order not in ['HWC', 'CHW']:
        raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are "HWC" and "CHW"')
    
    img = img.transpose(1, 2, 0)
    img2 = img2.transpose(1, 2, 0)


    img = img.astype(np.float64)
    img2 = img2.astype(np.float64)

    mse = np.mean((img - img2)**2)
    if mse == 0:
        return float('inf')
    return 10. * np.log10(1. * 1. / mse)


def calculate_ssim(img, img2, input_order='HWC'):


    assert img.shape == img2.shape, (f'Image shapes are different: {img.shape}, {img2.shape}.')
    if input_order not in ['HWC', 'CHW']:
        raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are "HWC" and "CHW"')


    img = img.transpose(1, 2, 0)
    img2 = img2.transpose(1, 2, 0)


    img = img.astype(np.float64)
    img2 = img2.astype(np.float64)

    ssims = []
    for i in range(img.shape[2]):
        ssims.append(_ssim(img[..., i], img2[..., i]))
    return np.array(ssims).mean()
    
def _ssim(img, img2):
    """Calculate SSIM (structural similarity) for one channel images.

    It is called by func:`calculate_ssim`.

    Args:
        img (ndarray): Images with range [0, 255] with order 'HWC'.
        img2 (ndarray): Images with range [0, 255] with order 'HWC'.

    Returns:
        float: SSIM result.
    """

    c1 = (0.01 * 1)**2
    c2 = (0.03 * 1)**2
    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())

    mu1 = cv2.filter2D(img, -1, window)[5:-5, 5:-5]  # valid mode for window size 11
    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
    mu1_sq = mu1**2
    mu2_sq = mu2**2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = cv2.filter2D(img**2, -1, window)[5:-5, 5:-5] - mu1_sq
    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
    sigma12 = cv2.filter2D(img * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

    ssim_map = ((2 * mu1_mu2 + c1) * (2 * sigma12 + c2)) / ((mu1_sq + mu2_sq + c1) * (sigma1_sq + sigma2_sq + c2))
    return ssim_map.mean()