huggingface官网下载并处理ImageNet2012数据集

文章目录

• 一、下载imagenet2012数据集
• 二、转换imagenet数据集格式

ImageNet数据集可以直接从ImageNet官方网站获取数据，但通常需要注册并遵守使用协议。另外，由于数据集较大，往往下载需要花费大量的时间空间，而通过huggingface下载数据集的方法不仅速度相对较快，而且能够直接从服务器远程进行下载。其中在huggingface官网下载模型或数据集的方法可以参考我之前的文章：huggingface学习 | 云服务器使用hf_hub_download下载huggingface上的模型文件。

一、下载imagenet2012数据集

huggingface官网上的imagenet2012数据集：https://huggingface.co/datasets/ILSVRC/imagenet-1k

可以看到这个数据集里包括训练集、验证集和测试集的所有数据，其中训练集分成了四个压缩包：

运行下面代码即可将训练集下载到local_dir指定的文件夹中：

import os
# 注意os.environ得在import huggingface库相关语句之前执行。
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
from huggingface_hub import hf_hub_download

def download_model(repo_id,filename,subfolder):
    local_dir = r'imagenet2012'
    print(f'开始下载\n仓库：{repo_id}\n大模型：{filename}\n如超时不用管，会自定继续下载，直至完成。中途中断，再次运行将继续下载。')
    while True:   
        try:
            hf_hub_download(local_dir=local_dir,
            repo_id=repo_id,
            token='你的huggingface用户信息',
            filename=filename,
            subfolder=subfolder,
            repo_type='dataset',
            local_dir_use_symlinks=False,
            resume_download=True,
            etag_timeout=100
            )
        except Exception as e :
            print(e)
        else:
            print(f'下载完成，大模型保存在：{local_dir}\{filename}')
            break
            
if __name__ == '__main__':
    repo_id ='ILSVRC/imagenet-1k'
    filenames1 = ['train_images_0.tar.gz','train_images_1.tar.gz','train_images_2.tar.gz','train_images_3.tar.gz','train_images_4.tar.gz']
    subfolder = 'data'
    for filename in filenames1:
        download_model(repo_id,filename,subfolder)

二、转换imagenet数据集格式

将上面下载得到的压缩包解压缩即可得到所有的图片，即一个文件夹里包含所有的图片。然而在进行深度学习任务时，往往需要将每张图片与其类别对应起来。因此需要转换为以下格式：

使用以下代码即可将文件夹转换为上述格式：

import os
import shutil

def split_by_nth_separator(string, separator, n):
    parts = string.split(separator)
    return separator.join(parts[:n]), separator.join(parts[n:])

def move_and_rename_images(source_dir, target_dir):
    # 遍历源文件夹中的文件
    for filename in os.listdir(source_dir):
        source_file = os.path.join(source_dir, filename)
        # 检查是否为文件且为图片文件
        if os.path.isfile(source_file) and filename.lower().endswith('.jpeg'):
            print(f'filename={filename}')
            folder_name=split_by_nth_separator(filename,'_',2)[1].split('.')[0]
            new_filename=split_by_nth_separator(filename,'_',2)[0]+'.JPEG'
            target_folder=os.path.join(target_dir,folder_name)
            print(f'folder_name={folder_name}',f'new_filename={new_filename}',f'target_folder={target_folder}',sep='\n')
            # 确保目标文件夹存在
            if not os.path.exists(target_folder):
                os.makedirs(target_folder)
            # 构造目标文件路径及新文件名
            target_file = os.path.join(target_folder, new_filename)
            # 移动文件并重命名
            shutil.move(source_file, target_file)
            print(f"Moved and renamed: {source_file} -> {target_file}")


# 指定源文件夹和目标文件夹路径
source_directory = 'imagenet2012/train_1' #图片数据所在文件夹
target_directory = 'imagenet2012/train' #处理之后的文件夹

# 调用函数
move_and_rename_images(source_directory, target_directory)