import os.path as osp from .base import Base class Alignment(Base): """ Alignment configure file, which contains training parameters of alignment. """ def __init__(self, args): super(Alignment, self).__init__('alignment') self.ckpt_dir = '/mnt/workspace/humanAIGC/project/STAR/weights' self.net = 'stackedHGnet_v1' self.nstack = 4 self.loader_type = 'alignment' self.data_definition = '300W' # COFW, 300W, WFLW self.test_file = 'test.tsv' # image self.channels = 3 self.width = 256 self.height = 256 self.means = (127.5, 127.5, 127.5) self.scale = 1 / 127.5 self.aug_prob = 1.0 self.display_iteration = 10 self.val_epoch = 1 self.valset = 'test.tsv' self.norm_type = 'default' self.encoder_type = 'default' self.decoder_type = 'default' # scheduler & optimizer self.milestones = [200, 350, 450] self.max_epoch = 260 self.optimizer = 'adam' self.learn_rate = 0.001 self.weight_decay = 0.00001 self.betas = [0.9, 0.999] self.gamma = 0.1 # batch_size & workers self.batch_size = 32 self.train_num_workers = 16 self.val_batch_size = 32 self.val_num_workers = 16 self.test_batch_size = 16 self.test_num_workers = 0 # tricks self.ema = True self.add_coord = True self.use_AAM = True # loss self.loss_func = 'STARLoss_v2' # STAR Loss paras self.star_w = 1 self.star_dist = 'smoothl1' self.init_from_args(args) # COFW if self.data_definition == 'COFW': self.edge_info = ( (True, (0, 4, 2, 5)), # RightEyebrow (True, (1, 6, 3, 7)), # LeftEyebrow (True, (8, 12, 10, 13)), # RightEye (False, (9, 14, 11, 15)), # LeftEye (True, (18, 20, 19, 21)), # Nose (True, (22, 26, 23, 27)), # LowerLip (True, (22, 24, 23, 25)), # UpperLip ) if self.norm_type == 'ocular': self.nme_left_index = 8 # ocular self.nme_right_index = 9 # ocular elif self.norm_type in ['pupil', 'default']: self.nme_left_index = 16 # pupil self.nme_right_index = 17 # pupil else: raise NotImplementedError self.classes_num = [29, 7, 29] self.crop_op = True self.flip_mapping = ( [0, 1], [4, 6], [2, 3], [5, 7], [8, 9], [10, 11], [12, 14], [16, 17], [13, 15], [18, 19], [22, 23], ) self.image_dir = osp.join(self.image_dir, 'COFW') # 300W elif self.data_definition == '300W': self.edge_info = ( (False, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16)), # FaceContour (False, (17, 18, 19, 20, 21)), # RightEyebrow (False, (22, 23, 24, 25, 26)), # LeftEyebrow (False, (27, 28, 29, 30)), # NoseLine (False, (31, 32, 33, 34, 35)), # Nose (True, (36, 37, 38, 39, 40, 41)), # RightEye (True, (42, 43, 44, 45, 46, 47)), # LeftEye (True, (48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59)), # OuterLip (True, (60, 61, 62, 63, 64, 65, 66, 67)), # InnerLip ) if self.norm_type in ['ocular', 'default']: self.nme_left_index = 36 # ocular self.nme_right_index = 45 # ocular elif self.norm_type == 'pupil': self.nme_left_index = [36, 37, 38, 39, 40, 41] # pupil self.nme_right_index = [42, 43, 44, 45, 46, 47] # pupil else: raise NotImplementedError self.classes_num = [68, 9, 68] self.crop_op = True self.flip_mapping = ( [0, 16], [1, 15], [2, 14], [3, 13], [4, 12], [5, 11], [6, 10], [7, 9], [17, 26], [18, 25], [19, 24], [20, 23], [21, 22], [31, 35], [32, 34], [36, 45], [37, 44], [38, 43], [39, 42], [40, 47], [41, 46], [48, 54], [49, 53], [50, 52], [61, 63], [60, 64], [67, 65], [58, 56], [59, 55], ) self.image_dir = osp.join(self.image_dir, '300W') # self.image_dir = osp.join(self.image_dir, '300VW_images') # 300VW elif self.data_definition == '300VW': self.edge_info = ( (False, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16)), # FaceContour (False, (17, 18, 19, 20, 21)), # RightEyebrow (False, (22, 23, 24, 25, 26)), # LeftEyebrow (False, (27, 28, 29, 30)), # NoseLine (False, (31, 32, 33, 34, 35)), # Nose (True, (36, 37, 38, 39, 40, 41)), # RightEye (True, (42, 43, 44, 45, 46, 47)), # LeftEye (True, (48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59)), # OuterLip (True, (60, 61, 62, 63, 64, 65, 66, 67)), # InnerLip ) if self.norm_type in ['ocular', 'default']: self.nme_left_index = 36 # ocular self.nme_right_index = 45 # ocular elif self.norm_type == 'pupil': self.nme_left_index = [36, 37, 38, 39, 40, 41] # pupil self.nme_right_index = [42, 43, 44, 45, 46, 47] # pupil else: raise NotImplementedError self.classes_num = [68, 9, 68] self.crop_op = True self.flip_mapping = ( [0, 16], [1, 15], [2, 14], [3, 13], [4, 12], [5, 11], [6, 10], [7, 9], [17, 26], [18, 25], [19, 24], [20, 23], [21, 22], [31, 35], [32, 34], [36, 45], [37, 44], [38, 43], [39, 42], [40, 47], [41, 46], [48, 54], [49, 53], [50, 52], [61, 63], [60, 64], [67, 65], [58, 56], [59, 55], ) self.image_dir = osp.join(self.image_dir, '300VW_Dataset_2015_12_14') # WFLW elif self.data_definition == 'WFLW': self.edge_info = ( (False, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32)), # FaceContour (True, (33, 34, 35, 36, 37, 38, 39, 40, 41)), # RightEyebrow (True, (42, 43, 44, 45, 46, 47, 48, 49, 50)), # LeftEyebrow (False, (51, 52, 53, 54)), # NoseLine (False, (55, 56, 57, 58, 59)), # Nose (True, (60, 61, 62, 63, 64, 65, 66, 67)), # RightEye (True, (68, 69, 70, 71, 72, 73, 74, 75)), # LeftEye (True, (76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87)), # OuterLip (True, (88, 89, 90, 91, 92, 93, 94, 95)), # InnerLip ) if self.norm_type in ['ocular', 'default']: self.nme_left_index = 60 # ocular self.nme_right_index = 72 # ocular elif self.norm_type == 'pupil': self.nme_left_index = 96 # pupils self.nme_right_index = 97 # pupils else: raise NotImplementedError self.classes_num = [98, 9, 98] self.crop_op = True self.flip_mapping = ( [0, 32], [1, 31], [2, 30], [3, 29], [4, 28], [5, 27], [6, 26], [7, 25], [8, 24], [9, 23], [10, 22], [11, 21], [12, 20], [13, 19], [14, 18], [15, 17], # cheek [33, 46], [34, 45], [35, 44], [36, 43], [37, 42], [38, 50], [39, 49], [40, 48], [41, 47], # elbrow [60, 72], [61, 71], [62, 70], [63, 69], [64, 68], [65, 75], [66, 74], [67, 73], [55, 59], [56, 58], [76, 82], [77, 81], [78, 80], [87, 83], [86, 84], [88, 92], [89, 91], [95, 93], [96, 97]) self.image_dir = osp.join(self.image_dir, 'WFLW', 'WFLW_images') self.label_num = self.nstack * 3 if self.use_AAM else self.nstack self.loss_weights, self.criterions, self.metrics = [], [], [] for i in range(self.nstack): factor = (2**i) / (2**(self.nstack - 1)) if self.use_AAM: self.loss_weights += [ factor * weight for weight in [1.0, 10.0, 10.0] ] self.criterions += [self.loss_func, 'AWingLoss', 'AWingLoss'] self.metrics += ['NME', None, None] else: self.loss_weights += [factor * weight for weight in [1.0]] self.criterions += [ self.loss_func, ] self.metrics += [ 'NME', ] self.key_metric_index = (self.nstack - 1) * 3 if self.use_AAM else ( self.nstack - 1) # data self.folder = self.get_foldername() self.work_dir = osp.join(self.ckpt_dir, self.data_definition, self.folder) self.model_dir = osp.join(self.work_dir, 'model') self.log_dir = osp.join(self.work_dir, 'log') self.train_tsv_file = osp.join(self.annot_dir, self.data_definition, 'train.tsv') self.train_pic_dir = self.image_dir self.val_tsv_file = osp.join(self.annot_dir, self.data_definition, self.valset) self.val_pic_dir = self.image_dir self.test_tsv_file = osp.join(self.annot_dir, self.data_definition, self.test_file) self.test_pic_dir = self.image_dir # self.train_tsv_file = osp.join(self.annot_dir, '300VW', "train.tsv") # self.train_pic_dir = self.image_dir # self.val_tsv_file = osp.join(self.annot_dir, '300VW', self.valset) # self.val_pic_dir = self.image_dir # self.test_tsv_file = osp.join(self.annot_dir, '300VW', self.test_file) # self.test_pic_dir = self.image_dir def get_foldername(self): str = '' str += '{}_{}x{}_{}_ep{}_lr{}_bs{}'.format( self.data_definition, self.height, self.width, self.optimizer, self.max_epoch, self.learn_rate, self.batch_size) str += '_{}'.format(self.loss_func) str += '_{}_{}'.format( self.star_dist, self.star_w) if self.loss_func == 'STARLoss' else '' str += '_AAM' if self.use_AAM else '' str += '_{}'.format( self.valset[:-4]) if self.valset != 'test.tsv' else '' str += '_{}'.format(self.id) return str