pytorch数据处理样例

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@desc: 这是一个pytorch数据处理样例
参考原文: https://blog.youkuaiyun.com/weixin_41560402/article/details/108121344

Ⅰ、加载数据 
    在pytorch中,提供了一种十分方便的数据读取机制,即
    使用torch.utils.data.Dataset与torch.utils.data.DataLoader组合得到数据迭代器。
    在每次训练时,利用这个迭代器输出每一个batch数据,并能在输出时对数据进行相应的预处理或数据增强等操作。
Ⅱ、数据预处理
    pytorch官方API, torchvision.transforms主要实现对数据集的预处理、数据增强、转换成tensor等一系列操作。

@date: 2023-11-28 15:25:48
@software: vscode
'''
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms

# ———————————————————————— 示例1：张量数据 ——————————————————————————
class TensorDataset(Dataset):
    """
    TensorDataset继承Dataset, 重载了__init__(), __getitem__(), __len__()魔法方法
    实现将一组Tensor数据对封装成Tensor数据集,能够通过index得到数据集的数据,能够通过len得到数据集大小
        __len__()：返回的是数据集的大小。
                构建的数据集是一个对象,它不像序列类型列表、元组、字符串那样可以直接用len()来获取序列的长度 
                魔法方法__len__()的目的就是方便像序列那样直接获取对象的长度。
                如果A是一个类, a是类A的实例化对象, 当A中定义了魔法方法__len__() , len(a)则返回对象的大小。
        __getitem__()：实现索引数据集中的某一个数据。
                序列可以通过索引的方法获取序列中的任意元素
                __getitem__()则实现了能够通过索引的方法获取对象中的任意元素。
                同时,可以在__getitem__()中实现数据预处理。

    """
    def __init__(self, data_tensor, target_tensor):
        self.data_tensor = data_tensor
        self.target_tensor = target_tensor

    def __getitem__(self, index):
        return self.data_tensor[index], self.target_tensor[index]

    def __len__(self):
        return self.data_tensor.size(0)

# 生成数据
data_tensor = torch.randn(4, 3)
target_tensor = torch.rand(4)

# 将数据封装成Dataset
tensor_dataset = TensorDataset(data_tensor, target_tensor)

# 可使用索引调用数据
print(tensor_dataset[1])
# 输出：(tensor([-1.0351, -0.1004,  0.9168]), tensor(0.4977))

# 获取数据集大小
print(len(tensor_dataset))
# 输出：4


# —————————————————————— torch.utils.data.DataLoader ——————————————————————
tensor_dataloader = DataLoader(tensor_dataset,   # 封装的对象
                               batch_size=2,     # 输出的batch size
                               shuffle=True,     # 随机输出
                               num_workers=0)    # 只有1个进程

# 以for循环形式输出
for data, target in tensor_dataloader:
    print(data, target)



# ———————————————————————— 示例2：图像数据 ———————————————————————— 
import os
from PIL import Image
from torch.utils.data import Dataset


class PatchDataset(Dataset):
    def __init__(self, data_dir, transforms=None):
        """
        :param data_dir: 数据集所在路径
        :param transform: 数据预处理
        """

        self.data_info = self.get_img_info(data_dir)
        self.transforms = transforms

    def __getitem__(self, item):
        path_img, label = self.data_info[item]
        image = Image.open(path_img).convert('RGB')
        if self.transforms is not None:
            image = self.transforms(image)

        return image, label

    def __len__(self):
        return len(self.data_info)

    @staticmethod
    def get_img_info(data_dir):
        path_dir = os.path.join(data_dir, 'train_dataset.txt')
        data_info = []
        with open(path_dir) as file:
            lines = file.readlines()
            for line in lines:
                data_info.append(line.strip('\n').split(' '))
        return data_info



# ———————————————————————— transforms.Compose ———————————————————————— 
# Compose类的作用是组合多个transforms函数，
# Compose类的初始化函数中需要传入一个含有多种transform方法的列表，
# 随后将图像逐一通过这些transform方法。
# ———————————————————————— transforms.ToTensor ———————————————————————— 
# 这个类的作用是将PIL Image或numpy.ndarray转换成tensor，在转换前会将调整维度，并进行单位化：
# Converts a PIL Image or numpy.ndarray (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]

from torchvision import transforms

transforms_train = transforms.Compose([transforms.Resize(40),
                                       transforms.RandomResizedCrop(32, scale=(0.64, 1.0), ratio=(1.0, 1.0)),
                                       transforms.RandomHorizontalFlip(),
                                       transforms.ToTensor(),
                                       transforms.Normalize(mean=[0.4914, 0.4832, 0.4856],
                                                            std=[0.2023, 0.2013, 0.2111])])

transforms_test = transforms.Compose([transforms.ToTensor(),
                                      transforms.Normalize(mean=[0.4914, 0.4832, 0.4856],
                                                           std=[0.2023, 0.2013, 0.2111])])


# ———————————————————————— torchvision.transforms ———————————————————————— 
# transforms主要实现对数据集的预处理、数据增强、转换成tensor等一系列操作。
# ———————————————————————— transforms主要用在Dataset类构建过程中，整个流程如下所示 ————————————————————————
from PIL import Image
from torchvision import transforms
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, data_dir, transforms=None):
        self.data_info = self.get_img_info(data_dir)
        self.transforms = transforms

    def __getitem__(self, item):
        path_img, label = self.data_info[item]
        image = Image.open(path_img).convert('RGB')
        # 使用定义好的transforms，对数据进行处理
        if self.transforms is not None:
            image = self.transforms(image)

        return image, label

    def __len__(self):
        return len(self.data_info)

train_transforms = transforms.Compose([transforms.ToTensor(),
                                       transforms.RandomHorizontalFlip(0.5)])
data_dir = "root_dir"   # 数据集路径
train_dataset = MyDataset(data_dir, train_transforms)


# —————————————————————————————————— 自定义transforms ——————————————————————————————————
# 对于目标检测，在对原始图像进行数据增强时，需要同时对目标的边界框坐标做相应的调整；
# 或者我们需要构建自己的数据增强方法，这个时候我们就需要自己定义transforms。
import random
from torchvision.transforms import functional as F


class Compose(object):
    """组合多个transform函数"""
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target


class ToTensor(object):
    """将PIL图像转为Tensor"""
    def __call__(self, image, target):
        image = F.to_tensor(image)
        # target不需要对维度进行调整或单位化
        target = torch.as_tensor(np.array(target), dtype=torch.int64)
        return image, target


class RandomHorizontalFlip(object):
    """随机水平翻转图像以及bboxes"""
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            height, width = image.shape[-2:]
            image = image.flip(-1)  # 水平翻转图片
            bbox = target["boxes"]
            # bbox: xmin, ymin, xmax, ymax
            bbox[:, [0, 2]] = width - bbox[:, [2, 0]]  # 翻转对应bbox坐标信息
            target["boxes"] = bbox
        return image, target

# —————————————————————————————————— 自定义transforms ——————————————————————————————————
# 对于图像分割，我们在做数据增强时同样需要自己定义transforms。
import numpy as np
from PIL import Image
import random

import torch
from torchvision import transforms as T
from torchvision.transforms import functional as F


def pad_if_smaller(img, size, fill=0):
    min_size = min(img.size)
    if min_size < size:
        ow, oh = img.size
        padh = size - oh if oh < size else 0
        padw = size - ow if ow < size else 0
        img = F.pad(img, (0, 0, padw, padh), fill=fill)
    return img


class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target


class RandomResize(object):
    def __init__(self, min_size, max_size=None):
        self.min_size = min_size
        if max_size is None:
            max_size = min_size
        self.max_size = max_size

    def __call__(self, image, target):
        size = random.randint(self.min_size, self.max_size)
        image = F.resize(image, size)
        target = F.resize(target, size, interpolation=Image.NEAREST)
        return image, target


class RandomHorizontalFlip(object):
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = F.hflip(image)
            target = F.hflip(target)
        return image, target


class RandomCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = pad_if_smaller(image, self.size)
        target = pad_if_smaller(target, self.size, fill=255)
        crop_params = T.RandomCrop.get_params(image, (self.size, self.size))
        image = F.crop(image, *crop_params)
        target = F.crop(target, *crop_params)
        return image, target


class CenterCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = F.center_crop(image, self.size)
        target = F.center_crop(target, self.size)
        return image, target


class ToTensor(object):
    def __call__(self, image, target):
        image = F.to_tensor(image)
        target = torch.as_tensor(np.array(target), dtype=torch.int64)
        return image, target


class Normalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, target):
        image = F.normalize(image, mean=self.mean, std=self.std)
        return image, target