-
Notifications
You must be signed in to change notification settings - Fork 1
/
dataset.py
80 lines (63 loc) · 2.24 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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import os
OUTPUT_DIR = './'
if not os.path.exists(OUTPUT_DIR):
os.makedirs(OUTPUT_DIR)
import numpy as np
import cv2
import torch
from torch.utils.data import DataLoader, Dataset
class TrainDataset(Dataset):
def __init__(self, df, tokenizer, transform=None):
super().__init__()
self.df = df
self.tokenizer = tokenizer
self.file_paths = df['file_path'].values
self.labels = df['captions'].values
self.transform = transform
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
file_path = self.file_paths[idx]
image = cv2.imread(file_path)
try:
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
except:
print('Error returned')
if self.transform:
augmented = self.transform(image = image)
image = augmented['image']
label = self.labels[idx]
label = self.tokenizer.text_to_sequence(label)
label_length = len(label)
label_length = torch.LongTensor([label_length])
return image, torch.LongTensor(label), label_length
class TestDataset(Dataset):
def __init__(self, df, transform=None):
super().__init__()
self.df = df
self.file_paths = df['file_path'].values
self.transform = transform
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
file_path = self.file_paths[idx]
image = cv2.imread(file_path)
# trick to overcome dead files
if image is None:
#check if jpg or png
decode_type = file_path[-3:]
new_path = ''
if decode_type == 'jpg':
new_path = file_path[:-4] + '.png'
image = cv2.imread(new_path)
else:
new_path = file_path[:-4] + '.jpg'
image = cv2.imread(new_path)
if image is None:
print(file_path)
print(new_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
if self.transform:
augmented = self.transform(image=image)
image = augmented['image']
return image