| from skimage.io import imread |
| from skimage.transform import resize |
| import numpy as np |
| import tensorflow as tf |
| import keras |
| from skimage.transform import resize |
| from tf_record_utility import TFRecordUtility |
| from configuration import DatasetName, DatasetType, \ |
| AffectnetConf, D300wConf, W300Conf, InputDataSize, LearningConfig |
| from numpy import save, load, asarray |
|
|
|
|
| class CustomHeatmapGenerator(keras.utils.Sequence): |
|
|
| def __init__(self, is_single, image_filenames, label_filenames, batch_size, n_outputs, accuracy=100): |
| self.image_filenames = image_filenames |
| self.label_filenames = label_filenames |
| self.batch_size = batch_size |
| self.n_outputs = n_outputs |
| self.is_single = is_single |
| self.accuracy = accuracy |
|
|
| def __len__(self): |
| _len = np.ceil(len(self.image_filenames) // float(self.batch_size)) |
| return int(_len) |
|
|
| def __getitem__(self, idx): |
| img_path = D300wConf.train_images_dir |
| tr_path_85 = D300wConf.train_hm_dir_85 |
| tr_path_90 = D300wConf.train_hm_dir_90 |
| tr_path_97 = D300wConf.train_hm_dir_97 |
| tr_path = D300wConf.train_hm_dir |
|
|
| batch_x = self.image_filenames[idx * self.batch_size:(idx + 1) * self.batch_size] |
| batch_y = self.label_filenames[idx * self.batch_size:(idx + 1) * self.batch_size] |
|
|
| img_batch = np.array([imread(img_path + file_name) for file_name in batch_x]) |
|
|
| if self.is_single: |
| if self.accuracy == 85: |
| lbl_batch = np.array([load(tr_path_85 + file_name) for file_name in batch_y]) |
| elif self.accuracy == 90: |
| lbl_batch = np.array([load(tr_path_90 + file_name) for file_name in batch_y]) |
| elif self.accuracy == 97: |
| lbl_batch = np.array([load(tr_path_97 + file_name) for file_name in batch_y]) |
| else: |
| lbl_batch = np.array([load(tr_path + file_name) for file_name in batch_y]) |
|
|
| lbl_out_array = lbl_batch |
| else: |
| lbl_batch_85 = np.array([load(tr_path_85 + file_name) for file_name in batch_y]) |
| lbl_batch_90 = np.array([load(tr_path_90 + file_name) for file_name in batch_y]) |
| lbl_batch_97 = np.array([load(tr_path_97 + file_name) for file_name in batch_y]) |
| lbl_out_array = [lbl_batch_85, lbl_batch_90, lbl_batch_97] |
|
|
| return img_batch, lbl_out_array |
|
|
