制备MNIST数据集（yolo格式）

数据集准备

MNIST数据集下载与制作

import torch
from torchvision import datasets, transforms
from PIL import Image
import os

from tqdm import tqdm

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为 Tensor 格式
    transforms.Normalize((0.5,), (0.5,))  # 标准化图像数据
])

# 下载 MNIST 数据集
mnist_train = datasets.MNIST(root='data', train=True, download=True, transform=transform)
mnist_test = datasets.MNIST(root='data', train=False, download=True, transform=transform)

# 创建 YOLO 格式的数据集目录
output_dir = 'MNIST'
os.makedirs(output_dir, exist_ok=True)
os.makedirs(os.path.join(output_dir,'train'), exist_ok=True)
os.makedirs(os.path.join(output_dir,'test'), exist_ok=True)

# 定义类别
classes = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

for id in classes:
    if not os.path.exists(os.path.join(output_dir,'train',id)): os.makedirs(os.path.join(output_dir,'train',id), exist_ok=True)
    if not os.path.exists(os.path.join(output_dir,'test',id)): os.makedirs(os.path.join(output_dir,'test',id), exist_ok=True)

image_nums = 0
# 将训练数据转换为 YOLO 格式
train_output_path = os.path.join(output_dir, 'train.txt')
with open(train_output_path, 'w') as f:
    for (image, label) in tqdm(mnist_train):
        image_nums += 1
        # 将 Tensor 转换为 PIL 图像
        image_path = os.path.join(os.getcwd(),output_dir,"train",str(label),f'image_{str(image_nums)}.png')
        image_pil = Image.fromarray((image.squeeze().numpy() * 255).astype('uint8'))
        image_pil.save(image_path)
        f.write(f'{image_path}\n')

# 将测试数据转换为 YOLO 格式
test_output_path = os.path.join(output_dir, 'test.txt')
with open(test_output_path, 'w') as f:
    for (image, label) in tqdm(mnist_test):
        image_nums += 1
        # 将 Tensor 转换为 PIL 图像
        image_path = os.path.join(os.getcwd(), output_dir, "test", str(label), f'image_{str(image_nums)}.png')
        image_pil = Image.fromarray((image.squeeze().numpy() * 255).astype('uint8'))
        image_pil.save(image_path)
        f.write(f'{image_path}\n')

print(f"MNIST dataset has been converted to YOLO format and saved to {output_dir}")