| 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 | |