【代码】将图片数据集转换为h5格式的数据集

  • 写这篇博客的起因是我在学习SRCNN(这里面给出了代码地址)这篇论文后,准备跑一跑其中的模型。在代码中所用到的数据集是h5格式的文件,并在代码中给出了数据集的下载链接。是在www.dropbox.com这个网址中,进不去,花了半天时间,各种方法试过了还是进不去。放弃了,花这时间还不如自己写一个代码生成h5格式的数据集文件。
  • 下面给出代码:
import glob
import h5py
import numpy as np
import PIL.Image as pImg

# 将rgb图像转换为灰度图像
def rgb2gray(img):
    return 16. + (64.738 * img[:, :, 0] + 129.057 * img[:, :, 1] + 25.064 * img[:, :, 2]) / 256.


def setTrianData(imgPath, h5Path, scale=3, pSize=33, pStride=14):
    """
    将图片格式的数据集转换为h5格式的数据集
    :param imgPath: 图片格式数据集的存储路径
    :param h5Path: h5格式数据集的存储路径
    :param scale: 图片尺度x2,x3,x4
    :param pSize: 采样子图像的尺寸大小,按照SRCNN原文默认为33
    :param pStride: 采样子图像的步幅大小,按照SRCNN原文默认为14
    :return:
    """
    # 用于存储低分辨率和高分辨率的采样子图像
    lr_sub_images, hr_sub_images = [], []
    for p in sorted(glob.glob(f'{imgPath}/*.*')):
        hr = pImg.open(p).convert('RGB')
        # 重新规划图像尺寸,使图像可以被scale整除,分别得到低分辨率图像的尺寸和高分辨率的尺寸
        lrWidth, lrHeight = hr.width // scale, hr.height // scale
        width, height = lrWidth*scale, lrHeight*scale

        # 将原始图像通过双三次插值重设尺寸,使之可被scale整除,作为高分辨图像数据HR
        hr = hr.resize((width, height), resample=pImg.BICUBIC)

        # 将HR通过双三次插值压缩scale倍,为低分辨图像的原始数据
        lr = hr.resize((lrWidth, lrHeight), resample=pImg.BICUBIC)

        # 将低分辨图像通过双三次插值放大scale倍,与HR图像维度相等,作为低分辨图像数据LR
        lr = lr.resize((width, height), resample=pImg.BICUBIC)

        hr = np.array(hr).astype(np.float32)
        lr = np.array(lr).astype(np.float32)
        hr = rgb2gray(hr)
        lr = rgb2gray(lr)

        # 按照33*33,步距为14进行采样
        for i in range(0, height - pSize + 1, pStride):
            for j in range(0, width - pSize + 1, pStride):
                lr_sub_images.append(lr[i:i + pSize, j:j + pSize])
                hr_sub_images.append(hr[i:i + pSize, j:j + pSize])
    h5_file = h5py.File(h5Path, 'w')
    h5_file.create_dataset('lr', data=np.array(lr_sub_images))
    h5_file.create_dataset('hr', data=np.array(hr_sub_images))
    h5_file.close()

if __name__ == '__main__':
    setTrianData(imgPath=".\T_91image", h5Path=".\T91image.h5")
  • 如下图所示为我的文件目录,其中T_91image文件夹下存放的是图片格式的数据集,运行上述代码将会在当前文件夹生成一个h5格式的文件。
    【代码】将图片数据集转换为h5格式的数据集_第1张图片
  • 生成的h5格式数据集就可以直接用于训练了

你可能感兴趣的:(图像处理,代码,计算机视觉,深度学习,opencv)