代码整理: PyTorch DL AI 模型测试

深度学习模型测试代码

个人看过觉得比较合适的代码部分记录于此，以后一些部分的代码抄就完事了。随缘更新

1. Multi-Stage Progressive Image Restoration (CVPR 2021) demo.py

Code:https://github.com/swz30/MPRNet

2022/6/5:图像增强方面的论文，输入数据都是图像格式。代码简洁明了，针对其他任务则添加关于模型路径的参数，修改模型读取处的代码；有需要则添加计算运行时间函数，输出运行平均时间。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms.functional as TF
from PIL import Image
import os
from runpy import run_path
from skimage import img_as_ubyte
from collections import OrderedDict
from natsort import natsorted
from glob import glob
import cv2
import argparse

# 输入路径，输出路径，所选择任务
parser = argparse.ArgumentParser(description='Demo MPRNet')
parser.add_argument('--input_dir', default='./samples/input/', type=str, help='Input images')
parser.add_argument('--result_dir', default='./samples/output/', type=str, help='Directory for results')
parser.add_argument('--task', required=True, type=str, help='Task to run', choices=['Deblurring', 'Denoising', 'Deraining'])

args = parser.parse_args()

def save_img(filepath, img):
	# 保存图像
    cv2.imwrite(filepath,cv2.cvtColor(img, cv2.COLOR_RGB2BGR))

def load_checkpoint(model, weights):
	# 加载权重
    checkpoint = torch.load(weights)
    try:
        model.load_state_dict(checkpoint["state_dict"])
    except:
        state_dict = checkpoint["state_dict"]
        new_state_dict = OrderedDict()
        for k, v in state_dict.items():
            name = k[7:] # remove `module.`
            new_state_dict[name] = v
        model.load_state_dict(new_state_dict)

task    = args.task
inp_dir = args.input_dir
out_dir = args.result_dir

# exist_ok=True,则目标路径存在不会触发异常
os.makedirs(out_dir, exist_ok=True)

# 自然排序，后缀相同时自然排序
files = natsorted(glob(os.path.join(inp_dir, '*.jpg'))
                + glob(os.path.join(inp_dir, '*.JPG'))
                + glob(os.path.join(inp_dir, '*.png'))
                + glob(os.path.join(inp_dir, '*.PNG')))

# 找不到文件抛出错误
if len(files) == 0:
    raise Exception(f"No files found at {inp_dir}")

# Load corresponding model architecture and weights
load_file = run_path(os.path.join(task, "MPRNet.py"))
# 把MPRNet.py文件中的MPRNet类取出来了
model = load_file['MPRNet']()
model.cuda()

# 读取参数路径，并设为eval
weights = os.path.join(task, "pretrained_models", "model_"+task.lower()+".pth")
load_checkpoint(model, weights)
model.eval()

# 因为网络结构中使用了U-Net结构，所以对于输入都要进行调整
img_multiple_of = 8

for file_ in files:
    img = Image.open(file_).convert('RGB')
    input_ = TF.to_tensor(img).unsqueeze(0).cuda()

    # Pad the input if not_multiple_of 8
    # 如果输入图像不满足8的倍数则padding
    h,w = input_.shape[2], input_.shape[3]
    H,W = ((h+img_multiple_of)//img_multiple_of)*img_multiple_of, ((w+img_multiple_of)//img_multiple_of)*img_multiple_of
    padh = H-h if h%img_multiple_of!=0 else 0
    padw = W-w if w%img_multiple_of!=0 else 0
    input_ = F.pad(input_, (0,padw,0,padh), 'reflect')

    # 得到输出，限制其范围[min,max]
    with torch.no_grad():
        restored = model(input_)
    restored = restored[0]  # 有多个阶段的输出，0号位置的是最终输出
    restored = torch.clamp(restored, 0, 1)

    # Unpad the output
    # 输出的图像又去掉padding的部分
    restored = restored[:,:,:h,:w]

    # [B,H,W,C]
    restored = restored.permute(0, 2, 3, 1).cpu().detach().numpy()
    # 将图像转换为8位无符号整数格式
    restored = img_as_ubyte(restored[0])

    # 先拆开前面的目录取出文件名，然后去除文件格式
    # "A/B/c.jpg"-->"c.jpg"-->"c"
    f = os.path.splitext(os.path.split(file_)[-1])[0]
    # 保存文件
    save_img((os.path.join(out_dir, f+'.png')), restored)

# 输出表示结束
print(f"Files saved at {out_dir}")

2. FFA-Net: Feature Fusion Attention Network for Single Image Dehazing (AAAI 2020) test.py

Code:https://github.com/zhilin007/FFA-Net

2022/6/6:FFA-Net是一篇图像去雾方向的论文，这篇论文曾经复现并进行改进，虽然在RESIDE数据集上有着非常好的指标，但在真实师姐场景几乎是不起作用的，或者说只有对近景的薄雾有一丢丢作用。代码的命名略显粗糙，tensorShow值得Copy一下。

import os,argparse
import numpy as np
from PIL import Image
from models import *
import torch
import torch.nn as nn
import torchvision.transforms as tfs 
import torchvision.utils as vutils
import matplotlib.pyplot as plt
from torchvision.utils import make_grid
# 当前工作目录绝对路径
abs=os.getcwd()+'/'
# 输入输出图像展示
def tensorShow(tensors,titles=['haze']):
        fig=plt.figure()
        for tensor,tit,i in zip(tensors,titles,range(len(tensors))):
            img = make_grid(tensor)
            npimg = img.numpy()
            ax = fig.add_subplot(221+i)
            ax.imshow(np.transpose(npimg, (1, 2, 0)))
            ax.set_title(tit)
        plt.show()

# 参数设置
parser=argparse.ArgumentParser()
parser.add_argument('--task',type=str,default='its',help='its or ots')
parser.add_argument('--test_imgs',type=str,default='test_imgs',help='Test imgs folder')
opt=parser.parse_args()
dataset=opt.task

# 模型的超参数
gps=3
blocks=19
# 测试图像路径和输出路径
img_dir=abs+opt.test_imgs+'/'
output_dir=abs+f'pred_FFA_{dataset}/'
print("pred_dir:",output_dir)

# 创建输出路径
if not os.path.exists(output_dir):
    os.mkdir(output_dir)

# 模型参数路径
model_dir=abs+f'trained_models/{dataset}_train_ffa_{gps}_{blocks}.pk'

# 加载模型，加载参数，调整为eval()
device='cuda' if torch.cuda.is_available() else 'cpu'
ckp=torch.load(model_dir,map_location=device)
net=FFA(gps=gps,blocks=blocks)
net=nn.DataParallel(net)
net.load_state_dict(ckp['model'])
net.eval()

# 读取路径中的图片
for im in os.listdir(img_dir):
    print(f'\r {im}',end='',flush=True)
    haze = Image.open(img_dir+im)
    haze1= tfs.Compose([
        tfs.ToTensor(),
        tfs.Normalize(mean=[0.64, 0.6, 0.58],std=[0.14,0.15, 0.152])
    ])(haze)[None,::]
    haze_no=tfs.ToTensor()(haze)[None,::]
    with torch.no_grad():
        pred = net(haze1)
    ts=torch.squeeze(pred.clamp(0,1).cpu())
    # 每一对输入输出图像都会展示出来
    tensorShow([haze_no,pred.clamp(0,1).cpu()],['haze','pred'])
    # "A/B/c.jpg"-->"A/B/c"-->"A/B/c_FFA.png"
    vutils.save_image(ts,output_dir+im.split('.')[0]+'_FFA.png')

3. COIN: COmpression with Implicit Neural representations 绘图代码 plots.py

Code:https://github.com/EmilienDupont/coin

# Based on https://github.com/InterDigitalInc/CompressAI/blob/master/compressai/utils/plot/__main__.py
import imageio
import json5 as json
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
from matplotlib import cm
from pathlib import Path


ours = 'COIN'

# Ensure consistent coloring across plots
# 为每个折线选择不同的颜色
name_to_color = {
    ours: mcolors.TABLEAU_COLORS['tab:blue'],
    'BMS': mcolors.TABLEAU_COLORS['tab:orange'],
    'MBT': mcolors.TABLEAU_COLORS['tab:green'],
    'CST': mcolors.TABLEAU_COLORS['tab:red'],
    'JPEG': mcolors.TABLEAU_COLORS['tab:purple'],
    'JPEG2000': mcolors.TABLEAU_COLORS['tab:brown'],
    'BPG': mcolors.TABLEAU_COLORS['tab:pink'],
    'VTM': mcolors.TABLEAU_COLORS['tab:gray'],
}

# Setup colormap for residuals plot
# 
viridis = cm.get_cmap('viridis', 100)


def parse_json_file(filepath, metric='psnr'):
    """Parses a json result file.

    Args:
        filepath (string): Path to results json file.
        metric (string): Metric to use for plot.
    """
    # 路径
    filepath = Path(filepath)
    name = filepath.name.split('.')[0]
    # 读取
    with filepath.open('r') as f:
        try:
            data = json.load(f)
        except json.decoder.JSONDecodeError as err:
            print(f'Error reading file {filepath}')
            raise err
	
	# 确保json文件格式没问题，这些文件都能再coin-main/baselines里面找到
    if 'results' not in data or 'bpp' not in data['results']:
        raise ValueError(f'Invalid file {filepath}')
	
	# 也是确保格式没问题
    if metric not in data['results']:
        raise ValueError(
            f'Error: metric {metric} not available.'
            f' Available metrics: {", ".join(data["results"].keys())}'
        )
	
	#查看ms-ssim的话转成db
    if metric == 'ms-ssim':
        # Convert to db
        values = np.array(data['results'][metric])
        data['results'][metric] = -10 * np.log10(1 - values)
	
	# 返回name,bpp和metric
    return {
        'name': data.get('name', name),
        'xs': data['results']['bpp'],
        'ys': data['results'][metric],
    }


def rate_distortion(scatters, title=None, ylabel='PSNR [dB]', output_file=None,
                    limits=None, show=False, figsize=None):
    """Creates a rate distortion plot based on scatters.

    Args:
        scatters (list of dicts): List of data to plot for each model.
        title (string):
        ylabel (string):
        output_file (string): If not None, save plot at output_file.
        limits (tuple of ints):
        show (bool): If True shows plot.
        figsize (tuple of ints):
    """
    if figsize is None:
        figsize = (7, 4)
    fig, ax = plt.subplots(figsize=figsize)
    for sc in scatters:
        if sc['name'] == ours:
            linewidth = 2.5
            markersize = 10
        else:
            linewidth = 1
            markersize = 6

        if sc['name'] in [ours, 'BMS', 'MBT', 'CST']:
            pattern = '.-'  # Learned algorithms
        else:
            pattern = '.--'  # Non learned algorithms
        ax.plot(sc['xs'], sc['ys'], pattern, label=sc['name'],
                c=name_to_color[sc['name']], linewidth=linewidth,
                markersize=markersize)

    ax.set_xlabel('Bit-rate [bpp]')
    ax.set_ylabel(ylabel)
    ax.grid()
    if limits is not None:
        ax.axis(limits)
    ax.legend(loc='lower right')

    if title:
        ax.title.set_text(title)

    if show:
        plt.show()

    if output_file:
        fig.savefig(output_file, dpi=300, bbox_inches='tight')
        plt.clf()
        plt.close()


def plot_rate_distortion(filepaths=['results.json',
                                    'baselines/compressai-bmshj2018-hyperprior.json',
                                    'baselines/compressai-mbt2018.json',
                                    'baselines/compressai-cheng2020-anchor.json',
                                    'baselines/jpeg.json', 'baselines/jpeg2000.json',
                                    'baselines/bpg_444_x265_ycbcr.json',
                                    'baselines/vtm.json'],
                         output_file=None, limits=None):
    """Creates rate distortion plot based on all results json files.

    Args:
        filepaths (list of string): List of paths to result json files.
        output_file (string): Path to save image.
        limits (tuple of float): Limits of plot.
    """
    # Read data
    scatters = []
    for f in filepaths:
        rv = parse_json_file(f, 'psnr')
        scatters.append(rv)
    # Create plot
    rate_distortion(scatters, output_file=output_file, limits=limits)


def plot_model_size(output_file=None, show=False):
    """Plots histogram of model sizes.

    Args:
        output_file (string): If not None, save plot at output_file.
        show (bool): If True shows plot.

    Notes:
        Data for all baselines was computed using the compressAI library
        https://github.com/InterDigitalInc/CompressAI
    """
    model_names = ['COIN', 'BMS', 'MBT', 'CST']
    model_sizes = [14.7455, 10135.868, 24764.604, 31834.464]  # in kB

    plt.grid(zorder=0, which="both", axis="y")  # Ensure grid is at the back

    barplot = plt.bar(model_names, model_sizes, log=True, zorder=10)
    for i in range(len(model_names)):
        barplot[i].set_color(name_to_color[model_names[i]])
    plt.ylabel("Model size [kB]")

    fig = plt.gcf()
    fig.set_size_inches(3, 4)

    if show:
        plt.show()

    if output_file:
        plt.savefig(output_file, format='png', dpi=400, bbox_inches='tight')
        plt.clf()
        plt.close()


def plot_residuals(path_original='kodak-dataset/kodim15.png',
                   path_coin='imgs/kodim15_coin_bpp_03.png',
                   path_jpeg='imgs/kodim15_jpeg_bpp_03.jpg',
                   output_file=None, show=False, max_residual=0.3,
                   title_fontsize=6):
    """Creates a plot comparing compression with COIN and JPEG both in terms of
    the compressed image and the residual between the compressed and original
    image.


    Args:
        path_original (string): Path to original image.
        path_coin (string): Path to image compressed with COIN.
        path_jpeg (string): Path to image compressed with JPEG.
        output_file (string): If not None, save plot at output_file.
        show (bool): If True shows plot.
        max_residual (float): Value between 0 and 1 to use for maximum residual
            on color scale. Usually set to a low value so residuals are clearer
            on plot.
    """
    # Load images and compute residuals
    img_original = imageio.imread(path_original) / 255.
    img_coin = imageio.imread(path_coin) / 255.
    img_jpeg = imageio.imread(path_jpeg) / 255.
    residual_coin = viridis(np.abs(img_coin - img_original).mean(axis=-1) / max_residual)[:, :, :3]
    residual_jpeg = viridis(np.abs(img_jpeg - img_original).mean(axis=-1) / max_residual)[:, :, :3]

    # Create plot
    plt.subplot(2, 3, 1)
    plt.imshow(img_original)
    plt.axis('off')
    plt.gca().set_title('Original', fontsize=title_fontsize)

    plt.subplot(2, 3, 2)
    plt.imshow(img_coin)
    plt.axis('off')
    plt.gca().set_title('COIN', fontsize=title_fontsize)

    plt.subplot(2, 3, 3)
    plt.imshow(residual_coin)
    plt.axis('off')
    plt.gca().set_title('COIN Residual', fontsize=title_fontsize)

    plt.subplot(2, 3, 5)
    plt.imshow(img_jpeg)
    plt.axis('off')
    plt.gca().set_title('JPEG', fontsize=title_fontsize)

    plt.subplot(2, 3, 6)
    plt.imshow(residual_jpeg)
    plt.axis('off')
    plt.gca().set_title('JPEG Residual', fontsize=title_fontsize)

    plt.subplots_adjust(wspace=0.1, hspace=0)

    if show:
        plt.show()

    if output_file:
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        plt.clf()
        plt.close()


if __name__ == '__main__':
    plot_rate_distortion(output_file='rate_distortion.png',
                         limits=(0, 1, 22, 38))
    plot_model_size(output_file='model_sizes.png')
    plot_residuals(output_file='residuals_kodim15_bpp_03.png')
    plot_residuals(output_file='residuals_kodim15_bpp_015.png',
                   path_coin='imgs/kodim15_coin_bpp_015.png',
                   path_jpeg='imgs/kodim15_jpeg_bpp_015.jpg')