# @Author  : lightXu
# @File    : tf_settings.py
# @Time    : 2018/11/22 0022 上午 10:41
import os
from segment.sheet_resolve.tools.utils import read_label

subject_list = ['math', 'math_zxhx', 'english', 'chinese',
                'physics', 'chemistry', 'biology', 'politics', 'history',
                'geography', 'science_comprehensive', 'arts_comprehensive', 'cloze', 'choice']

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

xml_template_path = os.path.join(BASE_DIR, 'labels', '000000-template.xml')
images_dir_path = os.path.join(BASE_DIR, 'images')
model_dir_path = os.path.join(BASE_DIR, 'model')
label_dir = os.path.join(BASE_DIR, 'model', 'subjects', 'labels')

# ssd
choice_m_ssd = os.path.join(model_dir_path, 'ssd_choice_m', 'frozen_inference_graph.pb')
choice_m_ssd_label = os.path.join(model_dir_path, 'ssd_choice_m', 'choice_m_label_map.pbtxt')

exam_number_ssd = os.path.join(model_dir_path, 'ssd_exam_number', 'frozen_inference_graph.pb')
exam_number_ssd_label = os.path.join(model_dir_path, 'ssd_exam_number', 'exam_number_label_map.pbtxt')

# faster_rcnn model
choice_m_model_path = os.path.join(model_dir_path, 'choice_m', 'choice_m' + '.ckpt')
choice_model_path = os.path.join(model_dir_path, 'choice', 'choice' + '.ckpt')
cloze_model_path = os.path.join(model_dir_path, 'cloze', 'cloze' + '.ckpt')
math_zxhx_model_path = os.path.join(model_dir_path, 'math_zxhx', 'sheet' + '.ckpt')
math_zxhx_detail_model_path = os.path.join(model_dir_path, 'math_zxhx_detail', 'sheet' + '.ckpt')
# 第三方
math_model_path = os.path.join(model_dir_path, 'subjects', 'math', 'sheet' + '.ckpt')
english_model_path = os.path.join(model_dir_path, 'subjects', 'english', 'sheet' + '.ckpt')
chinese_model_path = os.path.join(model_dir_path, 'subjects', 'chinese', 'sheet' + '.ckpt')
physics_model_path = os.path.join(model_dir_path, 'subjects', 'physics', 'sheet' + '.ckpt')
chemistry_model_path = os.path.join(model_dir_path, 'subjects', 'chemistry', 'sheet' + '.ckpt')
biology_model_path = os.path.join(model_dir_path, 'subjects', 'biology', 'sheet' + '.ckpt')
politics_model_path = os.path.join(model_dir_path, 'subjects', 'politics', 'sheet' + '.ckpt')
history_model_path = os.path.join(model_dir_path, 'subjects', 'history', 'sheet' + '.ckpt')
geography_model_path = os.path.join(model_dir_path, 'subjects', 'geography', 'sheet' + '.ckpt')
science_comprehensive_model_path = os.path.join(model_dir_path, 'subjects', 'science_comprehensive', 'sheet' + '.ckpt')
arts_comprehensive_model_path = os.path.join(model_dir_path, 'subjects', 'arts_comprehensive', 'sheet' + '.ckpt')

math_blank_model_path = os.path.join(model_dir_path, 'subjects', 'math_blank', 'sheet' + '.ckpt')
english_blank_model_path = os.path.join(model_dir_path, 'subjects', 'english_blank', 'sheet' + '.ckpt')
chinese_blank_model_path = os.path.join(model_dir_path, 'subjects', 'chinese_blank', 'sheet' + '.ckpt')
physics_blank_model_path = os.path.join(model_dir_path, 'subjects', 'physics_blank', 'sheet' + '.ckpt')
chemistry_blank_model_path = os.path.join(model_dir_path, 'subjects', 'chemistry_blank', 'sheet' + '.ckpt')
biology_blank_model_path = os.path.join(model_dir_path, 'subjects', 'biology_blank', 'sheet' + '.ckpt')
politics_blank_model_path = os.path.join(model_dir_path, 'subjects', 'politics_blank', 'sheet' + '.ckpt')
history_blank_model_path = os.path.join(model_dir_path, 'subjects', 'history_blank', 'sheet' + '.ckpt')
geography_blank_model_path = os.path.join(model_dir_path, 'subjects', 'geography_blank', 'sheet' + '.ckpt')
science_comprehensive_blank_model_path = os.path.join(model_dir_path, 'subjects', 'science_comprehensive_blank', 'sheet' + '.ckpt')
arts_comprehensive_blank_model_path = os.path.join(model_dir_path, 'subjects', 'arts_comprehensive_blank', 'sheet' + '.ckpt')

choice_m_classes = ('__background__', 'choice_m')
choice_classes = ('__background__', 'choice', 'choice', 'choice', 'choice', 'choice', 'choice', 'choice', 'choice',
                  'choice', 'choice', 'choice', 'choice')
cloze_classes = ('__background__', 'cloze')
math_zxhx_classes = ('__background__', 'info_title', 'print_info', 'exam_number',
                     'choice', 'cloze', 'solve', 'solve0', 'qr_code', 'bar_code', 'page',)

math_zxhx_detail_classes = (
    '__background__', 'alarm_info', 'info_title', 'attention', 'page', 'full_filling', 'print_info',
    'ban_area', 'type_score', 'time', 'total_score', 'executor', 'verify', 'name_w', 'school_w',
    'class_w', 'student_info_w', 'exam_number_w', 'room_w', 'cloze', 'cloze_s', 'cloze_score',
    'solve', 'solve0', 'seal_area', 'score_collect', 'seat_w', 'student_info', 'qr_code',
    'class', 'exam_number', 'exam_number_s', 'bar_code', 'choice', 'choice_s', 'lack',
    'select_s', 'select_b', 'type', 'mark',)

math_classes = read_label(label_dir, "math")
english_classes = read_label(label_dir, "english")
chinese_classes = read_label(label_dir, "chinese")
physics_classes = read_label(label_dir, "physics")
chemistry_classes = read_label(label_dir, "chemistry")
biology_classes = read_label(label_dir, "biology")
politics_classes = read_label(label_dir, "politics")
history_classes = read_label(label_dir, "history")
geography_classes = read_label(label_dir, "geography")
science_comprehensive_classes = read_label(label_dir, "science_comprehensive")
arts_comprehensive_classes = read_label(label_dir, "arts_comprehensive")

# math_blank_classes = read_label(label_dir, "math_blank")
# english_blank_classes = read_label(label_dir, "english_blank")
chinese_blank_classes = read_label(label_dir, "chinese_blank")
# physics_blank_classes = read_label(label_dir, "physics_blank")
# chemistry_blank_classes = read_label(label_dir, "chemistry_blank")
# biology_blank_classes = read_label(label_dir, "biology_blank")
# politics_blank_classes = read_label(label_dir, "politics_blank")
# history_blank_classes = read_label(label_dir, "history_blank")
# geography_blank_classes = read_label(label_dir, "geography_blank")
science_comprehensive_blank_classes = read_label(label_dir, "science_comprehensive_blank")
arts_comprehensive_blank_classes = read_label(label_dir, "arts_comprehensive_blank")

sheet_detail_bias_dict = {
    'alarm_info': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'info_title': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'attention': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'page': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'full_filling': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'print_info': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'ban_area': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'type_score': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'time': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'total_score': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'executor': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'verify': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'name_w': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'school_w': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'class_w': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'student_info_w': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'exam_number_w': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'room_w': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'cloze': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'cloze_s': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'cloze_score': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'solve': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'solve0': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'seal_area': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'score_collect': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'seat_w': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'student_info': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'qr_code': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'class': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'exam_number': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'exam_number_s': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'bar_code': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'choice': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'choice_s': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'lack': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'select_s': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'select_b': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'mark': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'type': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'exam_add': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'subject': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'judge_s': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'type_info': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'seat_number': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'judge': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'composition': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'correction': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'class_s': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'table': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'composition_info': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'table_s': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'composition0': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'name_s': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0},
    'choice_m': {"xmin_bias": 0, "ymin_bias": 0, "xmax_bias": 0, "ymax_bias": 0}, }

choice_coordinate_bias_dict = {}
cloze_coordinate_bias_dict = {}
math_zxhx_bias_dict = {}
math_zxhx_detail_bias_dict = {}
math_bias_dict = {}
english_bias_dict = {}
chinese_bias_dict = {}
physics_bias_dict = {}
chemistry_bias_dict = {}
biology_bias_dict = {}
politics_bias_dict = {}
history_bias_dict = {}
geography_bias_dict = {}
science_comprehensive_bias_dict = {}
arts_comprehensive_bias_dict = {}
choice_m_bias_dict = {}

model_dict = {
    'choice_ssd': {'path': choice_m_ssd, 'classes': choice_m_ssd_label},
    'exam_number_ssd': {'path': exam_number_ssd, 'classes': exam_number_ssd_label},
    'choice_m': {'path': choice_m_model_path, 'anchor_scales': (2, 4, 8, 16), 'anchor_ratios': (0.5, 1.0, 2, 4),
                 'classes': choice_m_classes, 'class_coordinate_bias': choice_m_bias_dict},
    'choice': {'path': choice_model_path, 'anchor_scales': (8, 16, 32), 'anchor_ratios': (0.5, 1.0, 1.5),
               'classes': choice_classes, 'class_coordinate_bias': choice_coordinate_bias_dict},
    'cloze': {'path': cloze_model_path, 'anchor_scales': (8, 16, 32), 'anchor_ratios': (0.5, 1.0, 1.5),
              'classes': cloze_classes, 'class_coordinate_bias': cloze_coordinate_bias_dict},
    'math_zxhx': {'path': math_zxhx_model_path, 'anchor_scales': (2, 4, 8, 16, 32),
                  'anchor_ratios': (0.5, 1, 1.5, 2, 4),
                  'classes': math_zxhx_classes, 'class_coordinate_bias': math_zxhx_bias_dict},
    'math_zxhx_detail': {'path': math_zxhx_detail_model_path, 'anchor_scales': (4, 8, 16, 32),
                         'anchor_ratios': (0.5, 1, 2, 3), 'classes': math_zxhx_detail_classes,
                         'class_coordinate_bias': math_zxhx_detail_bias_dict},
    # 'math': {'path': math_model_path, 'anchor_scales': (4, 8, 16, 32), 'anchor_ratios': (0.5, 1, 2, 3),
    #          'classes': math_classes, 'class_coordinate_bias': math_bias_dict},
    'math': dict(path=math_model_path, anchor_scales=(1, 2, 4, 8, 16), anchor_ratios=(0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                 classes=math_classes, class_coordinate_bias=math_bias_dict),
    'english': {'path': english_model_path,
                'anchor_scales': (1, 2, 4, 8, 16),
                'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                'classes': english_classes, 'class_coordinate_bias': english_bias_dict},
    'chinese': {'path': chinese_model_path,
                'anchor_scales': (1, 2, 4, 8, 16),
                'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                'classes': chinese_classes, 'class_coordinate_bias': chinese_bias_dict},
    'physics': {'path': physics_model_path,
                'anchor_scales': (1, 2, 4, 8, 16),
                'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                'classes': physics_classes, 'class_coordinate_bias': physics_bias_dict},
    'chemistry': {'path': chemistry_model_path,
                  'anchor_scales': (1, 2, 4, 8, 16),
                  'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                  'classes': chemistry_classes, 'class_coordinate_bias': chemistry_bias_dict},
    'biology': {'path': biology_model_path,
                'anchor_scales': (1, 2, 4, 8, 16),
                'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                'classes': biology_classes, 'class_coordinate_bias': biology_bias_dict},
    'politics': {'path': politics_model_path,
                 'anchor_scales': (1, 2, 4, 8, 16),
                 'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                 'classes': politics_classes,
                 'class_coordinate_bias': politics_bias_dict},
    'history': {'path': history_model_path,
                'anchor_scales': (1, 2, 4, 8, 16),
                'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                'classes': history_classes, 'class_coordinate_bias': history_bias_dict},
    'geography': {'path': geography_model_path,
                  'anchor_scales': (1, 2, 4, 8, 16),
                  'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                  'classes': geography_classes, 'class_coordinate_bias': geography_bias_dict},
    'science_comprehensive': {'path': science_comprehensive_model_path,
                              'anchor_scales': (1, 2, 4, 8, 16),
                              'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                              'classes': science_comprehensive_classes,
                              'class_coordinate_bias': science_comprehensive_bias_dict},
    'arts_comprehensive': {'path': arts_comprehensive_model_path,
                           'anchor_scales': (1, 2, 4, 8, 16),
                           'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                           'classes': arts_comprehensive_classes,
                           'class_coordinate_bias': arts_comprehensive_bias_dict},

    # 'math_blank': dict(path=math_blank_model_path, anchor_scales=(1, 2, 4, 8, 16),
    #                    anchor_ratios=(0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
    #                    classes=math_blank_classes, class_coordinate_bias=math_bias_dict),
    # 'english_blank': {'path': english_blank_model_path,
    #                   'anchor_scales': (1, 2, 4, 8, 16),
    #                   'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
    #                   'classes': english_blank_classes, 'class_coordinate_bias': english_bias_dict},
    'chinese_blank': {'path': chinese_blank_model_path,
                      'anchor_scales': (1, 2, 4, 8, 16),
                      'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                      'classes': chinese_blank_classes, 'class_coordinate_bias': chinese_bias_dict},
    # 'physics_blank': {'path': physics_blank_model_path,
    #                   'anchor_scales': (1, 2, 4, 8, 16),
    #                   'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
    #                   'classes': physics_blank_classes, 'class_coordinate_bias': physics_bias_dict},
    # 'chemistry_blank': {'path': chemistry_blank_model_path,
    #                     'anchor_scales': (1, 2, 4, 8, 16),
    #                     'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
    #                     'classes': chemistry_blank_classes, 'class_coordinate_bias': chemistry_bias_dict},
    # 'biology_blank': {'path': biology_blank_model_path,
    #                   'anchor_scales': (1, 2, 4, 8, 16),
    #                   'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
    #                   'classes': biology_blank_classes, 'class_coordinate_bias': biology_bias_dict},
    # 'politics_blank': {'path': politics_blank_model_path,
    #                    'anchor_scales': (1, 2, 4, 8, 16),
    #                    'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
    #                    'classes': politics_blank_classes,
    #                    'class_coordinate_bias': politics_bias_dict},
    # 'history_blank': {'path': history_blank_model_path,
    #                   'anchor_scales': (1, 2, 4, 8, 16),
    #                   'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
    #                   'classes': history_blank_classes, 'class_coordinate_bias': history_bias_dict},
    # 'geography_blank': {'path': geography_blank_model_path,
    #                     'anchor_scales': (1, 2, 4, 8, 16),
    #                     'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
    #                     'classes': geography_blank_classes, 'class_coordinate_bias': geography_bias_dict},
    'science_comprehensive_blank': {'path': science_comprehensive_blank_model_path,
                                    'anchor_scales': (1, 2, 4, 8, 16),
                                    'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                                    'classes': science_comprehensive_blank_classes,
                                    'class_coordinate_bias': science_comprehensive_bias_dict},
    'arts_comprehensive_blank': {'path': arts_comprehensive_blank_model_path,
                                 'anchor_scales': (1, 2, 4, 8, 16),
                                 'anchor_ratios': (0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4),
                                 'classes': arts_comprehensive_blank_classes,
                                 'class_coordinate_bias': arts_comprehensive_bias_dict},
}
