39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
|
from PIL import Image
|
|||
|
|
import torch
|
|||
|
|
from torch.utils.data import Dataset
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MyDataSet(Dataset):
|
|||
|
|
"""自定义数据集"""
|
|||
|
|
|
|||
|
|
def __init__(self, images_path: list, images_class: list, transform=None):
|
|||
|
|
self.images_path = images_path
|
|||
|
|
self.images_class = images_class
|
|||
|
|
self.transform = transform
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return len(self.images_path)
|
|||
|
|
|
|||
|
|
def __getitem__(self, item):
|
|||
|
|
img = Image.open(self.images_path[item])
|
|||
|
|
# RGB为彩色图片，L为灰度图片
|
|||
|
|
if img.mode != 'RGB':
|
|||
|
|
# img = img.convert('RGB')
|
|||
|
|
raise ValueError("image: {} isn't RGB mode.".format(self.images_path[item]))
|
|||
|
|
label = self.images_class[item]
|
|||
|
|
|
|||
|
|
if self.transform is not None:
|
|||
|
|
img = self.transform(img)
|
|||
|
|
|
|||
|
|
return img, label
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def collate_fn(batch):
|
|||
|
|
# 官方实现的default_collate可以参考
|
|||
|
|
# https://github.com/pytorch/pytorch/blob/67b7e751e6b5931a9f45274653f4f653a4e6cdf6/torch/utils/data/_utils/collate.py
|
|||
|
|
images, labels = tuple(zip(*batch))
|
|||
|
|
|
|||
|
|
images = torch.stack(images, dim=0)
|
|||
|
|
labels = torch.as_tensor(labels)
|
|||
|
|
return images, labels
|