图像数据管理DISK、LMDB和HDF5的存储与读取

随着数据量的不断增长，如何有效地存储和读取大量图像数据成为了数据处理和分析中的重要挑战。尤其在机器学习和深度学习领域，处理海量图像数据是一个常见需求。

为了应对这种情况，本文将探讨三种常见的图像存储与读取方式：DISK、LMDB 和 HDF5。通过对单一图像和多图像的存储与读取的讨论，以及对不同存储方式的性能比较，帮助理解每种方法的优缺点，并为实际应用提供参考。无论是在小规模数据处理还是在大规模数据处理场景下，合理选择存储方式可以显著提高数据处理效率。

数据准备

CIFAR-10数据集 由60000个32x32像素的彩色图像组成，涵盖不同的对象类别，如狗、猫和飞机等。虽然CIFAR-10不算是一个非常大的数据集，但使用完整的TinyImages数据集可能需要约400GB的可用磁盘空间。

文中的代码应用的数据集下载地址 CIFAR-10 数据集。

图像在Python中的处理通常需要将其加载到NumPy数组中，以便应用例如卷积神经网络（CNN）等算法。以下是如何将图像加载到NumPy数组中的示例代码：

import numpy as np
import pickle
from pathlib import Path

# 数据文件夹路径
data_dir = Path("Datasets/cifar-10-batches-py/")

# 解码函数
def unpickle(file):
    with open(file, "rb") as fo:
        dict = pickle.load(fo, encoding="bytes")
    return dict

images, labels = [], []
for batch in data_dir.glob("data_batch_*"):
    batch_data = unpickle(batch)
    for i, flat_im in enumerate(batch_data[b"data"]):
        im_channels = []
        # 每个图像都是扁平化的，通道按 R, G, B 的顺序排列  
        for j in range(3):
            im_channels.append(
                flat_im[j * 1024 : (j + 1) * 1024].reshape((32, 32))
            )
        # 重建原始图像
        images.append(np.dstack((im_channels)))
        # 保存标签
        labels.append(batch_data[b"labels"][i])

print("加载 CIFAR-10 训练集:")
print(f" - np.shape(images)     {np.shape(images)}")
print(f" - np.shape(labels)     {np.shape(labels)}")

输出结果中我们看到了图片的数量有50000张，像素是32*32，色道数量为3彩色。

加载 CIFAR-10 训练集:
 - np.shape(images)     (50000, 32, 32, 3)
 - np.shape(labels)     (50000,)

图像数据应用

图像存储的设置

在进行图像存储之前，需要安装一些第三方库，用于图像的不同存储方法。我们将使用Pillow库进行图像处理，并使用lmdb和h5py库来实现不同的存储方式。

安装Pillow库：

pip install Pillow

LMDB（闪电内存映射数据库）通过键值存储数据，适合大规模数据处理，使用前安装lmdb库：

pip install lmdb

HDF5（分层数据格式）是一种紧凑的科学数据格式，支持大规模数据存储，使用前安装h5py库：

pip install h5py

单一图像存储

在此部分，展示如何将单个图像存储为不同格式：DISK、LMDB 和 HDF5。首先定义存储路径：

from pathlib import Path

disk_dir = Path("Datasets/cifar-10-batches-py/disk/")
lmdb_dir = Path("Datasets/cifar-10-batches-py/lmdb/")
hdf5_dir = Path("Datasets/cifar-10-batches-py/hdf5/")

DISK存储

DISK存储将图像以.png格式保存到磁盘上，并使用唯一的图像ID命名，同时将图像的标签保存到.csv文件中。

from PIL import Image
import csv

def store_single_disk(image, image_id, label):
    """ 将单个图像作为 .png 文件存储在磁盘上。"""
    Image.fromarray(image).save(disk_dir / f"{image_id}.png")

    with open(disk_dir / f"{image_id}.csv", "wt") as csvfile:
        writer = csv.writer(
            csvfile, delimiter=" ", quotechar="|", quoting=csv.QUOTE_MINIMAL
        )
        writer.writerow([label])

LMDB存储

在LMDB存储中，每个条目保存为一个字节数组，键为图像ID，值为图像和标签的封装对象。

import lmdb
import pickle

class CIFAR_Image:
    def __init__(self, image, label):
        self.channels = image.shape[2]
        self.size = image.shape[:2]
        self.image = image.tobytes()
        self.label = label

    def get_image(self):
        """ 返回图像作为 numpy 数组 """
        image = np.frombuffer(self.image, dtype=np.uint8)
        return image.reshape(*self.size, self.channels)

def store_single_lmdb(image, image_id, label):
    """ 将单个图像存储到 LMDB """
    map_size = image.nbytes * 10

    env = lmdb.open(str(lmdb_dir / "single_lmdb"), map_size=map_size)

    with env.begin(write=True) as txn:
        value = CIFAR_Image(image, label)
        key = f"{image_id:08}"
        txn.put(key.encode("ascii"), pickle.dumps(value))
    env.close()

HDF5存储

在HDF5存储中，图像和标签以数据集形式存储在HDF5文件中。

import h5py

def store_single_hdf5(image, image_id, label):
    """ 将单个图像存储到 HDF5 文件 """
    file = h5py.File(hdf5_dir / f"{image_id}.h5", "w")
    file.create_dataset("image", np.shape(image), h5py.h5t.STD_U8BE, data=image)
    file.create_dataset("meta", np.shape(label), h5py.h5t.STD_U8BE, data=label)
    file.close()

存储方式对比

通过比较不同存储方式的存储时间，可以评估其效率和适用性。

以下代码展示了如何记录每种存储方法的耗时：

from timeit import timeit

_store_single_funcs = dict(
    disk=store_single_disk, 
    lmdb=store_single_lmdb, 
    hdf5=store_single_hdf5
)

store_single_timings = dict()

for method in ("disk", "lmdb", "hdf5"):
    t = timeit(
        "_store_single_funcs[method](image, 0, label)",
        setup="image=images[0]; label=labels[0]",
        number=1,
        globals=globals(),
    )
    store_single_timings[method] = t
    print(f"存储方法: {method}, 使用耗时: {t}")

以下是比较三种存储方法的表格：

存储方法	存储耗时	使用内存
Disk	0.011781秒	8 KB
LMDB	0.001933秒	32 KB
HDF5	0.001986秒	8 KB

通过对比可以发现，LMDB和HDF5存储方法的效率更高，尤其是在大数据量存储时表现更为出色。

多图像的存储

与单个图像的存储类似，多个图像的存储可以通过修改代码进行批量存储操作。在DISK存储中，多个图像保存为多个.png文件；对于LMDB和HDF5存储，可以将多个图像存储在单个数据库或文件中。

DISK存储

在DISK存储方式中，可以通过循环将图像逐个保存到磁盘，同时将对应的标签保存到一个.csv文件中。

from PIL import Image
import csv

def store_many_disk(images, labels):
    """存储多个图像到磁盘，每张图像单独存储为 PNG 文件，并将标签存入 CSV 文件。"""
    num_images = len(images)

    for i, image in enumerate(images):
        Image.fromarray(image).save(disk_dir / f"{i}.png")

    with open(disk_dir / f"{num_images}_labels.csv", "w") as csvfile:
        writer = csv.writer(csvfile, delimiter=" ", quotechar="|", quoting=csv.QUOTE_MINIMAL)
        for label in labels:
            writer.writerow([label])

LMDB存储

在LMDB中，将多个图像作为键值对存储在同一个数据库中，键为图像的唯一ID，值为图像和标签的封装对象。

def store_many_lmdb(images, labels):
    """存储多个图像到 LMDB 数据库。"""
    num_images = len(images)
    map_size = num_images * images[0].nbytes * 10

    env = lmdb.open(str(lmdb_dir / f"{num_images}_lmdb"), map_size=map_size)

    with env.begin(write=True) as txn:
        for i in range(num_images):
            value = CIFAR_Image(images[i], labels[i])
            key = f"{i:08}"
            txn.put(key.encode("ascii"), pickle.dumps(value))
    env.close()

HDF5存储

在HDF5中，可以将所有图像和标签以数据集的形式存储在同一个HDF5文件中。

import h5py

def store_many_hdf5(images, labels):
    """存储多个图像到 HDF5 文件。"""
    num_images = len(images)

    file = h5py.File(hdf5_dir / f"{num_images}_many.h5", "w")
    file.create_dataset("images", np.shape(images), h5py.h5t.STD_U8BE, data=images)
    file.create_dataset("meta", np.shape(labels), h5py.h5t.STD_U8BE, data=labels)
    file.close()

存储方式对比

在不同图像数量下，使用DISK、LMDB、HDF5存储方法进行比较，以评估存储效率。以下代码展示了如何记录不同存储方式的耗时：

from timeit import timeit

_store_many_funcs = dict(
    disk=store_many_disk, lmdb=store_many_lmdb, hdf5=store_many_hdf5
)

store_many_timings = {"disk": [], "lmdb": [], "hdf5": []}

cutoffs = [10, 100, 1000, 10000, 100000]

for cutoff in cutoffs:
    for method in ("disk", "lmdb", "hdf5"):
        t = timeit(
            "_store_many_funcs[method](images_, labels_)",
            setup="images_=images[:cutoff]; labels_=labels[:cutoff]",
            number=1,
            globals=globals(),
        )
        store_many_timings[method].append(t)
        print(f"存储方法: {method}, 图像数量: {cutoff}, 使用耗时: {t}")

以下为各存储方法的耗时对比表：

存储方法	存储 10 张图像耗时	存储 100 张图像耗时	存储 1000 张图像耗时	存储 10000 张图像耗时	存储 100000 张图像耗时
DISK	0.00655秒	0.05839秒	0.53844秒	5.32882秒	53.71001秒
LMDB	0.00204秒	0.00463秒	0.04158秒	0.25626秒	2.38738秒
HDF5	0.00145秒	0.00144秒	0.00257秒	0.01716秒	0.24997秒

从上表可以看出，随着图像数量增加，DISK存储的耗时显著增加，而LMDB和HDF5存储方法在处理大量图像时效率更高。

单一图像的读取

图像读取方式与存储方式类似，分别讨论使用DISK、LMDB和HDF5的读取方法。

DISK读取

在DISK存储中，图像以.png文件保存，通过Pillow库读取，同时从.csv文件中读取对应的标签。

from PIL import Image
import csv

def read_single_disk(image_id):
    """从磁盘读取单个图像及其标签。"""
    image = np.array(Image.open(disk_dir / f"{image_id}.png"))

    with open(disk_dir / f"{image_id}.csv", "r") as csvfile:
        reader = csv.reader(csvfile, delimiter=" ", quotechar="|", quoting=csv.QUOTE_MINIMAL)
        label = int(next(reader)[0])

    return image, label

LMDB读取

在LMDB中，图像和标签作为字节数组存储，读取时需要通过LMDB数据库获取存储的键值对，并解码成图像对象。

def read_single_lmdb(image_id):
    """从 LMDB 数据库读取单个图像及其标签。"""
    env = lmdb.open(str(lmdb_dir / "single_lmdb"), readonly=True)

    with env.begin() as txn:
        data = txn.get(f"{image_id:08}".encode("ascii"))
        cifar_image = pickle.loads(data)
        image = cifar_image.get_image()
        label = cifar_image.label
    env.close()

    return image, label

HDF5读取

在HDF5中，图像和标签作为数据集存储，读取时直接从对应的文件中检索图像和标签数据。

import h5py

def read_single_hdf5(image_id):
    """从 HDF5 文件读取单个图像及其标签。"""
    file = h5py.File(hdf5_dir / f"{image_id}.h5", "r+")
    image = np.array(file["/image"]).astype("uint8")
    label = int(np.array(file["/meta"]).astype("uint8"))
    file.close()

    return image, label

读取方式对比

与存储方式类似，通过记录不同读取方法的耗时来比较其性能。

_read_single_funcs = dict(
    disk=read_single_disk, lmdb=read_single_lmdb, hdf5=read_single_hdf5
)

read_single_timings = {"disk": [], "lmdb": [], "hdf5": []}

for method in ("disk", "lmdb", "hdf5"):
    t = timeit(
        "_read_single_funcs ",
        setup="image=images[0]; label=labels[0]",
        number=1,
        globals=globals(),
    )
    read_single_timings[method] = t
    print(f"读取方法: {method}, 使用耗时: {t}")

以下为各读取方法的性能对比表：

读取方法	耗时（秒）
Disk	0.0010008秒
LMDB	0.0005218秒
HDF5	0.0010201秒

可以看出，LMDB在读取速度上明显优于其他方法，尤其在读取单个图像时表现突出。

多个图像的读取

多图像读取的操作与单一图像类似，通过循环读取多个图像或从数据库一次性读取多个图像。接下来分别讨论不同存储方式下的多图像读取。

DISK读取

在DISK读取方式中，图像以.png文件的形式存储在磁盘上，读取时通过Pillow库逐个加载图像文件。对应的标签存储在一个.csv文件中，读取时通过CSV模块加载。这种方式适合处理较少量的图像数据，但随着图像数量增加，磁盘IO操作的耗时会显著增加，读取效率相对较低。

def read_many_disk(num_images):
    """从磁盘读取多个图像及其标签。"""
    images, labels = [], []

    for image_id in range(num_images):
        images.append(np.array(Image.open(disk_dir / f"{image_id}.png")))

    with open(disk_dir / f"{num_images}_labels.csv", "r") as csvfile:
        reader = csv.reader(csvfile, delimiter=" ", quotechar="|", quoting=csv.QUOTE_MINIMAL)
        for row in reader:
            if row:
                labels.append(int(row[0]))
            
    return images, labels

LMDB读取

在LMDB读取方式中，图像和标签以键值对的形式存储在LMDB数据库中。读取时，通过事务操作根据图像的唯一ID从数据库中提取存储的图像和标签数据，并将其解码为可使用的格式。LMDB具备较高的读取效率，特别适合大规模数据的随机访问和并行处理。

def read_many_lmdb(num_images):
    """从 LMDB 数据库读取多个图像及其标签。"""
    images, labels = []
    env = lmdb.open(str(lmdb_dir / f"{num_images}_lmdb"), readonly=True)

    with env.begin() as txn:
        for image_id in range(num_images):
            data = txn.get(f"{image_id:08}".encode("ascii"))
            cifar_image = pickle.loads(data)
            images.append(cifar_image.get_image())
            labels.append(cifar_image.label)
    env.close()

    return images, labels

HDF5读取

在HDF5读取方式中，多个图像和标签存储在HDF5文件的多个数据集中。读取时，通过HDF5文件访问接口加载所有图像和标签的数据集。HDF5存储格式能够有效处理多维数据，读取效率较高，尤其适用于处理大规模的科学数据或图像数据。

import h5py

def read_many_hdf5(num_images):
    """从 HDF5 文件读取多个图像及其标签。"""
    file = h5py.File(hdf5_dir / f"{num_images}_many.h5", "r+")
    images = np.array(file["/images"]).astype("uint8")
    labels = np.array(file["/meta"]).astype("uint8")
    file.close()

    return images, labels

读取方式对比

以下是多图像读取的性能对比表：

读取方法	读取 10 张图像耗时	读取 100 张图像耗时	读取 1000 张图像耗时	读取 10000 张图像耗时	读取 100000 张图像耗时
Disk	0.00788秒	0.04342秒	0.41955秒	4.09544秒	41.6931秒
LMDB	0.00049秒	0.00127秒	0.01066秒	0.10226秒	1.06638秒
HDF5	0.00093秒	0.00104秒	0.00424秒	0.02351秒	0.23247秒

可以看出，随着图像数量的增加，LMDB和HDF5在读取多个图像时的性能明显优于DISK存储方式。

读写操作综合比较

在存储和读取过程中，LMDB和HDF5在性能方面具有明显的优势。DISK存储虽然适合小规模数据的存储和读取，但在大规模数据场景下其性能劣势明显。

通过绘制图表，可以直观地展示不同存储方式下的读写性能对比。

并行操作

在大规模数据处理中，并行化处理可以显著提高读写效率。Python的multiprocessing库允许通过多进程执行多个任务，从而加速存储或读取操作。以下是并行存储数据的示例：

import multiprocessing

def process_image(image, label):
    """模拟图像处理操作"""
    return image, label * 2

def parallel_store(images, labels):
    """并行存储图像数据"""
    with multiprocessing.Pool(processes=4) as pool:
        results = pool.starmap(process_image, zip(images, labels))
    return results

# 示例并行处理
processed_images = parallel_store(images[:1000], labels[:1000])
print("并行存储完成")

通过并行化存储和读取操作，可以显著缩短处理时间，尤其在处理大量图像时，优势更加明显。

总结

通过对比三种存储方式，DISK、LMDB 和 HDF5 在不同场景下各有优劣。DISK 适合小规模的数据存储与读取，但随着数据量的增加，其性能会显著下降。LMDB 和 HDF5 则在处理大规模数据时表现出色，尤其是 LMDB，在读取性能上具备明显的优势。而 HDF5 则在多维数据集的存储和操作上展现了较高的灵活性。因此，在实际应用中，用户可以根据数据量和操作需求，选择合适的存储方式来实现高效的数据管理。此外，通过并行化的处理方式，还可以进一步提升数据读写的效率，满足更加复杂的应用需求。