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