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

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

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_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.analysis.sheet.choice_infer import infer_choice_m
from segment.sheet_resolve.analysis.sheet.ocr_sheet import tell_columns, sheet_sorted
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, adjust_exam_number, exam_number_infer_by_s
from segment.sheet_resolve.analysis.sheet.sheet_infer import infer_exam_number, infer_solve, box_infer_and_complete
from segment.sheet_resolve.analysis.sheet.sheet_infer import exam_number_adjust_infer
from segment.sheet_resolve.tools import utils
from segment.sheet_resolve.tools.tf_settings import xml_template_path, model_dict
from segment.sheet_resolve.tools.utils import create_xml
from segment.sheet_resolve.analysis.sheet.ocr_sheet import ocr2sheet

logger = logging.getLogger(settings.LOGGING_TYPE)

sheet_infer_dict = dict(bar_code=True,
                        choice_m=True,
                        exam_number=True,
                        common_sheet=False,
                        cloze=True,
                        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']

    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]
    infer_box_list = []

    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)
                if len(exam_number_list) > 0:
                    regions.extend(exam_number_list)

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

    if sheet_infer_dict['choice_m']:
        try:
            col_split = col_split_x.copy()
            if not col_split:
                col_split = [w - 1]
            infer_box_list = ocr2sheet(image, col_split_x, ocr)
            choice_m_list = infer_choice_m(image, regions, infer_box_list, col_split)
            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.05:
                            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['cloze']:

        cloze_list = []
        for infer_box in infer_box_list:
            # {'loc': [240, 786, 1569, 1368]}
            loc = infer_box['loc']
            xmin, ymin, xmax, ymax = loc[0], loc[1], loc[2], loc[3]

            for ele in regions:
                if ele['class_name'] in ['cloze_s', 'cloze']:
                    tf_loc = ele['bounding_box']
                    tf_loc_l = tf_loc['xmin']
                    tf_loc_t = tf_loc['ymin']
                    if xmin < tf_loc_l < xmax and ymin < tf_loc_t < ymax:
                        cloze_box = {'class_name': 'cloze',
                                     'bounding_box': {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax}}
                        cloze_list.append(cloze_box)
                        break

        regions.extend(cloze_list)

    if sheet_infer_dict['solve']:
        try:
            include_class = ['info_title',
                             'bar_code',
                             'choice_m',
                             'cloze',
                             'cloze_s',
                             'exam_number',
                             'solve',
                             'solve0',
                             'composition',
                             'composition0',
                             'correction',
                             'alarm_info',
                             'page',
                             'mark'
                             ]
            if 'math' not in subject:
                include_class.remove('cloze_s')
            regions_subset = [ele for ele in regions if ele['class_name'] in include_class]

            col_split = col_split_x.copy()
            if not col_split:
                col_split = [w - 1]
            col_regions = sheet_sorted(regions_subset, col_split.copy())

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

            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, h, w, col_regions, col_split_x.copy())
            regions.extend(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, 'col_split': col_split_x})

    #  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'

    # 没有adjust 的regions
    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)

    for i, ele in enumerate(sheets_dict['col_split']):
        name = 'line_{}'.format(i)
        xmin = ele
        ymin = 1
        xmax = ele + 2
        ymax = img_shape[0]
        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, subject, xml_path):
    choice_m_dict_tf = []
    direction_list = []

    for ele in regions:
        if ele['class_name'] == 'choice_m':
            choice_m_dict_tf.append(ele)
        if ele['class_name'] == 'choice_n':
            loc = ele['bounding_box']
            xmin, ymin, xmax, ymax = loc['xmin'], loc['ymin'], loc['xmax'], loc['ymax']
            if ymax - ymin > 2 * (xmax - xmin):
                direction = 180
            else:
                direction = 90
            direction_list.append(direction)
        if ele['class_name'] == 'choice_s':
            loc = ele['bounding_box']
            xmin, ymin, xmax, ymax = loc['xmin'], loc['ymin'], loc['xmax'], loc['ymax']
            if ymax - ymin > 2 * (xmax - xmin):
                direction = 90
            else:
                direction = 180
            direction_list.append(direction)

    c180 = direction_list.count(180)
    c90 = direction_list.count(90)

    if c180 == c90 == 0:
        direction = 0
    else:
        direction = 180 if c180 >= c90 else 90

    # 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, direction, subject,
                                                       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"] or 'correction' in ele["class_name"]:
            solve_dict = {'number': -1, 'default_points': -1, 'span': False, 'span_id': 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