|
|
import os |
|
|
from PIL import Image |
|
|
from torch.utils.data import Dataset, DataLoader |
|
|
from torchvision import transforms |
|
|
|
|
|
class TinyImageNetDataset(Dataset): |
|
|
def __init__(self, root_dir, transform=None, train=True): |
|
|
self.root_dir = root_dir |
|
|
self.transform = transform |
|
|
self.image_paths = [] |
|
|
|
|
|
if train: |
|
|
|
|
|
train_dir = os.path.join(root_dir, 'train') |
|
|
for cls in os.listdir(train_dir): |
|
|
cls_dir = os.path.join(train_dir, cls, 'images') |
|
|
for img_name in os.listdir(cls_dir): |
|
|
if img_name.endswith('.JPEG'): |
|
|
self.image_paths.append(os.path.join(cls_dir, img_name)) |
|
|
else: |
|
|
|
|
|
val_dir = os.path.join(root_dir, 'val') |
|
|
images_dir = os.path.join(val_dir, 'images') |
|
|
for img_name in os.listdir(images_dir): |
|
|
if img_name.endswith('.JPEG'): |
|
|
self.image_paths.append(os.path.join(images_dir, img_name)) |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.image_paths) |
|
|
|
|
|
def __getitem__(self, idx): |
|
|
img = Image.open(self.image_paths[idx]).convert('RGB') |
|
|
if self.transform: |
|
|
img = self.transform(img) |
|
|
return img, 0 |
|
|
|
|
|
def get_dataloaders(config): |
|
|
transform = transforms.Compose([ |
|
|
transforms.Resize(config.image_size), |
|
|
transforms.RandomHorizontalFlip(), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) |
|
|
]) |
|
|
|
|
|
train_dataset = TinyImageNetDataset(config.dataset_path, transform=transform, train=True) |
|
|
val_dataset = TinyImageNetDataset(config.dataset_path, transform=transform, train=False) |
|
|
|
|
|
train_loader = DataLoader( |
|
|
train_dataset, |
|
|
batch_size=config.batch_size, |
|
|
shuffle=True, |
|
|
num_workers=config.num_workers, |
|
|
pin_memory=True |
|
|
) |
|
|
|
|
|
val_loader = DataLoader( |
|
|
val_dataset, |
|
|
batch_size=config.batch_size, |
|
|
shuffle=False, |
|
|
num_workers=config.num_workers |
|
|
) |
|
|
|
|
|
return train_loader, val_loader |
