-
Notifications
You must be signed in to change notification settings - Fork 0
/
get_data.py
124 lines (91 loc) · 3.86 KB
/
get_data.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
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import torch
from torch import nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import Dataset, DataLoader
import albumentations as A
from albumentations.pytorch import ToTensor
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from addict import Dict
from PIL import Image
import os
def get_data(train_config, validation_config):
class BirdsDataset(Dataset):
def __init__(self, path, augmentations=None, shape=(224, 224)):
self.__augmentations = augmentations
self.__shape = shape
self.__images_labels = []
self.labels = []
if os.path.exists(path):
self.__path = path
self.__labels = os.listdir(self.__path)
for label in self.__labels:
label_path = os.path.join(self.__path, label)
if os.path.isdir(label_path):
images = os.listdir(label_path)
for image in images:
if image.endswith("png") or image.endswith("jpg"):
image_path = os.path.join(label_path, image)
info = (image_path, label)
self.__images_labels.append(info)
self.labels.append(label)
else:
pass
else:
raise Exception(f"Path '{path}' for DataFrame doesn't exist!")
def __getitem__(self, index):
image_path, label = self.__images_labels[index]
image = self.__load_image(image_path)
if self.__augmentations is not None:
image = self.__augmentations(image=image.permute(1, 2, 0).numpy())["image"]
label = self.labels.index(label)
return Dict({
"image": image,
"label": label,
})
def __len__(self):
return len(self.__images_labels)
def __load_image(self, path, channels="RGB"):
width, height = self.__shape
loader = A.Compose([
A.Resize(width, height),
ToTensor()
])
image_array = np.array(Image.open(path).convert(channels))
return loader(image=image_array)["image"]
train_dataset = BirdsDataset(path=train_config["path"],
augmentations=train_config["augmentations"],
shape=train_config["shape"])
validation_dataset = BirdsDataset(path=validation_config["path"],
augmentations=validation_config["augmentations"],
shape=validation_config["shape"])
print(f"Train Size: {len(train_dataset)}")
print(f"Validation Size: {len(validation_dataset)}")
print(f"Num Classes: {len(train_dataset.labels)}")
def collate_fn(batch):
images, labels = [], []
for item in batch:
image = item["image"].tolist()
images.append(image)
label = item["label"]
labels.append(label)
return Dict({
"images": torch.tensor(images),
"labels": torch.tensor(labels)
})
train_loader = DataLoader(train_dataset,
shuffle=True,
num_workers=train_config["num_workers"],
pin_memory=True,
batch_size=train_config["batch_size"],
collate_fn=collate_fn)
validation_loader = DataLoader(validation_dataset,
shuffle=True,
num_workers=validation_config["num_workers"],
pin_memory=True,
batch_size=validation_config["batch_size"],
collate_fn=collate_fn)
return train_loader, validation_loader