from config import DatasetName, W300Conf, DatasetType, LearningConfig, InputDataSize, CofwConf from cnn import CNNModel import tensorflow as tf from tensorflow import keras import numpy as np import matplotlib.pyplot as plt import math from datetime import datetime from sklearn.utils import shuffle from sklearn.model_selection import train_test_split from numpy import save, load, asarray import csv from skimage.io import imread import pickle from tqdm import tqdm import os from data_util import DataUtil from acr_loss import ACRLoss class Train: def __init__(self, arch, dataset_name, save_path, lambda_weight): self.lambda_weight = lambda_weight self.dataset_name = dataset_name self.save_path = save_path self.arch = arch self.base_lr = 1e-3 self.max_lr = 5e-3 if dataset_name == DatasetName.w300: self.num_landmark = W300Conf.num_of_landmarks * 2 self.img_path = W300Conf.train_image self.annotation_path = W300Conf.train_annotation '''evaluation path:''' self.eval_img_path = W300Conf.test_image_path + 'challenging/' self.eval_annotation_path = W300Conf.test_annotation_path + 'challenging/' if dataset_name == DatasetName.cofw: self.num_landmark = CofwConf.num_of_landmarks * 2 self.img_path = CofwConf.train_image self.annotation_path = CofwConf.train_annotation '''evaluation path:''' self.eval_img_path = CofwConf.test_image_path self.eval_annotation_path = CofwConf.test_annotation_path def train(self, weight_path): """ :param weight_path: :return: """ '''create loss''' c_loss = ACRLoss() '''create summary writer''' summary_writer = tf.summary.create_file_writer( "./train_logs/fit/" + datetime.now().strftime("%Y%m%d-%H%M%S")) '''create sample generator''' x_train_filenames, y_train_filenames = self._create_generators() '''making models''' model = self.make_model(arch=self.arch, w_path=weight_path) '''create train configuration''' step_per_epoch = len(x_train_filenames) // LearningConfig.batch_size lr = 1e-3 for epoch in range(LearningConfig.epochs): '''calculate Learning rate''' optimizer = self._get_optimizer(lr=lr) '''''' x_train_filenames, y_train_filenames = self._shuffle_data(x_train_filenames, y_train_filenames) for batch_index in range(step_per_epoch): '''load annotation and images''' images, annotation_gr = self._get_batch_sample( batch_index=batch_index, x_train_filenames=x_train_filenames, y_train_filenames=y_train_filenames) phi = self.calculate_adoptive_weight(epoch=epoch, batch_index=batch_index, y_train_filenames=y_train_filenames, weight_path=weight_path) '''convert to tensor''' images = tf.cast(images, tf.float32) annotation_gr = tf.cast(annotation_gr, tf.float32) '''train step''' loss_total, loss_low, loss_high = self.train_step( epoch=epoch, step=batch_index, total_steps=step_per_epoch, images=images, model=model, annotation_gr=annotation_gr, phi=phi, lambda_weight=self.lambda_weight, optimizer=optimizer, summary_writer=summary_writer, c_loss=c_loss) with summary_writer.as_default(): tf.summary.scalar('loss_total', loss_total, step=epoch) tf.summary.scalar('loss_low', loss_low, step=epoch) tf.summary.scalar('loss_high', loss_high, step=epoch) '''save weights''' model.save(self.save_path + str(epoch) + '_' + self.dataset_name + '.h5') # @tf.function def train_step(self, epoch, step, total_steps, images, model, annotation_gr, phi, optimizer, summary_writer, c_loss, lambda_weight): with tf.GradientTape() as tape: '''create annotation_predicted''' annotation_predicted = model(images, training=True) '''calculate loss''' loss_total, loss_low, loss_high = c_loss.acr_loss(x_pr=annotation_predicted, x_gt=annotation_gr, phi=phi, lambda_weight=lambda_weight, ds_name=self.dataset_name) '''calculate gradient''' gradients_of_model = tape.gradient(loss_total, model.trainable_variables) '''apply Gradients:''' optimizer.apply_gradients(zip(gradients_of_model, model.trainable_variables)) '''printing loss Values: ''' tf.print("->EPOCH: ", str(epoch), "->STEP: ", str(step) + '/' + str(total_steps), ' -> : LOSS: ', loss_total, ' -> : loss_low: ', loss_low, ' -> : loss_high: ', loss_high ) # print('==--==--==--==--==--==--==--==--==--') with summary_writer.as_default(): tf.summary.scalar('loss_total', loss_total, step=epoch) tf.summary.scalar('loss_low', loss_low, step=epoch) tf.summary.scalar('loss_high', loss_high, step=epoch) return loss_total, loss_low, loss_high def calculate_adoptive_weight(self, epoch, batch_index, y_train_filenames, weight_path): dt_utils = DataUtil(self.num_landmark) batch_y = y_train_filenames[ batch_index * LearningConfig.batch_size:(batch_index + 1) * LearningConfig.batch_size] asm_acc = None if 0 <= epoch <= 15: asm_acc = 80 elif 15 < epoch <= 30: asm_acc = 85 elif 30 < epoch <= 70: asm_acc = 90 elif 70 < epoch <= 100: asm_acc = 95 pn_batch = np.array([self._load_and_normalize(self.annotation_path + file_name) for file_name in batch_y]) pn_batch_asm = np.array([dt_utils.get_asm(input=self._load_and_normalize(self.annotation_path + file_name), dataset_name=self.dataset_name, accuracy=asm_acc) for file_name in batch_y]) delta = np.array(abs(pn_batch - pn_batch_asm)) phi = np.array([delta[i] / np.max(delta[i]) for i in range(len(pn_batch))]) # bs * num_lnd return phi def _get_optimizer(self, lr=1e-1, beta_1=0.9, beta_2=0.999, decay=1e-5): return tf.keras.optimizers.Adam(lr=lr, beta_1=beta_1, beta_2=beta_2, decay=decay) def make_model(self, arch, w_path): cnn = CNNModel() model = cnn.get_model(arch=arch, output_len=self.num_landmark) if w_path is not None: model.load_weights(w_path) return model def _shuffle_data(self, filenames, labels): filenames_shuffled, y_labels_shuffled = shuffle(filenames, labels) return filenames_shuffled, y_labels_shuffled def _create_generators(self, img_path=None, annotation_path=None): tf_utils = DataUtil(number_of_landmark=self.num_landmark) if img_path is None: filenames, labels = tf_utils.create_image_and_labels_name(img_path=self.img_path, annotation_path=self.annotation_path) else: filenames, labels = tf_utils.create_image_and_labels_name(img_path=img_path, annotation_path=annotation_path) filenames_shuffled, y_labels_shuffled = shuffle(filenames, labels) return filenames_shuffled, y_labels_shuffled def _get_batch_sample(self, batch_index, x_train_filenames, y_train_filenames, is_eval=False, batch_size=None): if is_eval: batch_x = x_train_filenames[ batch_index * batch_size:(batch_index + 1) * batch_size] batch_y = y_train_filenames[ batch_index * batch_size:(batch_index + 1) * batch_size] img_batch = np.array([imread(self.eval_img_path + file_name) for file_name in batch_x]) / 255.0 pn_batch = np.array([load(self.eval_annotation_path + file_name) for file_name in batch_y]) else: img_path = self.img_path pn_tr_path = self.annotation_path '''create batch data and normalize images''' batch_x = x_train_filenames[ batch_index * LearningConfig.batch_size:(batch_index + 1) * LearningConfig.batch_size] batch_y = y_train_filenames[ batch_index * LearningConfig.batch_size:(batch_index + 1) * LearningConfig.batch_size] '''create img and annotations''' img_batch = np.array([imread(img_path + file_name) for file_name in batch_x]) / 255.0 pn_batch = np.array([self._load_and_normalize(pn_tr_path + file_name) for file_name in batch_y]) return img_batch, pn_batch def _load_and_normalize(self, point_path): annotation = load(point_path) width = InputDataSize.image_input_size height = InputDataSize.image_input_size x_center = width / 2 y_center = height / 2 annotation_norm = [] for p in range(0, len(annotation), 2): annotation_norm.append((x_center - annotation[p]) / width) annotation_norm.append((y_center - annotation[p + 1]) / height) return annotation_norm