|
from configuration import DatasetName, WflwConf, W300Conf, DatasetType, LearningConfig, InputDataSize |
|
from cnn_model import CNNModel |
|
import tensorflow as tf |
|
import tensorflow.keras as 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 image_utility import ImageUtility |
|
from tqdm import tqdm |
|
import os |
|
from Asm_assisted_loss import ASMLoss |
|
from cnn_model import CNNModel |
|
|
|
|
|
class Train: |
|
def __init__(self, arch, dataset_name, save_path, asm_accuracy=90): |
|
""" |
|
:param arch: |
|
:param dataset_name: |
|
:param save_path: |
|
:param asm_accuracy: |
|
""" |
|
|
|
self.dataset_name = dataset_name |
|
self.save_path = save_path |
|
self.arch = arch |
|
self.asm_accuracy = asm_accuracy |
|
|
|
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 |
|
self.pose_path = W300Conf.train_pose |
|
|
|
if dataset_name == DatasetName.wflw: |
|
self.num_landmark = WflwConf.num_of_landmarks * 2 |
|
self.img_path = WflwConf.train_image |
|
self.annotation_path = WflwConf.train_annotation |
|
self.pose_path = WflwConf.train_pose |
|
|
|
def train(self, weight_path): |
|
""" |
|
|
|
:param weight_path: |
|
:return: |
|
""" |
|
|
|
'''create loss''' |
|
c_loss = ASMLoss(dataset_name=self.dataset_name, accuracy=90) |
|
cnn = CNNModel() |
|
'''making models''' |
|
model = cnn.get_model(arch=self.arch, output_len=self.num_landmark) |
|
if weight_path is not None: |
|
model.load_weights(weight_path) |
|
|
|
'''create sample generator''' |
|
image_names, landmark_names, pose_names = self._create_generators() |
|
|
|
'''create train configuration''' |
|
step_per_epoch = len(image_names) // LearningConfig.batch_size |
|
|
|
'''start train:''' |
|
optimizer = tf.keras.optimizers.Adam(lr=1e-2, decay=1e-5) |
|
for epoch in range(LearningConfig.epochs): |
|
image_names, landmark_names, pose_names = shuffle(image_names, landmark_names, pose_names) |
|
for batch_index in range(step_per_epoch): |
|
'''load annotation and images''' |
|
images, annotation_gr, poses_gr = self._get_batch_sample( |
|
batch_index=batch_index, |
|
img_filenames=image_names, |
|
landmark_filenames=landmark_names, |
|
pose_filenames=pose_names) |
|
|
|
'''convert to tensor''' |
|
images = tf.cast(images, tf.float32) |
|
annotation_gr = tf.cast(annotation_gr, tf.float32) |
|
poses_gr = tf.cast(poses_gr, tf.float32) |
|
|
|
'''train step''' |
|
self.train_step(epoch=epoch, |
|
step=batch_index, |
|
total_steps=step_per_epoch, |
|
model=model, |
|
images=images, |
|
annotation_gt=annotation_gr, |
|
poses_gt=poses_gr, |
|
optimizer=optimizer, |
|
c_loss=c_loss) |
|
'''save weights''' |
|
model.save(self.save_path + self.arch + str(epoch) + '_' + self.dataset_name) |
|
|
|
def train_step(self, epoch, step, total_steps, model, images, annotation_gt, poses_gt, optimizer, c_loss): |
|
""" |
|
|
|
:param epoch: |
|
:param step: |
|
:param total_steps: |
|
:param model: |
|
:param images: |
|
:param annotation_gt: |
|
:param poses_gt: |
|
:param optimizer: |
|
:param c_loss: |
|
:return: |
|
""" |
|
|
|
with tf.GradientTape() as tape: |
|
'''create annotation_predicted''' |
|
annotation_predicted, pose_predicted = model(images, training=True) |
|
'''calculate loss''' |
|
mse_loss, asm_loss = c_loss.calculate_landmark_ASM_assisted_loss(landmark_pr=annotation_predicted, |
|
landmark_gt=annotation_gt, |
|
current_epoch=epoch, |
|
total_steps=total_steps) |
|
pose_loss = c_loss.calculate_pose_loss(x_pr=pose_predicted, x_gt=poses_gt) |
|
|
|
'''calculate loss''' |
|
total_loss = mse_loss + asm_loss + pose_loss |
|
|
|
'''calculate gradient''' |
|
gradients_of_model = tape.gradient(total_loss, 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), ' -> : total_loss: ', |
|
total_loss) |
|
|
|
def _create_generators(self): |
|
""" |
|
:return: |
|
""" |
|
image_names, landmark_filenames, pose_names = \ |
|
self._create_image_and_labels_name(img_path=self.img_path, |
|
annotation_path=self.annotation_path, |
|
pose_path=self.pose_path) |
|
return image_names, landmark_filenames, pose_names |
|
|
|
def _create_image_and_labels_name(self, img_path, annotation_path, pose_path): |
|
""" |
|
|
|
:param img_path: |
|
:param annotation_path: |
|
:param pose_path: |
|
:return: |
|
""" |
|
img_filenames = [] |
|
landmark_filenames = [] |
|
poses_filenames = [] |
|
|
|
for file in os.listdir(img_path): |
|
if file.endswith(".jpg") or file.endswith(".png"): |
|
lbl_file = str(file)[:-3] + "npy" |
|
pose_file = str(file)[:-3] + "npy" |
|
if os.path.exists(annotation_path + lbl_file) and os.path.exists(pose_path + lbl_file): |
|
img_filenames.append(str(file)) |
|
landmark_filenames.append(lbl_file) |
|
poses_filenames.append(pose_file) |
|
|
|
return np.array(img_filenames), np.array(landmark_filenames), np.array(poses_filenames) |
|
|
|
def _get_batch_sample(self, batch_index, img_filenames, landmark_filenames, pose_filenames): |
|
""" |
|
:param batch_index: |
|
:param img_filenames: |
|
:param landmark_filenames: |
|
:param pose_filenames: |
|
:return: |
|
""" |
|
|
|
'''create batch data and normalize images''' |
|
batch_img = img_filenames[ |
|
batch_index * LearningConfig.batch_size:(batch_index + 1) * LearningConfig.batch_size] |
|
batch_lnd = landmark_filenames[ |
|
batch_index * LearningConfig.batch_size:(batch_index + 1) * LearningConfig.batch_size] |
|
batch_pose = pose_filenames[ |
|
batch_index * LearningConfig.batch_size:(batch_index + 1) * LearningConfig.batch_size] |
|
'''create img and annotations''' |
|
img_batch = np.array([imread(self.img_path + file_name) for file_name in batch_img]) / 255.0 |
|
lnd_batch = np.array([self._load_and_normalize(self.annotation_path + file_name) for file_name in batch_lnd]) |
|
pose_batch = np.array([load(self.pose_path + file_name) for file_name in batch_pose]) |
|
|
|
return img_batch, lnd_batch, pose_batch |
|
|
|
def _load_and_normalize(self, point_path): |
|
""" |
|
:param point_path: |
|
:return: |
|
""" |
|
|
|
annotation = load(point_path) |
|
'''normalize landmarks''' |
|
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 |
|
|
