30 lines
853 B
Python
30 lines
853 B
Python
import torch
|
|
from PIL import Image
|
|
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])
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
label = self.images_class[item]
|
|
if self.transform is not None:
|
|
img = self.transform(img)
|
|
return img, label
|
|
|
|
@staticmethod
|
|
def collate_fn(batch):
|
|
images, labels = tuple(zip(*batch))
|
|
images = torch.stack(images, dim=0)
|
|
labels = torch.as_tensor(labels)
|
|
return images, labels
|