# @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 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 logger = logging.getLogger(settings.LOGGING_TYPE) sheet_infer_dict = dict(bar_code=True, choice_m=True, exam_number=True, common_sheet=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)) 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: 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, 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.85 or iou[1] > 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) infer_choice_m_flag = True 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']) else: tree = utils.create_xml('exam_number', tree, exam_number_box['xmin'], exam_number_box['ymin'], exam_number_box['xmax'], exam_number_box['ymax']) exam_number_row_col_dict = {} tree.write(xml_path) return [exam_number_row_col_dict] 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))