exam-segment-django/segment/sheet_resolve/analysis/resolve.py

# @Author  : lightXu
# @File    : resolve.py
# @Time    : 2018/12/3 0003 上午 10:16

import time
import traceback
import xml.etree.cElementTree as ET
from django.conf import settings
import segment.logging_config as logging
import segment.sheet_resolve.analysis.choice.analysis_choice as resolve_choice
import segment.sheet_resolve.analysis.choice.choice_box as choice_box
import segment.sheet_resolve.analysis.choice.choice_line_box as choice_line_box
import segment.sheet_resolve.analysis.cloze.analysis_cloze as resolve_cloze
import segment.sheet_resolve.analysis.cloze.cloze_line_box as resolve_cloze_line_box
import segment.sheet_resolve.analysis.exam_number.exam_number_box as resolve_exam_number_box
import segment.sheet_resolve.analysis.exam_number.exam_number_row_column as exam_number_row_column
import segment.sheet_resolve.analysis.sheet.analysis_sheet as resolve_sheet
import segment.sheet_resolve.analysis.solve.mark_box as resolve_mark_box
import segment.sheet_resolve.analysis.solve.mark_line_box as resolve_mark_line_box
from segment.sheet_resolve.tools import utils
from segment.sheet_resolve.tools.tf_sess import TfSess
from segment.sheet_resolve.tools.tf_settings import xml_template_path, model_dict
from segment.sheet_resolve.tools.utils import read_single_img, read_xml_to_json, create_xml
from segment.sheet_resolve.analysis.sheet.sheet_adjust import adjust_item_edge_by_gray_image
from segment.sheet_resolve.analysis.sheet.sheet_infer import infer_bar_code, box_infer_and_complete, infer_solve
from segment.sheet_resolve.analysis.sheet.sheet_infer import infer_exam_number, adjust_exam_number, exam_number_infer_by_s
from segment.sheet_resolve.analysis.sheet.choice_infer import infer_choice_m
from segment.sheet_resolve.analysis.sheet.ocr_sheet import tell_columns, sheet_sorted

logger = logging.getLogger(settings.LOGGING_TYPE)


sheet_infer_dict = dict(bar_code=True,
                        choice_m=True,
                        exam_number=True,
                        common_sheet=False,
                        solve=True)
infer_choice_m_flag = False


def sheet(series_number, image_path, image, conf_thresh, mns_thresh, subject, sheet_sess, ocr=''):
    global infer_choice_m_flag
    model_type = subject
    classes = list(model_dict[model_type]['classes'])
    coordinate_bias_dict = model_dict[model_type]['class_coordinate_bias']

    if '_blank' in model_type:
        model_type = model_type.replace("_blank", "")

    sheets_dict = resolve_sheet.get_single_image_sheet_regions(model_type, image_path, image, classes,
                                                               sheet_sess.sess, sheet_sess.net,
                                                               conf_thresh, mns_thresh, coordinate_bias_dict)

    h, w = image.shape[0], image.shape[1]
    regions = sheets_dict['regions']
    fetched_class = [ele['class_name'] for ele in regions]

    try:
        regions = adjust_item_edge_by_gray_image(image, regions)
    except Exception as e:
        traceback.print_exc()
        logger.info('试卷：{} 自适应边框失败: {}'.format(image_path, e))

    # 分栏
    col_split_x = tell_columns(image, regions)

    if sheet_infer_dict['bar_code']:
        try:
            if ('bar_code' not in fetched_class) and ocr:
                attention_region = [ele for ele in regions if ele['class_name'] == 'attention']
                bar_code_list = infer_bar_code(image, ocr, attention_region)
                regions.extend(bar_code_list)
        except Exception as e:
            traceback.print_exc()
            logger.info('试卷：{} 条形码推断失败: {}'.format(image_path, e))

    if sheet_infer_dict['exam_number']:
        try:
            cond1 = 'exam_number' in fetched_class
            tmp = ['info_title', 'qr_code', 'bar_code', 'choice', 'choice_m', 'exam_number_w']
            cond2 = True in [True for ele in tmp if ele in fetched_class]  # 第一面特征
            cond3 = 'exam_number_w' in fetched_class
            cond4 = 'exam_number_s' in fetched_class

            # if cond1 and cond3 and not cond4:
            if cond1 and cond3:
                regions = adjust_exam_number(regions)
            if not cond1 and cond4:
                exam_number_list = exam_number_infer_by_s(image, regions)
                regions.extend(exam_number_list)

            if not cond1 and not cond4 and cond2 and ocr:
                exam_number_list = infer_exam_number(image, ocr, regions)
                regions.extend(exam_number_list)

        except Exception as e:
            traceback.print_exc()
            logger.info('试卷：{} 考号推断失败: {}'.format(image_path, e))

    if sheet_infer_dict['choice_m']:

        try:
            choice_m_list = infer_choice_m(image, regions, col_split_x, ocr)
            #remain_choice_m = []
            if len(choice_m_list) > 0:
                choice_m_old_list = [ele for ele in regions if 'choice_m' == ele['class_name']]
                for infer_box in choice_m_list.copy():
                    infer_loc = infer_box['bounding_box']

                    for tf_box in choice_m_old_list:
                        tf_loc = tf_box['bounding_box']
                        iou = utils.cal_iou(infer_loc, tf_loc)
                        # if iou[0] > 0.70 or iou[1] > 0.70 or iou[2] > 0.70:
                        # if iou[0] > 0.70 or iou[2] > 0.70:
                        if iou[0] > 0.85:
                            # if infer_box not in remain_choice_m:
                            #     remain_choice_m.append(infer_box)
                            #     choice_m_list.remove(infer_box)
                            regions.remove(tf_box)
                            # break
                        elif iou[0] > 0:
                            choice_m_list.remove(infer_box)
                            break

                #remain_choice_m.extend(choice_m_list)

                # regions = [ele for ele in regions if 'choice_m' != ele['class_name']]
                # regions.extend(remain_choice_m)
                regions.extend(choice_m_list)
                infer_choice_m_flag = True

        except Exception as e:
            traceback.print_exc()
            logger.info('试卷：{} 选择题推断失败: {}'.format(image_path, e))

    if sheet_infer_dict['solve']:
        try:
            include_class = ['info_title',
                             'bar_code',
                             'choice_m',
                             'cloze',
                             'cloze_s',
                             'exam_number',
                             'solve',
                             'composition',
                             'correction'
                             ]
            regions_subset = [ele for ele in regions if ele['class_name'] in include_class]
            col_regions = sheet_sorted(regions_subset, col_split_x)

            top = min([ele['bounding_box']['ymin'] for ele in regions])
            bottom = max([ele['bounding_box']['ymax'] for ele in regions])

            seal_area = [ele for ele in regions if 'seal' in ele['class_name']]
            if len(seal_area) > 0:
                right, left = w, 1
                for ele in seal_area:
                    if ele['bounding_box']['xmax'] > w // 2:
                        right = ele['bounding_box']['xmin']
                    if ele['bounding_box']['xmax'] < w // 2:
                        left = ele['bounding_box']['xmax']
            else:
                left = min([ele['bounding_box']['xmin'] for ele in regions])
                right = max([ele['bounding_box']['xmax'] for ele in regions])

            solve_regions = infer_solve(regions, left, right, top, bottom, col_regions, col_split_x)
            regions.append(solve_regions)
        except Exception as e:
            traceback.print_exc()
            logger.info('试卷：{} 解答题补全推断失败: {}'.format(image_path, e))

    if sheet_infer_dict['common_sheet']:

        try:
            regions = box_infer_and_complete(image, regions, ocr)
        except Exception as e:
            traceback.print_exc()
            logger.info('试卷：{} 识别框补全推断失败: {}'.format(image_path, e))

    try:
        adjust_regions = adjust_item_edge_by_gray_image(image, regions)
    except Exception as e:
        adjust_regions = regions

        traceback.print_exc()
        logger.info('试卷：{} 自适应边框失败: {}'.format(image_path, e))

    sheets_dict.update({'regions': adjust_regions})

    #  generate xml
    tree = ET.parse(xml_template_path)
    xml_save_path = sheets_dict['img_name'].replace('.jpg', '.xml')
    root = tree.getroot()
    series = ET.SubElement(root, 'paper_id')
    series.text = series_number

    img_shape = image.shape
    project = ET.SubElement(root, 'size', {})
    width = ET.SubElement(project, 'width')
    width.text = str(img_shape[1])
    height = ET.SubElement(project, 'height')
    height.text = str(img_shape[0])
    depth = ET.SubElement(project, 'depth')
    if len(img_shape) >= 3:
        depth.text = '3'
    else:
        depth.text = '1'

    for ele in regions:
        name = ele['class_name']
        xmin = ele['bounding_box']['xmin']
        ymin = ele['bounding_box']['ymin']
        xmax = ele['bounding_box']['xmax']
        ymax = ele['bounding_box']['ymax']
        tree = create_xml(name, tree, xmin, ymin, xmax, ymax)

    tree.write(xml_save_path)
    return sheets_dict, xml_save_path


def choice(image, regions, xml_path, conf_thresh, mns_thresh, choice_sess):
    model_type = 'choice'
    classes = model_dict[model_type]['classes']
    coordinate_bias_dict = model_dict[model_type]['class_coordinate_bias']

    choice_list = []
    for ele in regions:
        if ele["class_name"] == 'choice':

            choice_bbox = ele['bounding_box']
            left = choice_bbox['xmin']
            top = choice_bbox['ymin']
            choice_img = utils.crop_region(image, choice_bbox)

            choice_dict_tf = resolve_choice. \
                get_single_image_sheet_regions('choice', choice_img, classes,
                                               choice_sess.sess, choice_sess.net, conf_thresh, mns_thresh,
                                               coordinate_bias_dict)

            choice_list = choice_list + choice_line_box.choice_line(left, top, choice_img, choice_dict_tf, xml_path)

    return choice_list


def choice_row_col(image, regions, xml_path, conf_thresh, mns_thresh, choice_sess):
    model_type = 'choice_m'
    classes = model_dict[model_type]['classes']
    coordinate_bias_dict = model_dict[model_type]['class_coordinate_bias']

    choice_list = []
    for ele in regions:
        if ele["class_name"] == 'choice':

            choice_box = ele['bounding_box']
            left = choice_box['xmin']
            top = choice_box['ymin']
            choice_img = utils.crop_region(image, choice_box)

            choice_m_dict_tf = resolve_choice. \
                get_single_image_sheet_regions('choice_m', choice_img, classes,
                                               choice_sess.sess, choice_sess.net, conf_thresh, mns_thresh,
                                               coordinate_bias_dict)

            choice_list = choice_list + choice_line_box.choice_line_with_number(left, top, choice_img, choice_m_dict_tf, xml_path)

    return choice_list


def choice_m_row_col(image, regions, xml_path):

    choice_m_dict_tf = [ele for ele in regions if ele['class_name'] == 'choice_m']
    # choice_m_row_col_with_number
    choice_list = []
    try:
        # choice_list = choice_box.get_number_by_enlarge_choice_m(image, choice_m_dict_tf, xml_path)
        # if infer_choice_m_flag:
        #     choice_list = choice_line_box.choice_m_adjust(image, choice_m_dict_tf)
        #
        # else:
        #     choice_list = choice_line_box.choice_m_row_col(image, choice_m_dict_tf, xml_path)  # 找选择题行列、分数

        choice_list = choice_line_box.choice_m_row_col(image, choice_m_dict_tf, xml_path)  # 找选择题行列、分数
        tree = ET.parse(xml_path)  # xml tree
        for index_num, box in enumerate(choice_list):
            if len(box['bounding_box']) > 0:
                abcd = box['bounding_box']
                number = str(box['number'])
                name = '{}_{}*{}_{}_{}'.format('choice_m', box['rows'], box['cols'], box['direction'], number)
                tree = utils.create_xml(name, tree,
                                        abcd['xmin'], abcd['ymin'],
                                        abcd['xmax'], abcd['ymax'])

        tree.write(xml_path)
    except Exception as e:
        traceback.print_exc()
        print(e)

    return choice_list


def exam_number(image, regions, xml_path):
    exam_number_dict = {}
    for ele in regions:
        if ele["class_name"] == 'exam_number':
            exam_number_dict = ele

    exam_number_box = exam_number_dict['bounding_box']
    left = exam_number_box['xmin']
    top = exam_number_box['ymin']
    exam_number_img = utils.crop_region(image, exam_number_box)

    # exam_number_dict = resolve_exam_number_box.exam_number(left, top, exam_number_img, xml_path)
    exam_number_dict = resolve_exam_number_box.exam_number_whole(left, top, exam_number_img, xml_path)

    # print(exam_number_dict)
    return exam_number_dict


def exam_number_row_col(image, regions, xml_path):
    exam_number_dict = {}
    for ele in regions:
        if ele["class_name"] == 'exam_number':
            exam_number_dict = ele

    exam_number_box = exam_number_dict['bounding_box']
    left = exam_number_box['xmin']
    top = exam_number_box['ymin']
    exam_number_img = utils.crop_region(image, exam_number_box)

    exam_number_row_col_dict = exam_number_row_column.get_exam_number_row_and_col(left, top, exam_number_img)

    tree = ET.parse(xml_path)  # xml tree
    if len(exam_number_row_col_dict) > 0:
        exam_number_box = exam_number_row_col_dict['bounding_box']
        name = '{}_{}*{}_{}'.format('exam_number',
                                    exam_number_row_col_dict['rows'],
                                    exam_number_row_col_dict['cols'],
                                    exam_number_row_col_dict['direction'])
        tree = utils.create_xml(name, tree,
                                exam_number_box['xmin'], exam_number_box['ymin'],
                                exam_number_box['xmax'], exam_number_box['ymax'])
        tree.write(xml_path)
        return [exam_number_row_col_dict]

    else:
        tree = utils.create_xml('exam_number', tree,
                                exam_number_box['xmin'], exam_number_box['ymin'],
                                exam_number_box['xmax'], exam_number_box['ymax'])

        tree.write(xml_path)

        return []


def cloze(image, regions, xml_path, conf_thresh, mns_thresh, cloze_sess):
    classes = model_dict['cloze']['classes']
    coordinate_bias_dict = model_dict['cloze']['class_coordinate_bias']

    cloze_list = []
    for ele in regions:
        if ele["class_name"] == 'cloze':
            cloze_box = ele['bounding_box']
            left = cloze_box['xmin']
            top = cloze_box['ymin']
            cloze_img = utils.crop_region(image, cloze_box)
            cloze_dict_tf = resolve_cloze.get_single_image_sheet_regions('cloze', cloze_img, classes,
                                                                         cloze_sess.sess, cloze_sess.net, conf_thresh,
                                                                         mns_thresh, coordinate_bias_dict)
            cloze_list = cloze_list + resolve_cloze_line_box.cloze_line(left, top, cloze_img, cloze_dict_tf['regions'], xml_path)

    return cloze_list


def solve_with_mark(image, regions, xml_path):
    solve_list = []
    mark_list = []
    for ele in regions.copy():
        if 'solve' in ele["class_name"]:
            exam_number_box = ele['bounding_box']
            left = exam_number_box['xmin']
            top = exam_number_box['ymin']
            exam_number_img = utils.crop_region(image, exam_number_box)
            solve_mark_dict = resolve_mark_box.solve_mark(left, top, exam_number_img, xml_path)
            if len(solve_mark_dict) > 0:
                ele['class_name'] = 'solve_'+str(solve_mark_dict['number'])
                solve_list.append(ele)
                mark_list.append(solve_mark_dict)

    return solve_list, mark_list


def solve(image, regions, xml_path):
    solve_list = []
    tree = ET.parse(xml_path)
    for ele in regions.copy():
        if 'solve' in ele["class_name"]:
            exam_number_box = ele['bounding_box']
            exam_number_img = utils.crop_region(image, exam_number_box)
            number = resolve_mark_line_box.solve_line(exam_number_img)
            solve_dict = {'number': number, 'location': exam_number_box, 'default_points': 12}
            solve_list.append(solve_dict)

            tree = utils.create_xml(str(number), tree,
                                    exam_number_box['xmin'], exam_number_box['ymin'],
                                    exam_number_box['xmax'], exam_number_box['ymax'])
    tree.write(xml_path)
    return solve_list


def solve_with_number(regions, xml_path):
    solve_list = []
    for ele in regions:
        if 'solve' in ele["class_name"] or 'composition' in ele["class_name"]:
            solve_dict = {'number': -1, 'default_points': -1}
            ele.update(solve_dict)
            solve_list.append(ele)

    tree = ET.parse(xml_path)  # xml tree
    for index_num, box in enumerate(solve_list):
        if len(box['bounding_box']) > 0:
            abcd = box['bounding_box']
            number = str(box['number'])
            default_points = box["default_points"]
            name = '{}_{}_{}'.format(box["class_name"], number, default_points)
            tree = utils.create_xml(name, tree,
                                    abcd['xmin'], abcd['ymin'],
                                    abcd['xmax'], abcd['ymax'])

    tree.write(xml_path)
    return solve_list


def cloze_with_number(regions, xml_path):
    cloze_list = []
    for ele in regions:
        if 'cloze' == ele["class_name"] or "cloze_s" == ele["class_name"]:
            cloze_dict = {'number': -1, 'default_points': -1}
            ele.update(cloze_dict)
            cloze_list.append(ele)

    tree = ET.parse(xml_path)  # xml tree
    for index_num, box in enumerate(cloze_list):
        if len(box['bounding_box']) > 0:
            abcd = box['bounding_box']
            number = str(box['number'])
            default_points = box["default_points"]
            name = '{}_{}_{}'.format(box["class_name"], number, default_points)
            tree = utils.create_xml(name, tree,
                                    abcd['xmin'], abcd['ymin'],
                                    abcd['xmax'], abcd['ymax'])

    tree.write(xml_path)
    return cloze_list


def make_together(image_path):

    sheet_sess = TfSess('sheet')
    choice_sess = TfSess('choice')
    cloze_sess = TfSess('cloze')

    raw_img = read_single_img(image_path)
    conf_thresh_0 = 0.7
    mns_thresh_0 = 0.3

    series_number = 123456789
    subject = 'english'
    sheets_dict_0, xml_save_path = sheet(series_number, image_path, raw_img, conf_thresh_0, mns_thresh_0, subject, sheet_sess)
    # 手动修改faster_rcnn识别生成的框

    sheets_dict_0 = read_xml_to_json(xml_save_path)
    regions = sheets_dict_0['regions']
    classes_name = str([ele['class_name'] for ele in regions])

    if 'choice' in classes_name:
        try:
            sheets_dict_0['choice'] = choice(raw_img, regions, xml_save_path, conf_thresh_0, mns_thresh_0, choice_sess)
        except Exception:
            traceback.print_exc()

    if 'exam_number' in classes_name:
        try:
            sheets_dict_0['exam_number'] = exam_number(raw_img, regions, xml_save_path)
        except Exception:
            traceback.print_exc()

    if 'cloze' in classes_name:
        try:
            sheets_dict_0['cloze'] = cloze(raw_img, regions, xml_save_path, conf_thresh_0, mns_thresh_0, cloze_sess)
        except Exception:
            traceback.print_exc()

    if 'solve' in classes_name:
        try:
            solve_list, mark_list = solve(raw_img, regions, xml_save_path,)
            sheets_dict_0['solve'] = solve_list
            sheets_dict_0['mark'] = mark_list
        except Exception:
            traceback.print_exc()

    # print(sheets_dict_0)
    return sheets_dict_0


# if __name__ == '__main__':
#     start_time = time.time()
#
#     image_path_0 = os.path.join(r'C:\Users\Administrator\Desktop\sheet\correct\back_sizes\template',
#                                 '20180719004308818_0020.jpg')
#     make_together(image_path_0)
#     end_time = time.time()
#     print('time cost: ', (end_time - start_time))