external/landmark_detection/conf/alignment.py

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