Skip to content

图片预处理

零散的图片,需要处理成大文件才能加快读取的速度。所以就需要一个预处理

图片处理成HDF5格式的文件

2025年8月12日 重新编写了代码,将图片先加载到内存中,再保存到hdf5这个文件中去,测试了一下速度还行,贴上代码

800张图片运行了6.6s,有空测试一下这个代码在imagenet上的速度。128W的图片理论上33分钟能处理完成

alt text

python

import os
import h5py
import numpy as np
import cv2
from tqdm import tqdm

def load_images_to_hdf5(hotdog_dir, not_hotdog_dir, output_file, img_size=(128, 128)):
    """
    将hotdog和not-hotdog图片加载并保存到HDF5文件(分离加载和存储步骤)
    
    参数:
        hotdog_dir: hotdog图片所在目录
        not_hotdog_dir: not-hotdog图片所在目录
        output_file: 输出的HDF5文件名
        img_size: 图片 resize 的尺寸,默认(128, 128)
    """
    # 获取所有图片路径和对应的标签
    image_paths = []
    labels = []  # 1表示hotdog,0表示not-hotdog
    
    # 处理hotdog图片
    for img_name in os.listdir(hotdog_dir):
        img_path = os.path.join(hotdog_dir, img_name)
        if os.path.isfile(img_path) and img_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
            image_paths.append(img_path)
            labels.append(1)
    
    # 处理not-hotdog图片
    for img_name in os.listdir(not_hotdog_dir):
        img_path = os.path.join(not_hotdog_dir, img_name)
        if os.path.isfile(img_path) and img_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
            image_paths.append(img_path)
            labels.append(0)
    
    # 计算总图片数量
    total_images = len(image_paths)
    if total_images == 0:
        print("没有找到任何图片文件")
        return
    
    # 第一步:先将所有图片加载到内存
    print("开始加载图片到内存...")
    images = []
    valid_labels = []
    valid_paths = []
    
    for i in tqdm(range(total_images), desc="加载图片"):
        try:
            # 使用OpenCV读取图片(默认BGR格式)
            img = cv2.imread(image_paths[i])
            
            if img is None:
                print(f"无法读取图片 {image_paths[i]}")
                continue
            
            # 调整图片大小
            img = cv2.resize(img, img_size)
            
            # 转换为RGB格式(因为OpenCV默认是BGR)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # 存储到内存列表
            images.append(img)
            valid_labels.append(labels[i])
            valid_paths.append(image_paths[i])
            
        except Exception as e:
            print(f"处理图片 {image_paths[i]} 时出错: {e}")
    
    # 转换为numpy数组以便高效存储
    images_np = np.array(images, dtype=np.uint8)
    labels_np = np.array(valid_labels, dtype=np.int32)
    
    # 第二步:将内存中的所有数据写入HDF5文件
    print(f"开始将 {len(images)} 张图片写入HDF5文件...")
    with h5py.File(output_file, 'w') as hf:
        # 创建数据集来存储图片和标签
        images_ds = hf.create_dataset('images', 
                                     shape=images_np.shape,
                                     dtype=np.uint8,
                                     compression='gzip',
                                     compression_opts=9)
        
        labels_ds = hf.create_dataset('labels', 
                                     shape=labels_np.shape,
                                     dtype=np.int32)
        
        # 存储图片路径以便后续参考
        str_type = h5py.special_dtype(vlen=str)
        paths_ds = hf.create_dataset('paths', 
                                    shape=(len(valid_paths),),
                                    dtype=str_type)
        
        # 一次性写入所有数据
        images_ds[:] = images_np
        labels_ds[:] = labels_np
        
        # 写入路径
        for i, path in enumerate(valid_paths):
            paths_ds[i] = path
    
    print(f"成功将 {len(images)} 张图片保存到 {output_file}")
    print(f"其中 hotdog 图片: {sum(valid_labels)} 张")
    print(f"其中 not-hotdog 图片: {len(valid_labels) - sum(valid_labels)} 张")
    print(f"跳过了 {total_images - len(images)} 张无法处理的图片")

if __name__ == "__main__":
    # 定义图片目录路径
    hotdog_directory = "hotdog/test/hotdog"
    not_hotdog_directory = "hotdog/test/not-hotdog"
    
    # 定义输出HDF5文件路径
    output_hdf5 = "123.h5"
    
    # 调用函数处理图片,可根据需要调整尺寸
    load_images_to_hdf5(hotdog_directory, not_hotdog_directory, output_hdf5, img_size=(224, 224))

后续的改善

添加并行处理的代码,这样理论上可以降低到几分钟,之前计算的33分钟是单线程的计算结果。

将文件处理成JSION文件,效果应该的hdf5是相同的,看李宏毅的homework数据也改成这个格式的了。

SON(JavaScript Object Notation)是一种轻量级的数据交换格式,用于表示结构化数据。它以简单的键值对形式组织数据,易于阅读和编写,同时具有跨平台和语言的兼容性。