-
Notifications
You must be signed in to change notification settings - Fork 1
/
dataset.py
152 lines (128 loc) · 5.28 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
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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import torch.utils.data as data
from PIL import Image
import PIL
import os
import os.path
import pickle
import random
import numpy as np
import pandas as pd
class TextDataset(data.Dataset):
def __init__(self, data_dir, split='train', embedding_type='cnn-rnn',
imsize=64, transform=None, target_transform=None):
self.transform = transform
self.target_transform = target_transform
self.imsize = imsize
self.data = []
self.data_dir = data_dir
if data_dir.find('birds') != -1:
self.bbox = self.load_bbox()
else:
self.bbox = None
split_dir = os.path.join(data_dir, split)
self.filenames = self.load_filenames(split_dir)
self.embeddings = self.load_embedding(split_dir, embedding_type)
self.captions = self.load_all_captions()
def get_img(self, img_path, bbox):
img = Image.open(img_path).convert('RGB')
width, height = img.size
if bbox is not None:
R = int(np.maximum(bbox[2], bbox[3]) * 0.75)
center_x = int((2 * bbox[0] + bbox[2]) / 2)
center_y = int((2 * bbox[1] + bbox[3]) / 2)
y1 = np.maximum(0, center_y - R)
y2 = np.minimum(height, center_y + R)
x1 = np.maximum(0, center_x - R)
x2 = np.minimum(width, center_x + R)
img = img.crop([x1, y1, x2, y2])
load_size = int(self.imsize * 76 / 64)
img = img.resize((load_size, load_size), PIL.Image.BILINEAR)
if self.transform is not None:
img = self.transform(img)
return img
def load_bbox(self):
data_dir = self.data_dir
bbox_path = os.path.join(data_dir, 'CUB_200_2011/bounding_boxes.txt')
df_bounding_boxes = pd.read_csv(bbox_path,
delim_whitespace=True,
header=None).astype(int)
#
filepath = os.path.join(data_dir, 'CUB_200_2011/images.txt')
df_filenames = \
pd.read_csv(filepath, delim_whitespace=True, header=None)
filenames = df_filenames[1].tolist()
print('Total filenames: ', len(filenames), filenames[0])
#
filename_bbox = {img_file[:-4]: [] for img_file in filenames}
numImgs = len(filenames)
for i in range(0, numImgs):
# bbox = [x-left, y-top, width, height]
bbox = df_bounding_boxes.iloc[i][1:].tolist()
key = filenames[i][:-4]
filename_bbox[key] = bbox
#
return filename_bbox
def load_all_captions(self):
caption_dict = {}
for key in self.filenames:
caption_name = '%s/text/%s.txt' % (self.data_dir, key)
captions = self.load_captions(caption_name)
caption_dict[key] = captions
return caption_dict
def load_captions(self, caption_name):
cap_path = caption_name
with open(cap_path, "r") as f:
#captions = f.read().decode('utf8').split('\n')
captions = f.read().split('\n')
captions = [cap.replace("\ufffd\ufffd", " ")
for cap in captions if len(cap) > 0]
return captions
def load_embedding(self, data_dir, embedding_type):
if embedding_type == 'cnn-rnn':
embedding_filename = '/char-CNN-RNN-embeddings.pickle'
elif embedding_type == 'cnn-gru':
embedding_filename = '/char-CNN-GRU-embeddings.pickle'
elif embedding_type == 'skip-thought':
embedding_filename = '/skip-thought-embeddings.pickle'
print("data dir:", data_dir)
print("embidding filename:",embedding_filename)
with open(data_dir + embedding_filename, 'rb') as f:
# for python 3
u = pickle._Unpickler(f)
u.encoding = 'latin1'
embeddings = u.load()
# suppress for python 3.x
#embeddings = pickle.load(f)
embeddings = np.array(embeddings)
# embedding_shape = [embeddings.shape[-1]]
print('embeddings: ', embeddings.shape)
return embeddings
def load_filenames(self, data_dir):
filepath = os.path.join(data_dir, 'filenames.pickle')
with open(filepath, 'rb') as f:
filenames = pickle.load(f)
print('Load filenames from: %s (%d)' % (filepath, len(filenames)))
return filenames
def __getitem__(self, index):
key = self.filenames[index]
if self.bbox is not None:
bbox = self.bbox[key]
data_dir = '%s/CUB_200_2011' % self.data_dir
else:
bbox = None
data_dir = self.data_dir
captions = self.captions[key]
embeddings = self.embeddings[index, :, :]
img_name = '%s/images/%s.jpg' % (data_dir, key)
img = self.get_img(img_name, bbox)
rand_ix = random.randint(0, embeddings.shape[0]-1)
embedding = embeddings[rand_ix, :]
if self.target_transform is not None:
embedding = self.target_transform(embedding)
return img, embedding,captions[rand_ix]
def __len__(self):
return len(self.filenames)