torch之从.datasets.CIFAR10解压出训练与测试图片（附带网盘链接）

news2024/11/15 2:21:26

前言
从官网上下载的是长这个样子的
在这里插入图片描述
想看图片，咋办咧，看下面代码

import torch
import torchvision
import numpy as np
import os
import cv2
batch_size = 50

transform_predict = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
])
#-----#
# train 为True 则是解压出训练图片 为Fasle的时候则解压出测试图片
#------#
image_data = torchvision.datasets.CIFAR10(
    root='/home/netted/img_process_ml/temp', train=True, download=False, transform=transform_predict)
image_loader = torch.utils.data.DataLoader(
    image_data, batch_size, shuffle=True, num_workers=0)

path = '/home/netted/img_process_ml/temp/train'
os.makedirs(path,exist_ok=True)
for i in range(10):
    os.makedirs(f'{path}/{i}',exist_ok=True)


def format(image):
    image = image.clone().detach().cpu().squeeze(0)
    image = np.around(image.mul(255))
    image = np.uint8(image).transpose(1, 2, 0)
    return image


def data(image_loader):
    idx0 = 0
    idx1 = 0
    idx2 = 0
    idx3 = 0
    idx4 = 0
    idx5 = 0
    idx6 = 0
    idx7 = 0
    idx8 = 0
    idx9 = 0

    for i, (data, target) in enumerate(image_loader):

        for idx in range(len(data)):
            label = target[idx].item()
            image = format(data[idx])

            if label == 0:
                cv2.imwrite(f'{path}/{label}/plane_{idx0}.png',image)
                idx0 += 1

            if label == 1:
                cv2.imwrite(f'{path}/{label}/car_{idx1}.png', image)
                idx1 += 1

            if label == 2:
                cv2.imwrite(f'{path}/{label}/bird_{idx2}.png', image)
                idx2 += 1

            if label == 3:
                cv2.imwrite(f'{path}/{label}/cat_{idx3}.png', image)
                idx3 += 1

            if label == 4:
                cv2.imwrite(f'{path}/{label}/deer_{idx4}.png', image)
                idx4 += 1

            if label == 5:
                cv2.imwrite(f'{path}/{label}/dog_{idx5}.png', image)
                idx5 += 1

            if label == 6:
                cv2.imwrite(f'{path}/{label}/frog_{idx6}.png', image)
                idx6 += 1

            if label == 7:
                cv2.imwrite(f'{path}/{label}/horse_{idx7}.png', image)
                idx7 += 1

            if label == 8:
                cv2.imwrite(f'{path}/{label}/ship_{idx8}.png', image)
                idx8 += 1

            if label == 9:
                cv2.imwrite(f'{path}/{label}/truck_{idx9}.png', image)
                idx9 += 1

data(image_loader)