File size: 2,849 Bytes
ce254e2 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | from datasets import load_dataset, load_from_disk
from src.paths import RAW_DATA_DIR, REPO_ID
from torch.utils.data import Dataset, TensorDataset, DataLoader as TorchDataLoader
from PIL import Image
import numpy as np
import torch
import torchvision.transforms as T
class HuggingFaceImageDataset(Dataset):
def __init__(self, hf_dataset):
self.data = hf_dataset
self.transform = T.Compose([
T.Grayscale(num_output_channels=3),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
item = self.data[idx]
img = item['image']
if not isinstance(img, Image.Image):
img = Image.fromarray(np.uint8(img))
x = self.transform(img)
y = torch.tensor(item['label'], dtype=torch.long)
return x, y
class DataLoader:
def __init__(self, test_size = 0.2, batch_size = 32, seed = 42, num_workers = 2):
self.local_path = RAW_DATA_DIR
self.repository_id = REPO_ID
self.test_size = test_size
self.batch_size = batch_size
self.seed = seed
self.num_workers = num_workers
def _raw_is_empty(self) -> bool:
return not any(f for f in self.local_path.iterdir() if not f.name.startswith("."))
def download(self):
# only download if there's not data in raw folder
if self._raw_is_empty():
dataset = load_dataset(self.repository_id)
dataset.save_to_disk(str(self.local_path))
def load(self):
if self._raw_is_empty():
raise FileNotFoundError("Data folder is empty")
return load_from_disk(str(self.local_path))
def get_train_data(self):
train_data = self.load()['train']
X_train = train_data['image']
y_train = train_data['label']
return X_train, y_train
def get_train_loader(self):
train_data = self.load()['train']
dataset = HuggingFaceImageDataset(train_data)
return TorchDataLoader(
dataset,
batch_size = self.batch_size,
shuffle = True,
num_workers = self.num_workers,
pin_memory = True
)
def get_test_data(self):
test_data = self.load()['test']
X_test = test_data['image']
y_test = test_data['label']
return X_test, y_test
def get_test_loader(self):
test_data = self.load()['test']
dataset = HuggingFaceImageDataset(test_data)
return TorchDataLoader(
dataset,
batch_size = self.batch_size,
shuffle = False,
num_workers = self.num_workers,
pin_memory = True
)
|