| 199 | |
| 200 | |
| 201 | class TextDataset(object): |
| 202 | def __init__(self, workdir, embedding_type, hr_lr_ratio): |
| 203 | lr_imsize = 64 |
| 204 | self.hr_lr_ratio = hr_lr_ratio |
| 205 | if self.hr_lr_ratio == 1: |
| 206 | self.image_filename = '/76images.pickle' |
| 207 | elif self.hr_lr_ratio == 4: |
| 208 | self.image_filename = '/304images.pickle' |
| 209 | |
| 210 | self.image_shape = [lr_imsize * self.hr_lr_ratio, |
| 211 | lr_imsize * self.hr_lr_ratio, 3] |
| 212 | self.image_dim = self.image_shape[0] * self.image_shape[1] * 3 |
| 213 | self.embedding_shape = None |
| 214 | self.train = None |
| 215 | self.test = None |
| 216 | self.workdir = workdir |
| 217 | if embedding_type == 'cnn-rnn': |
| 218 | self.embedding_filename = '/char-CNN-RNN-embeddings.pickle' |
| 219 | elif embedding_type == 'skip-thought': |
| 220 | self.embedding_filename = '/skip-thought-embeddings.pickle' |
| 221 | |
| 222 | def get_data(self, pickle_path, aug_flag=True): |
| 223 | with open(pickle_path + self.image_filename, 'rb') as f: |
| 224 | images = pickle.load(f) |
| 225 | images = np.array(images) |
| 226 | print('images: ', images.shape) |
| 227 | |
| 228 | with open(pickle_path + self.embedding_filename, 'rb') as f: |
| 229 | embeddings = pickle.load(f) |
| 230 | embeddings = np.array(embeddings) |
| 231 | self.embedding_shape = [embeddings.shape[-1]] |
| 232 | print('embeddings: ', embeddings.shape) |
| 233 | with open(pickle_path + '/filenames.pickle', 'rb') as f: |
| 234 | list_filenames = pickle.load(f) |
| 235 | print('list_filenames: ', len(list_filenames), list_filenames[0]) |
| 236 | with open(pickle_path + '/class_info.pickle', 'rb') as f: |
| 237 | class_id = pickle.load(f) |
| 238 | |
| 239 | return Dataset(images, self.image_shape[0], embeddings, |
| 240 | list_filenames, self.workdir, None, |
| 241 | aug_flag, class_id) |
no outgoing calls
no test coverage detected