小工具合集

import matplotlib.pyplot as plt
        # 可视化first conv layer filters
        plt.figure(figsize=(8 * 2, 8 * 2))
        cnt = 0
        for j in range(inp_enc_level1.size()[1]):
            cnt = cnt + 1
            plt.subplot(int(np.floor(np.sqrt(inp_enc_level1.size()[1]))), int(np.floor(np.sqrt(inp_enc_level1.size()[1]))), cnt)
            plt.imshow(inp_enc_level1[0][cnt - 1].detach().cpu().numpy(), cmap='gray')
        plt.show()

特征可视化

import matplotlib.pyplot as plt
        import numpy as np
        # 可视化first conv layer filters
        # plt.figure(figsize=(8 * 2, 8 * 2))
        cnt = 0
        x_show = []
        for j in range(x1.size()[1]):
            cnt = cnt + 1
            # plt.subplot(int(np.floor(np.sqrt(x1.size()[1]))),
            #             int(np.floor(np.sqrt(x1.size()[1]))), cnt)
            x_show.append(x1[0][cnt - 1].detach().cpu().numpy())
            #plt.imshow(x1[0][cnt - 1].detach().cpu().numpy(), cmap='gray')
        # plt.imshow(x1[0].detach().cpu().numpy(), cmap='gray')
        # print(x_show)
        plt.imshow(sum(x_show)/len(x_show), cmap='gray')
        plt.show()

求特征Mean

from thop import profile
if __name__ == '__main__':
    net = dehaze_net().cuda()
    # net = Deblur_L(num_feat=64, num_block=30).cuda()
    input_dim = torch.randn(1, 3, 256, 256).cuda()
    flops, params = profile(model=net, inputs=(input_dim,))
    print(flops / 10 ** 9, params / 10 ** 6)

计算参数和Flops

import numpy as np
import cv2
import math

def calculate_psnr(img1, img2, border=0):
    # img1 and img2 have range [0, 255]
    img1 = img1.squeeze()
    img2 = img2.squeeze()
    if not img1.shape == img2.shape:
        raise ValueError('Input images must have the same dimensions.')
    h, w = img1.shape[:2]
    img1 = img1[border:h-border, border:w-border]
    img2 = img2[border:h-border, border:w-border]

    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    mse = np.mean((img1 - img2)**2)
    if mse == 0:
        return float('inf')
    return 20 * math.log10(255.0 / math.sqrt(mse))


# --------------------------------------------
# SSIM
# --------------------------------------------
def calculate_ssim(img1, img2, border=0):
    '''calculate SSIM
    the same outputs as MATLAB's
    img1, img2: [0, 255]
    '''

    img1 = img1.squeeze()
    img2 = img2.squeeze()
    if not img1.shape == img2.shape:
        raise ValueError('Input images must have the same dimensions.')

    h, w = img1.shape[:2]
    img1 = img1[border:h-border, border:w-border]
    img2 = img2[border:h-border, border:w-border]

    if img1.ndim == 2:
        return ssim(img1, img2)
    elif img1.ndim == 3:
        if img1.shape[2] == 3:
            ssims = []
            for i in range(3):
                ssims.append(ssim(img1[:,:,i], img2[:,:,i]))
            return np.array(ssims).mean()
        elif img1.shape[2] == 1:
            return ssim(np.squeeze(img1), np.squeeze(img2))
    else:
        raise ValueError('Wrong input image dimensions.')


def ssim(img1, img2):
    C1 = (0.01 * 255)**2
    C2 = (0.03 * 255)**2

    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())

    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
    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(img1**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(img1 * 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()

if __name__ == '__main__':
	img_blur = Image.open(img_blurpath).convert('RGB')
	img_clear = Image.open(img_clearpath).convert('RGB')
	 ssim_ori = calculate_ssim(img_clear.squeeze().permute(1, 2, 0).cpu().numpy()*255, img_blur.squeeze().permute(1, 2, 0).cpu().numpy()*255,0)
        psnr_ori = calculate_psnr(img_clear.squeeze().permute(1, 2, 0).cpu().numpy()*255, img_blur.squeeze().permute(1, 2, 0).cpu().numpy()*255,0)

计算PSNR和SSIM