"""
@author: JiXuan Xu, Jun Wang
@date: 20201023
@contact: jun21wangustc@gmail.com 
"""
import sys
import logging.config
logging.config.fileConfig("FaceX-Zoo/face_sdk/config/logging.conf")
logger = logging.getLogger('sdk')

import torch

sys.path.append('/content/FaceX-Zoo/face_sdk/core/model_loader')
from BaseModelLoader import BaseModelLoader

class FaceAlignModelLoader(BaseModelLoader):
    def __init__(self, model_path, model_category, model_name, meta_file='model_meta.json'):
        logger.info('Start to analyze the face landmark model, model path: %s, model category: %s，model name: %s' %
                    (model_path, model_category, model_name))
        super().__init__(model_path, model_category, model_name, meta_file)
        self.cfg['img_size'] = self.meta_conf['input_width']
        
    # https://github.com/pytorch/pytorch/issues/3678
    def load_model(self):
        try:
            model = torch.load(self.cfg['model_file_path']) 
        except Exception as e:
            logger.error('The model failed to load, please check the model path: %s!'
                         % self.cfg['model_file_path'])
            raise e
        else:
            logger.info('Successfully loaded the face landmark model!')
            return model, self.cfg
