Spaces:
Sleeping
Sleeping
File size: 2,120 Bytes
1ffed57 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import numpy as np
import albumentations
from torchvision import datasets
from albumentations.pytorch import ToTensorV2
from torch.utils.data import Dataset, DataLoader
class CIFAR10Data(Dataset):
def __init__(self, dataset, transforms=None) -> None:
self.dataset = dataset
self.transforms = transforms
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
image, label = self.dataset[index]
image = np.array(image)
if self.transforms:
image = self.transforms(image=image)['image']
return image, label
def _get_test_transforms():
test_transforms = albumentations.Compose([albumentations.Normalize([0.49139968, 0.48215841, 0.44653091],
[0.24703223, 0.24348513, 0.26158784]),
ToTensorV2()])
return test_transforms
def _get_data(is_train, is_download):
"""Method to get data for training or testing
Args:
is_train (bool): True if data is for training else false
is_download (bool): True to download dataset from iternet
Returns:
object: Oject of dataset
"""
data = datasets.CIFAR10('../data', train=is_train, download=is_download)
return data
def _get_data_loader(data, **kwargs):
"""Method to get data loader.
Args:
data (object): Oject of dataset
Returns:
object: Object of DataLoader class used to feed data to neural network model
"""
loader = DataLoader(data, **kwargs)
return loader
def get_test_data_loader(**kwargs):
"""Method to get data loader for testing
Args:
batch_size (int): Number of images in a batch
Returns:
object: Object of DataLoader class used to feed data to neural network model
"""
test_transforms = _get_test_transforms()
test_data = _get_data(is_train=False, is_download=True)
test_data = CIFAR10Data(test_data, test_transforms)
test_loader = _get_data_loader(data=test_data, **kwargs)
return test_loader |