-
Notifications
You must be signed in to change notification settings - Fork 0
/
Dataset.py
61 lines (50 loc) · 1.97 KB
/
Dataset.py
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
import pickle
import numpy as np
from PIL import Image
from torchvision import transforms
from sklearn.preprocessing import LabelEncoder
from torch.utils.data import Dataset, DataLoader
RESCALE_SIZE=224
class KTDataset(Dataset):
"""
Датасет с картинками, который паралельно подгружает их из папок
производит скалирование и превращение в торчевые тензоры
"""
def __init__(self, files, mode):
super().__init__()
self.files = sorted(files)
self.mode = mode
self.len_ = len(self.files)
self.label_encoder = LabelEncoder()
if self.mode != 'test':
self.labels = [path.parent.name for path in self.files]
self.label_encoder.fit(self.labels)
with open('label_encoder_v6.pkl', 'wb') as le_dump_file:
pickle.dump(self.label_encoder, le_dump_file)
def __len__(self):
return self.len_
def load_sample(self, file):
image = Image.open(file)
image.load()
return image
def __getitem__(self, index):
# для преобразования изображений в тензоры PyTorch и нормализации входа
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
x = self.load_sample(self.files[index])
x = self._prepare_sample(x)
x = np.array(x / 255, dtype='float32')
x = transform(x)
if self.mode == 'test':
return x
else:
label = self.labels[index]
label_id = self.label_encoder.transform([label])
y = label_id.item()
return x, y
def _prepare_sample(self, image):
image = image.crop((470, 0, 1449, 979))
image = image.resize((RESCALE_SIZE, RESCALE_SIZE))
return np.array(image)