1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273 |
- # @Author : lightXu
- # @File : sheet_infer.py
- # @Time : 2019/9/26 0026 上午 10:18
- import itertools
- import os
- import re
- import traceback
- import xml.etree.cElementTree as ET
- from itertools import combinations
- import cv2
- import numpy as np
- from shapely.geometry import LineString, Polygon
- from segment.sheet_resolve.tools.utils import create_xml, crop_region_direct, crop_region, image_hash_detection_simple
- from segment.sheet_resolve.tools.brain_api import get_ocr_text_and_coordinate
- ASPECT_FLAG = 4.0
- REMAIN_RATIO = 0.1
- PIX_VALUE_LOW = 15.0
- PIX_VALUE_HIGH = 245
- TYPE_SCORE_MNS = 0.5
- def _get_char_near_img(char_location, near):
- left = char_location['left']
- top = char_location['top']
- width = char_location['width']
- height = char_location['height']
- next_location = char_location
- if near == 'left':
- next_location = {'left': int(left - 1.5 * width), 'top': top, 'width': width, 'height': height}
- if near == 'right':
- next_location = {'left': int(left + 1.5 * width), 'top': top, 'width': width, 'height': height}
- if near == 'up':
- next_location = {'left': left, 'top': int(top - 1.5 * height), 'width': width, 'height': height}
- if near == 'down':
- next_location = {'left': left, 'top': int(top + 1.5 * height), 'width': width, 'height': height}
- return next_location
- def _get_board(image, location, direction):
- std = 0
- next_location = location
- while std < 10:
- next_location = _get_char_near_img(next_location, direction)
- box = (next_location['left'], next_location['top'],
- next_location['left'] + next_location['width'],
- next_location['top'] + next_location['height'],)
- region = crop_region_direct(image, box)
- std = np.var(region)
- return next_location
- def infer_bar_code(image, ocr_dict_list, attention_region):
- attention_polygon_list = []
- for attention in attention_region:
- coordinates = attention['bounding_box']
- xmin = coordinates['xmin']
- ymin = coordinates['ymin']
- xmax = coordinates['xmax']
- ymax = coordinates['ymax']
- attention_polygon = Polygon([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
- attention_polygon_list.append(attention_polygon)
- img_cols, img_rows = image.shape[0], image.shape[1]
- pattern = r'条形码|条码|条形|形码'
- bar_code_dict_list = []
- for index, ele in enumerate(ocr_dict_list):
- words = ele['words'].replace(' ', '')
- chars_list = ele['chars']
- length = len(chars_list)
- match_list = [(m.group(), m.span()) for m in re.finditer(pattern, words) if m]
- if match_list: # 不为空
- for match in match_list:
- start_index = match[1][0]
- end_index = match[1][1] - 1
- for i in range(start_index - 1, -1, -1):
- xmin_start = chars_list[start_index]['location']['left']
- start_tmp = chars_list[i]['location']['left'] + 2 * chars_list[i]['location']['width']
- if xmin_start <= start_tmp:
- start_index = i
- for i in range(end_index, length):
- xmax_end = chars_list[end_index]['location']['left'] + 2 * chars_list[i]['location']['width']
- end_tmp = chars_list[i]['location']['left']
- if xmax_end >= end_tmp:
- end_index = i
- bar_code_char_xmin = chars_list[start_index]['location']["left"]
- bar_code_char_xmax = chars_list[end_index]['location']["left"]+chars_list[end_index]['location']["width"]
- bar_code_char_ymin = chars_list[start_index]['location']["top"]
- bar_code_char_ymax = chars_list[end_index]['location']["top"]+chars_list[end_index]['location']["height"]
- bar_code_char_polygon = Polygon([(bar_code_char_xmin, bar_code_char_ymin),
- (bar_code_char_xmax, bar_code_char_ymin),
- (bar_code_char_xmax, bar_code_char_ymax),
- (bar_code_char_xmin, bar_code_char_ymax)])
- contain_cond = [False]*len(attention_polygon_list)
- for i, attention_ele in enumerate(attention_polygon_list):
- if attention_ele.contains(bar_code_char_polygon):
- contain_cond[i] = True
- if True not in contain_cond: # 条形码文字不在attention里面
- left_board_location = _get_board(image, chars_list[start_index]['location'], 'left')
- right_board_location = _get_board(image, chars_list[end_index]['location'], 'right')
- up_board_location = _get_board(image, chars_list[start_index]['location'], 'up')
- down_board_location = _get_board(image, chars_list[end_index]['location'], 'down')
- xmin = left_board_location['left']
- ymin = up_board_location['top']
- xmax = right_board_location['left'] + right_board_location['width']
- ymax = down_board_location['top'] + down_board_location['height']
- bias = 5
- xmin = max(1, int(xmin)-bias)
- ymin = max(1, int(ymin)-bias)
- xmax = min(int(xmax)+bias, img_rows - 1)
- ymax = min(int(ymax)+bias, img_cols - 1)
- bar_code_dict = {'class_name': 'bar_code',
- 'bounding_box': {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax}}
- bar_code_dict_list.append(bar_code_dict)
- # print(bar_code_dict)
- break # 默认只有一个条形码
- else:
- continue
- # 过滤attention 区域存在条形码的文字
- for bar_code in bar_code_dict_list.copy():
- coordinates = bar_code['bounding_box']
- xmin = coordinates['xmin']
- ymin = coordinates['ymin']
- xmax = coordinates['xmax']
- ymax = coordinates['ymax']
- bar_code_polygon = Polygon([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
- for attention_polygon in attention_polygon_list:
- cond1 = bar_code_polygon.within(attention_polygon) or bar_code_polygon.contains(attention_polygon)
- cond2 = False
- cond3 = bar_code_polygon.overlaps(attention_polygon)
- if cond3:
- intersection_poly = bar_code_polygon.intersection(attention_polygon)
- cond2 = intersection_poly.area / bar_code_polygon.area >= 0.01
- cond3 = intersection_poly.area / attention_polygon.area >= 0.01
- if cond1 or cond2 or cond3:
- bar_code_dict_list.remove(bar_code)
- break
- return bar_code_dict_list
- def infer_exam_number(image, ocr_dict_list, existed_regions, times_threshold=5):
- # existed_polygon_list = []
- # for region in existed_regions:
- # coordinates = region['bounding_box']
- # xmin = coordinates['xmin']
- # ymin = coordinates['ymin']
- # xmax = coordinates['xmax']
- # ymax = coordinates['ymax']
- # existed_polygon = Polygon([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
- # existed_polygon_list.append(existed_polygon)
- img_rows, img_cols = image.shape[0], image.shape[1]
- exam_number_dict_list = []
- xmin, ymin, xmax, ymax = 9999, 9999, 0, 0
- pattern = r'[0oO]|[2-9]' # 除去1,避免[]被识别为1
- exclude = r'分|题|[ABD]'
- key_digital = []
- all_height = []
- cols = []
- for index, ele in enumerate(ocr_dict_list):
- words = ele['words'].replace(' ', '')
- match_list = [(m.group(), m.span()) for m in re.finditer(pattern, words) if m]
- exclude_list = [(m.group(), m.span()) for m in re.finditer(exclude, words, re.I) if m]
- match_digital_arr = np.asarray([int(char[0].replace('o', '0').replace('O', '0')) for char in match_list])
- if len(match_digital_arr) > 0:
- counts = np.bincount(match_digital_arr)
- mode_times = np.max(counts)
- if mode_times >= times_threshold and len(exclude_list) < 1:
- mode_value = np.argmax(counts) # 众数,避免考号末尾出现的其他数字
- key_index = np.where(match_digital_arr == mode_value)[0]
- cols.append(len(key_index))
- start_index = match_list[key_index[0]][1][0]
- end_index = match_list[key_index[-1]][1][0]
- xmin_t = ele['chars'][start_index]['location']['left']
- ymin_t = ele['chars'][start_index]['location']['top']
- xmax_t = ele['chars'][end_index]['location']['left'] + ele['chars'][end_index]['location']['width']
- ymax_t = ele['chars'][end_index]['location']['top'] + ele['chars'][end_index]['location']['height']
- mean_width = sum([int(ele['chars'][match_list[i][1][0]]['location']['width'])
- for i in key_index]) // len(key_index)
- mean_height = sum([int(ele['chars'][match_list[i][1][0]]['location']['height'])
- for i in key_index]) // len(key_index)
- all_height.append(mean_height)
- xmin = min(xmin, xmin_t-mean_width)
- ymin = min(ymin, ymin_t)
- xmax = max(xmax, xmax_t+mean_width)
- ymax = max(ymax, ymax_t)
- xmin = int(xmin) if xmin >= 1 else 1
- ymin = int(ymin) if ymin >= 1 else 1
- xmax = int(xmax) if xmax <= img_cols - 1 else img_cols - 1
- ymax = int(ymax) if ymax <= img_rows - 1 else img_rows - 1
- key_digital.append(mode_value)
- if 9 in key_digital:
- break
- if 0 in key_digital and 9 in key_digital and len(key_digital) > 4:
- mean_height = sum(all_height)//10
- exam_number_dict = {'class_name': 'exam_number',
- 'bounding_box': {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax+mean_height},
- 'rows': 10,
- 'cols': max(cols)
- }
- exam_number_dict_list.append(exam_number_dict)
- return exam_number_dict_list
- else:
- if len(key_digital) > 1:
- dgt_min = min(key_digital)
- dgt_max = max(key_digital)
- mean_height = sum(all_height)//len(all_height)
- dif = dgt_max - dgt_min
- blank_height = ymax - ymin - mean_height * (dif+1)
- mean_blank = blank_height // dif
- upper_height = dgt_min * (mean_blank + mean_height) + mean_blank//2
- downward_height = (9-dgt_max) * (mean_blank + mean_height) + mean_blank
- exam_number_dict = {'class_name': 'exam_number',
- 'bounding_box': {'xmin': xmin, 'ymin': ymin-upper_height,
- 'xmax': xmax, 'ymax': ymax+downward_height},
- 'rows': 10,
- 'cols': max(cols)}
- exam_number_dict_list.append(exam_number_dict)
- if len(key_digital) == 1:
- dgt_min = dgt_max = min(key_digital)
- eval_height = sum(all_height)//len(all_height) * 1.5
- upper_height = dgt_min * eval_height
- downward_height = (9-dgt_max) * eval_height
- exam_number_dict = {'class_name': 'exam_number',
- 'bounding_box': {'xmin': xmin, 'ymin': ymin-upper_height,
- 'xmax': xmax, 'ymax': ymax+downward_height},
- 'rows': 10,
- 'cols': max(cols)}
- exam_number_dict_list.append(exam_number_dict)
- iou_cond = True
- exam_number_dict_list_check = []
- if len(exam_number_dict_list) > 0:
- for exam_number_dict in exam_number_dict_list:
- exam_number_polygon = Polygon([(exam_number_dict["xmin"], exam_number_dict["ymin"]),
- (exam_number_dict["xmax"], exam_number_dict["ymin"]),
- (exam_number_dict["xmax"], exam_number_dict["ymax"]),
- (exam_number_dict["xmin"], exam_number_dict["ymax"])])
- for region in existed_regions:
- class_name = region["class_name"]
- if class_name in ["attention", "solve", "choice", "choice_m", 'choice_s', "cloze", 'cloze_s',
- 'bar_code', 'qr_code', 'composition', 'solve0']:
- coordinates = region['bounding_box']
- xmin = coordinates['xmin']
- ymin = coordinates['ymin']
- xmax = coordinates['xmax']
- ymax = coordinates['ymax']
- existed_polygon = Polygon([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
- overlab_area = existed_polygon.intersection(exam_number_polygon).area
- iou = overlab_area / (exam_number_polygon.area + existed_polygon.area - overlab_area)
- if iou > 0:
- iou_cond = False
- break
- if iou_cond:
- exam_number_dict_list_check.append(exam_number_polygon)
- return exam_number_dict_list_check
- def adjust_exam_number(regions):
- exam_number_w_regions = list()
- exam_number_regions = list()
- for i in range(len(regions) - 1, -1, -1):
- region = regions[i]
- if region['class_name'] == 'exam_number_w':
- exam_number_w_regions.append(region)
- if region['class_name'] == 'exam_number':
- exam_number_regions.append(region)
- regions.pop(i)
- exam_number_region = exam_number_regions[0]
- if len(exam_number_regions) > 1:
- exam_number_regions = sorted(exam_number_regions, key=lambda x: x['bounding_box']['ymin'])
- exam_number_region = exam_number_regions[0]
- exam_number_w_index = 0
- if len(exam_number_w_regions) > 1:
- distance = [abs(int(ele['bounding_box']['ymax']) - int(exam_number_region['bounding_box']['ymin']))
- for ele in exam_number_w_regions]
- exam_number_w_index = distance.index(min(distance))
- exam_number_w_region = exam_number_w_regions[exam_number_w_index]
- standard = exam_number_w_region['bounding_box']
- exam_number_region['bounding_box'].update({'xmin': standard['xmin'], 'xmax': standard['xmax']})
- regions.append(exam_number_region)
- return regions
- def exam_number_infer_by_s(image, regions):
- exam_number_s_list = [ele for ele in regions if ele['class_name'] == 'exam_number_s'
- and (int(ele['bounding_box']['xmax'])-int(ele['bounding_box']['xmin']) <
- int(ele['bounding_box']['ymax'])-int(ele['bounding_box']['ymin']))]
- # 找边界
- exam_number_s_list = sorted(exam_number_s_list, key=lambda x: x['bounding_box']['xmin'])
- left_limit = exam_number_s_list[0]['bounding_box']['xmin']
- right_limit = exam_number_s_list[-1]['bounding_box']['xmax']
- left_image = crop_region(image, exam_number_s_list[0]['bounding_box'])
- right_image = crop_region(image, exam_number_s_list[-1]['bounding_box'])
- mean_width = sum([int(ele['bounding_box']['xmax'])-int(ele['bounding_box']['xmin'])
- for ele in exam_number_s_list]) // len(exam_number_s_list)
- top_limit = min([ele['bounding_box']['ymin'] for ele in exam_number_s_list])
- bottom_limit = max([ele['bounding_box']['ymax'] for ele in exam_number_s_list])
- left_infer = True
- while left_infer:
- infer_box_xmin = int(left_limit - 1.5*mean_width)
- infer_box_xmax = int(left_limit - 0.5*mean_width)
- infer_box_ymin = int(exam_number_s_list[0]['bounding_box']['ymin'])
- infer_box_ymax = int(exam_number_s_list[0]['bounding_box']['ymax'])
- infer_image = crop_region_direct(image, [infer_box_xmin, infer_box_ymin, infer_box_xmax, infer_box_ymax])
- simi = image_hash_detection_simple(left_image, infer_image)
- print('l:', simi)
- if simi >= 0.85:
- left_limit = infer_box_xmin
- else:
- left_infer = False
- right_infer = True
- while right_infer:
- infer_box_xmin = int(right_limit + 0.5 * mean_width)
- infer_box_xmax = int(right_limit + 1.5 * mean_width)
- infer_box_ymin = int(exam_number_s_list[-1]['bounding_box']['ymin'])
- infer_box_ymax = int(exam_number_s_list[-1]['bounding_box']['ymax'])
- infer_image = crop_region_direct(image, [infer_box_xmin, infer_box_ymin, infer_box_xmax, infer_box_ymax])
- simi = image_hash_detection_simple(right_image, infer_image)
- print('r:', simi)
- if simi >= 0.70:
- right_limit = infer_box_xmax
- else:
- right_infer = False
- infer_exam_number_region = {'xmin': left_limit, 'xmax': right_limit, 'ymin': top_limit, 'ymax': bottom_limit, }
- exam_dict_list = [{'class_name': 'exam_number', 'bounding_box': infer_exam_number_region}]
- # print(exam_dict_list)
- return exam_dict_list
- def draw_blank(image, infer_box):
- x_bias, y_bias = infer_box[:2]
- exam_number_infer_region = crop_region_direct(image, infer_box)
- ocr = get_ocr_text_and_coordinate(exam_number_infer_region)
- pattern = r'[\u4e00-\u9fa5]'
- for index, ele in enumerate(ocr):
- # words = ele['words'].replace(' ', '')
- for char in ele['chars']:
- if re.match(pattern, char['char']):
- loc = char['location']
- left, top, width, height = loc['left']+x_bias, loc['top']+y_bias, loc['width'], loc['height']
- image[top:top+2*height, left:left+width, :] = 255
- return image
- def exam_number_adjust_infer(image, regions):
- box = []
- box1 = []
- for i in range(len(regions) - 1, -1, -1):
- region = regions[i]
- if region['class_name'] == 'exam_number_s':
- loc = region['bounding_box']
- if loc['ymax'] - loc['ymin'] > loc['xmax'] - loc['xmin']:
- box.append([loc['xmin'], loc['ymin'], loc['xmax'], loc['ymax']])
- if region['class_name'] == 'exam_number':
- loc = region['bounding_box']
- box.append([loc['xmin'], loc['ymin'], loc['xmax'], loc['ymax']])
- regions.pop(i)
- if region['class_name'] == 'exam_number_w':
- loc = region['bounding_box']
- if loc['xmax'] - loc['xmin'] > loc['ymax'] - loc['ymin']:
- box1.append([loc['xmin'], loc['ymax'], loc['xmax'], loc['ymax']]) # exam_number_w 只取下界
- if not box:
- return image, regions
- box_array = np.array(box+box1)
- xmin, ymin = np.min(box_array, axis=0)[:2]
- xmax, ymax = np.max(box_array, axis=0)[2:]
- images = draw_blank(image, (xmin, ymin, xmax, ymax))
- exam_number_infer = {'class_name': 'exam_number',
- 'bounding_box': {'xmin': xmin, 'ymin': ymin,
- 'xmax': xmax, 'ymax': ymax}}
- regions.append(exam_number_infer)
- return images, regions
- def gen_xml_new(path, ocr_list):
- tree = ET.parse(r'../../tools/000000-template.xml') # xml tree
- for index, ele in enumerate(ocr_list):
- words = ele['words']
- location = ele['location']
- xmin = location['xmin']
- ymin = location['ymin']
- xmax = location['xmax']
- ymax = location['ymax']
- tree = create_xml('{}'.format(words), tree, str(xmin), str(ymin), str(xmax), str(ymax))
- # print(exam_items_bbox)
- tree.write(path.replace('.jpg', '.xml'))
- def subfield_answer_sheet(img0, answer_sheet):
- h, w = img0.shape[:2]
- one_part = 0
- line_xmax_1 = 0
- line_xmax_2 = 0
- modules = []
- modules11 = []
- w_int_1 = w
- w_int_2 = round(w / 2)
- w_int_3 = round(w / 3)
- w_int_4 = round(w / 4)
- w_int_8 = round(w / 8)
- if w_int_8 < 50:
- w_int_8 = 50
- key_modules_classes = ['choice', 'cloze', 'solve', 'solve0', 'composition0', 'composition', 'correction',
- 'ban_area', ]
- if h > w: # 暂定答题卡高大于宽的为单栏
- one_part = 1
- else:
- temp1 = 0
- temp2 = 0
- for ele in answer_sheet:
- if ele["class_name"] in key_modules_classes:
- modules.append(ele)
- modules_xmin = sorted(modules, key=lambda x: (x['bounding_box']['xmin']))
- modules_xmax = sorted(modules, key=lambda x: (x['bounding_box']['xmax']))
- for i in range(len(modules_xmin) - 1):
- if i == 0 and modules_xmin[0]['bounding_box']['xmin'] - 0 > w_int_4:
- temp1 = 1
- else:
- if modules_xmin[i + 1]['bounding_box']['xmin'] - modules_xmin[i]['bounding_box']['xmax'] > w_int_4:
- if modules11 == []:
- line_xmax_1 = modules_xmin[i]['bounding_box']['xmax'] + 20
- line_xmax_2 = modules_xmin[i + 1]['bounding_box']['xmin'] - 20
- else:
- modules11.append(modules_xmin[i]['bounding_box']['xmax'])
- modules11_xmax = sorted(modules11)[-1]
- line_xmax_1 = modules11_xmax + 20
- line_xmax_2 = modules_xmin[i + 1]['bounding_box']['xmin'] - 20
- modules11 = []
- temp1 = 1
- temp2 = 1
- break
- elif modules_xmin[i + 1]['bounding_box']['xmin'] - modules_xmin[i]['bounding_box']['xmax'] > -w_int_8:
- if temp1 == 0:
- if modules11 == []:
- line_xmax_1 = int((modules_xmin[i + 1]['bounding_box']['xmin'] +
- modules_xmin[i]['bounding_box']['xmax']) / 2)
- else:
- modules11.append(modules_xmin[i]['bounding_box']['xmax'])
- modules11_xmax = sorted(modules11)[-1]
- line_xmax_1 = int((modules_xmin[i + 1]['bounding_box']['xmin'] +
- modules11_xmax) / 2)
- modules11 = []
- temp1 = 1
- elif temp1 == 1:
- if modules11 == []:
- line_xmax_2 = int((modules_xmin[i + 1]['bounding_box']['xmin'] +
- modules_xmin[i]['bounding_box']['xmax']) / 2)
- else:
- modules11.append(modules_xmin[i]['bounding_box']['xmax'])
- modules11_xmax = sorted(modules11)[-1]
- line_xmax_2 = int((modules_xmin[i + 1]['bounding_box']['xmin'] +
- modules11_xmax) / 2)
- temp2 = 1
- else:
- modules11.append(modules_xmin[i]['bounding_box']['xmax'])
- if temp1 == 0 and temp2 == 0:
- if modules_xmax[-1]['bounding_box']['xmax'] - w < -(2 * w_int_4):
- line_xmax_1 = modules_xmax[-1]['bounding_box']['xmax'] + 20
- line_xmax_2 = 2 * w_int_3
- elif modules_xmax[-1]['bounding_box']['xmax'] - w < -w_int_4:
- line_xmax_1 = modules_xmax[-1]['bounding_box']['xmax'] + 20
- elif temp1 == 1 and temp2 == 0:
- if modules_xmax[-1]['bounding_box']['xmax'] - w < -w_int_4:
- line_xmax_2 = 2 * w_int_3
- return line_xmax_1, line_xmax_2
- def get_intersection_point(lines, orthogonal_lines, border):
- intersect_point_list = []
- for line in lines:
- width_min, height_min, width_max, height_max = border
- (x_l, y_u), (x_r, y_d) = line.coords
- x_l = x_l if x_l > width_min else width_min + 1 # 避免边界
- x_r = x_r if x_r < width_max else width_max - 1
- y_u = y_u if y_u > height_min else height_min + 1
- y_d = y_d if y_d < height_max else height_max - 1
- points_list = []
- if x_l == x_r:
- line_direction = 'lon'
- raw_line = LineString([(x_l, y_u), (x_r, y_d)])
- extend_line = LineString([(x_l, height_min), (x_r, height_max)])
- points_list.extend([height_min + 1, height_max - 1]) # 延长线与边界交点,并避免key_point位于现有边界上
- line_start, line_end = y_u, y_d
- else:
- line_direction = 'lat'
- raw_line = LineString([(x_l, y_u), (x_r, y_d)])
- extend_line = LineString([(width_min, y_u), (width_max, y_d)])
- points_list.extend([width_min + 1, width_max - 1]) # 延长线与边界交点,并避免key_point位于现有边界上
- line_start, line_end = x_l, x_r
- for ele in orthogonal_lines:
- cond1 = extend_line.intersects(ele) # T, L, 十交叉
- cond2 = extend_line.crosses(ele) # 十字交叉
- cond3 = raw_line.intersects(ele)
- cond4 = raw_line.crosses(ele)
- if line_direction == 'lat':
- if cond3:
- (xp, yp) = raw_line.intersection(ele).bounds[:2]
- intersect_point_list.append((xp, yp))
- elif cond1:
- (xp, yp) = extend_line.intersection(ele).bounds[:2]
- points_list.append(xp)
- if line_direction == 'lon':
- if cond3:
- (xp, yp) = raw_line.intersection(ele).bounds[:2]
- intersect_point_list.append((xp, yp))
- elif cond1:
- (xp, yp) = extend_line.intersection(ele).bounds[:2]
- points_list.append(yp)
- points_array = np.asarray(points_list, dtype=np.uint)
- left_key = np.max(points_array[points_array <= line_start])
- right_key = np.min(points_array[points_array >= line_end]) # 延长线两边延长并取得第一个交点
- if line_direction == 'lat':
- intersect_point = [(left_key, y_u), (right_key, y_d)]
- else:
- intersect_point = [(x_l, left_key), (x_r, right_key)]
- # print(intersect_point)
- intersect_point_list.extend(intersect_point)
- return intersect_point_list
- def infer_sheet_box(image, sheet_dict, lon_split_line, exclude_classes):
- height_max, width_max = image.shape[0], image.shape[1]
- height_min, width_min = 0, 0
- latitude = []
- longitude = []
- lines = []
- sheet_polygons = []
- all_sheet_polygons = []
- choice_polygon = []
- # exclude_classes = ['cloze_s', 'exam_number_s', 'choice_s', 'type_score',
- # 'mark', 'page', 'exam_number_s', 'cloze_score', 'name_w',
- # 'class_w',]
- h_min = []
- h_max = []
- for index, region_box in enumerate(sheet_dict):
- coordinates = region_box['bounding_box']
- xmin = coordinates['xmin']
- ymin = coordinates['ymin']
- xmax = coordinates['xmax']
- ymax = coordinates['ymax']
- if region_box['class_name'] == 'info_title': # 上限
- h_min.append(ymin)
- if region_box['class_name'] == 'page': # 下限
- h_max.append(ymin)
- if region_box['class_name'] == 'alarm_info':
- h_min.append(ymin)
- h_max.append(ymin)
- if h_min:
- hgt_min = min(h_min)
- if hgt_min < height_max / 4:
- height_min = hgt_min
- if h_max:
- hgt_max = max(h_max)
- if hgt_max > 3 * height_max / 4:
- height_max = hgt_max
- # height_min = h_min if h_min != 9999 else height_min
- # height_max = h_max if h_max != 0 else height_max
- for index, region_box in enumerate(sheet_dict):
- coordinates = region_box['bounding_box']
- xmin = coordinates['xmin']
- ymin = coordinates['ymin']
- xmax = coordinates['xmax']
- ymax = coordinates['ymax']
- box_polygon = Polygon([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
- if region_box['class_name'] not in exclude_classes:
- if region_box['class_name'] not in ['choice', 'cloze']: # 推断选择题区域内的choice_m
- sheet_polygons.append(box_polygon)
- if region_box['class_name'] == 'choice':
- choice_polygon.append(box_polygon)
- all_sheet_polygons.append(box_polygon)
- line1 = LineString([(xmin, ymin), (xmin, ymax)])
- line2 = LineString([(xmax, ymin), (xmax, ymax)])
- line3 = LineString([(xmin, ymin), (xmax, ymin)])
- line4 = LineString([(xmin, ymax), (xmax, ymax)])
- lines.extend([line1, line2, line3, line4])
- longitude.extend([line1, line2])
- latitude.extend([line3, line4])
- # sheet_polygons 去除包裹的情况
- sheet_polygons_ = list(combinations(sheet_polygons, 2))
- for polygons in sheet_polygons_:
- if polygons[0].within(polygons[1]) or polygons[0].contains(polygons[1]):
- area_list = [polygons[0].area, polygons[1].area]
- min_polygon = polygons[area_list.index(min(area_list))]
- if min_polygon in sheet_polygons:
- sheet_polygons.remove(min_polygon)
- min_polygon = sorted(all_sheet_polygons, key=lambda p: p.area)[0]
- avg_area = sum([polygon.area for polygon in sheet_polygons]) / len(sheet_polygons)
- # 所有矩形框的延长线与矩形框集图像边界的交点
- latitude = sorted(latitude, key=lambda x: x.bounds[1]) # y
- longitude = sorted(longitude, key=lambda x: x.bounds[0]) # x
- lat_intersect_point_list = get_intersection_point(latitude, longitude,
- (width_min, height_min, width_max, height_max))
- lon_intersect_point_list = get_intersection_point(longitude, latitude,
- (width_min, height_min, width_max, height_max))
- raw_corner = [(width_min + 1, height_min + 1), (width_min + 1, height_max - 1), (width_max - 1, 1),
- (width_max - 1, height_max - 1)]
- # raw_corner = []
- intersect_point_list = lat_intersect_point_list + lon_intersect_point_list + raw_corner
- intersect_point_list = list(set(intersect_point_list))
- intersect_point_dict = {k: index + 1 for index, k in enumerate(intersect_point_list)}
- def _filter_rect(p_list):
- flag = 0
- for ele in p_list:
- try:
- flag = intersect_point_dict[ele]
- except KeyError:
- flag = 0
- break
- if flag > 0:
- x_c = sum([ele[0] for ele in p_list]) / 4
- y_c = sum([ele[1] for ele in p_list]) / 4
- d1, d2, d3, d4 = [LineString([p, (x_c, y_c)]).length for p in p_list]
- return (0 not in [d1, d2, d3, d4]) and d1 == d2 and d1 == d3 and d1 == d4
- else:
- return False
- def _find_rect(point):
- (x1, y1) = point[0]
- (x2, y2) = point[1]
- if x1 != x2 and y1 != y2:
- xmin, ymin = min(x1, x2), min(y1, y2)
- xmax, ymax = max(x1, x2), max(y1, y2)
- points_4 = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
- w, h = xmax - xmin, ymax - ymin
- aspect_flag_extreme = max(w / h, h / w) < 1.5 * ASPECT_FLAG # 解决极端情况
- rect_flag = _filter_rect(points_4)
- if aspect_flag_extreme and rect_flag:
- gen_polygon = Polygon([(points_4[0]), (points_4[1]), (points_4[2]), (points_4[3])])
- flags = set()
- for polygon in sheet_polygons:
- decision = [gen_polygon.contains(polygon),
- gen_polygon.within(polygon),
- gen_polygon.overlaps(polygon)]
- if True in decision: # 边界问题
- flags.add(False)
- break
- else:
- flags.add(True)
- if False in flags:
- pass
- else:
- return gen_polygon
- def _filter_none(p):
- if p is not None:
- return True
- points_2 = combinations(intersect_point_list, 2)
- gen_polygon_list = map(_find_rect, points_2)
- gen_polygon_list = list(filter(_filter_none, gen_polygon_list))
- gen_polygon_list = sorted(gen_polygon_list, key=lambda p: p.area, reverse=True)
- # gen_polygon_list = [polygon for index, polygon in enumerate(gen_polygon_list) if index % 2 == 0]
- it = itertools.groupby(gen_polygon_list)
- gen_polygon_list = [k for k, g in it]
- # 在选择题区域的infer polygon
- # gen_choice = []
- # for ele in gen_polygon_list:
- # for choice_p in choice_polygon:
- # if ele.within(choice_p):
- # gen_choice.append(ele)
- sheet_box_area = sum([polygon.area for polygon in sheet_polygons])
- image_area = width_max * height_max
- blank_ratio = 1 - sheet_box_area / image_area
- polygon_index = 0
- include_polygon = []
- while blank_ratio > REMAIN_RATIO and polygon_index < len(gen_polygon_list):
- polygon = gen_polygon_list[polygon_index]
- blank_ratio = blank_ratio - polygon.area / image_area
- include_polygon.append(polygon)
- polygon_index += 1
- # gen_polygon_list = [polygon for index, polygon in enumerate(gen_polygon_list)
- # if polygon.area > 1.5 * min_polygon.area]
- for polygon in gen_polygon_list.copy():
- xi, yi, xx, yx = polygon.bounds
- w, h = xx - xi, yx - yi
- if polygon.area <= 1.5 * min_polygon.area or h / w > 2 and polygon.area < avg_area:
- gen_polygon_list.remove(polygon)
- polygon_2 = list(combinations(gen_polygon_list, 2))
- for polygons in polygon_2:
- try:
- cond2 = polygons[0].overlaps(polygons[1]) # 叠置关系二次分段
- if cond2:
- area_list = [polygons[0].area, polygons[1].area]
- min_index = area_list.index(min(area_list))
- smaller_polygon = polygons[min_index]
- larger_polygon = polygons[1 - min_index]
- new_polygon = smaller_polygon.difference(larger_polygon)
- if smaller_polygon in gen_polygon_list:
- gen_polygon_list.remove(smaller_polygon)
- if 'MultiPolygon' in str(type(new_polygon)):
- for ele in new_polygon:
- xm, ym, xx, yx = ele.bounds
- w, h = xx - xm, yx - ym
- if max(w / h, h / w) < 1.5 * ASPECT_FLAG and ele.area > 1.5 * min_polygon.area:
- gen_polygon_list.append(ele)
- elif len(set(new_polygon.exterior.coords)) == 4:
- xm, ym, xx, yx = new_polygon.bounds
- w, h = xx - xm, yx - ym
- if max(w / h, h / w) < 1.5 * ASPECT_FLAG and new_polygon.area > 1.5 * min_polygon.area:
- gen_polygon_list.append(new_polygon)
- except Exception as polygon_e:
- print(polygon_e)
- continue
- polygon_2 = list(combinations(gen_polygon_list, 2)) # 包含关系取大值
- for polygons in polygon_2:
- cond1 = polygons[0].equals(polygons[1])
- if cond1 and polygons[1] in gen_polygon_list:
- gen_polygon_list.remove(polygons[1])
- polygon_2 = list(combinations(gen_polygon_list, 2))
- for polygons in polygon_2:
- cond2 = polygons[0].contains(polygons[1]) or polygons[0].within(polygons[1])
- if cond2:
- area_list = [polygons[0].area, polygons[1].area]
- min_index = area_list.index(min(area_list))
- smaller_polygon = polygons[min_index]
- larger_polygon = polygons[1 - min_index]
- sxi, syi, sxx, syx = smaller_polygon.bounds
- bxi, byi, bxx, byx = larger_polygon.bounds
- # inner_touch_cond = '212F11FF2' == larger_polygon.relate(smaller_polygon)
- two_side_touch_cond = (sxi == bxi and sxx == bxx) or (syi == byi and syx == byx)
- if two_side_touch_cond:
- dif_polygon = larger_polygon.difference(smaller_polygon)
- if larger_polygon in gen_polygon_list:
- gen_polygon_list.remove(larger_polygon)
- if 'MultiPolygon' in str(type(dif_polygon)):
- for ele in dif_polygon:
- xm, ym, xx, yx = ele.bounds
- w, h = xx - xm, yx - ym
- if max(w / h, h / w) < 1.5 * ASPECT_FLAG and ele.area > 1.5 * min_polygon.area:
- gen_polygon_list.append(ele)
- elif len(set(dif_polygon.exterior.coords)) == 4: # empty
- xm, ym, xx, yx = dif_polygon.bounds
- w, h = xx - xm, yx - ym
- if max(w / h, h / w) < 1.5 * ASPECT_FLAG and dif_polygon.area > 1.5 * min_polygon.area:
- gen_polygon_list.append(dif_polygon)
- else:
- if smaller_polygon in gen_polygon_list:
- gen_polygon_list.remove(smaller_polygon)
- polygon_2 = list(combinations(gen_polygon_list, 2)) # 包含关系取大值
- for polygons in polygon_2:
- cond1 = polygons[0].equals(polygons[1])
- if cond1 and polygons[1] in gen_polygon_list:
- gen_polygon_list.remove(polygons[1])
- if len(lon_split_line) > 0:
- for line in lon_split_line:
- # line = LineString([(286, 1), (286, 599)])
- for poly in gen_polygon_list.copy():
- cond1 = line.intersects(poly)
- cond2 = line.touches(poly)
- if cond1 and not cond2:
- dif_polygons = poly.difference(line)
- corner_list = list(set(dif_polygons.exterior.coords))
- sorted_corner_list = sorted(corner_list, key=lambda x: x[0])
- if len(sorted_corner_list) == 6:
- left = sorted(sorted_corner_list[0:2], key=lambda x: x[1])
- middle = sorted(sorted_corner_list[2:4], key=lambda x: x[1])
- right = sorted(sorted_corner_list[4:6], key=lambda x: x[1])
- tmp_corner_list = [middle[0], left[0], left[1], middle[1], right[1], right[0], middle[0]]
- polygon1 = Polygon(tmp_corner_list[:4])
- polygon2 = Polygon(tmp_corner_list[3:])
- gen_polygon_list.remove(poly)
- for p in [polygon1, polygon2]:
- xi, yi, xx, yx = p.bounds
- w, h = xx - xi, yx - yi
- aspect_flag = max(w / h, h / w) < ASPECT_FLAG
- if aspect_flag:
- gen_polygon_list.append(p)
- gen_polygon_list = [polygon for index, polygon in enumerate(gen_polygon_list) if polygon.area > min_polygon.area]
- # if gen_choice:
- # gen_choice = sorted(gen_choice, key=lambda x: x.area)[-1]
- # gen_polygon_list.append(gen_choice)
- return gen_polygon_list
- def infer_class(image, sheet_dict_list, infer_polygon, image_cols, ocr_dict_list=''):
- res = []
- all_type_score_polygon = []
- all_choice_polygon = []
- all_cloze_polygon = []
- all_solve_polygon = []
- all_choice_s_width = []
- for region_box in sheet_dict_list:
- if region_box['class_name'] in ['type_score', 'choice', 'cloze', 'solve', 'choice_s']:
- coordinates = region_box['bounding_box']
- xmin = coordinates['xmin']
- ymin = coordinates['ymin']
- xmax = coordinates['xmax']
- ymax = coordinates['ymax']
- box_polygon = Polygon([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
- if region_box['class_name'] == 'type_score':
- all_type_score_polygon.append(box_polygon)
- if region_box['class_name'] == 'choice':
- all_choice_polygon.append(box_polygon)
- if region_box['class_name'] == 'cloze':
- all_cloze_polygon.append(box_polygon)
- if region_box['class_name'] == 'solve':
- all_solve_polygon.append(box_polygon)
- if region_box['class_name'] == 'choice_s':
- all_choice_s_width.append(int(xmax)-int(xmin))
- for poly in infer_polygon.copy(): # infer type_score solve
- p_xmin, p_ymin, p_xmax, p_ymax = poly.bounds
- type_score_num = 0
- type_score_ymin = []
- for type_score_polygon in all_type_score_polygon:
- cond1 = type_score_polygon.within(poly)
- cond2 = False
- cond3 = type_score_polygon.overlaps(poly)
- if cond3:
- intersection_poly = type_score_polygon.intersection(poly)
- d1 = intersection_poly.area / type_score_polygon.area >= TYPE_SCORE_MNS
- print('type_score:', intersection_poly.area / type_score_polygon.area)
- d2 = type_score_polygon.area < 0.2 * poly.area
- cond2 = d1 and d2
- if cond1 or cond2:
- type_score_num += 1
- t_xmin, t_ymin, t_xmax, t_ymax = type_score_polygon.bounds
- type_score_ymin.append(t_ymin)
- t_height = t_ymax - t_ymin
- if t_ymin - p_ymin > 3 * t_height:
- type_score_num += 1
- type_score_ymin.append(p_ymin)
- if type_score_num == 1:
- in_xmin, in_ymin, in_xmax, in_ymax = poly.bounds
- solve_box = {'class_name': 'solve',
- 'bounding_box': {'xmin': int(in_xmin), 'ymin': int(in_ymin),
- 'xmax': int(in_xmax), 'ymax': int(in_ymax)}}
- sheet_dict_list.append(solve_box)
- infer_polygon.remove(poly)
- res.append(solve_box)
- if type_score_num > 1: # 多type_score
- type_score_ymin = sorted(type_score_ymin)
- type_score_ymin[0] = min(p_ymin, type_score_ymin[0])
- type_score_ymin.append(p_ymax)
- for i in range(0, len(type_score_ymin) - 1):
- w = p_xmax - p_xmin
- h = type_score_ymin[i + 1] - type_score_ymin[i]
- if max(w / h, h / w) < ASPECT_FLAG:
- solve_box = {'class_name': 'solve',
- 'bounding_box': {'xmin': int(p_xmin), 'ymin': int(type_score_ymin[i]),
- 'xmax': int(p_xmax), 'ymax': int(type_score_ymin[i + 1])}}
- sheet_dict_list.append(solve_box)
- res.append(solve_box)
- infer_polygon.remove(poly)
- # for poly in infer_polygon.copy(): # infer choice_m
- # for choice_polygon in all_choice_polygon:
- # cond1 = choice_polygon.within(poly) or choice_polygon.contains(poly)
- # cond2 = False
- # cond3 = choice_polygon.overlaps(poly)
- # if cond3:
- # intersection_poly = choice_polygon.intersection(poly)
- # cond2 = intersection_poly.area / poly.area >= 0.8
- #
- # if cond1 or cond2:
- # in_xmin, in_ymin, in_xmax, in_ymax = poly.bounds
- # choice_m_img = crop_region_direct(image, (int(in_xmin), int(in_ymin),
- # int(in_xmax), int(in_ymax)))
- # # cv2.imshow('m', choice_m_img)
- # # cv2.waitKey(0)
- # ocr_res = get_ocr_text_and_coordinate(choice_m_img)
- # char_a_min = []
- # char_d_max = []
- # for index, chars in enumerate(ocr_res):
- # for char in chars['chars']:
- # left, top = char['location']['left'], char['location']['top']
- # width, height = char['location']['width'], char['location']['height']
- # if char['char'] in 'abcdlABCD[]aabbccddAABBCCDD[[]]':
- # xm, ym = int(left - width / 2), int(top - height / 2)
- # char_a_min.append((xm, ym))
- # xx, yx = int(left + 3 * width / 2), int(top + 3 * height / 2)
- # char_d_max.append((xx, yx))
- # if char_a_min and char_d_max:
- # char_a_min_arr, char_d_max_arr = np.array(char_a_min), np.array(char_d_max)
- # tmp_min = np.min(char_a_min_arr, axis=0)
- # tmp_max = np.max(char_d_max_arr, axis=0)
- #
- # m_xmin, m_ymin, m_xmax, m_ymax = tmp_min[0], tmp_min[1], tmp_max[0], tmp_max[1]
- # dif_width = sum(all_choice_s_width) // len(all_choice_s_width) - (m_xmax - m_xmin)
- # choice_box = {'class_name': 'choice_m',
- # 'bounding_box': {'xmin': int(m_xmin) + int(in_xmin) - dif_width // 2,
- # 'ymin': int(m_ymin) + int(in_ymin),
- # 'xmax': int(m_xmax) + int(in_xmin) + dif_width // 2,
- # 'ymax': int(m_ymax) + int(in_ymin)
- # }}
- #
- # sheet_dict_list.append(choice_box)
- # infer_polygon.remove(poly)
- # res.append(choice_box)
- # break
- for poly in infer_polygon.copy(): # infer ocr blank
- flag = []
- for ocr in ocr_dict_list:
- location = ocr['location']
- xmin = location['left']
- ymin = location['top']
- xmax = location['left'] + location['width']
- ymax = location['top'] + location['height']
- box_polygon = Polygon([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
- cond1 = poly.within(box_polygon) or poly.contains(box_polygon)
- cond2 = False
- cond3 = box_polygon.overlaps(poly)
- if cond3:
- intersection_poly = box_polygon.intersection(poly)
- cond2 = intersection_poly.area / poly.area >= 0.2
- flag.append(cond1 or cond2 or False) # True 不是blank
- if True not in flag:
- in_xmin, in_ymin, in_xmax, in_ymax = poly.bounds
- blank_box = {'class_name': 'blank',
- 'bounding_box': {'xmin': int(in_xmin), 'ymin': int(in_ymin),
- 'xmax': int(in_xmax), 'ymax': int(in_ymax)}}
- # sheet_dict_list.append(solve_box)
- infer_polygon.remove(poly)
- res.append(blank_box)
- for poly in infer_polygon.copy(): # infer blank
- bounds = [int(ele) for ele in poly.bounds]
- img_region = crop_region_direct(image, bounds)
- img = cv2.threshold(img_region, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
- img_mean = np.mean(img)
- img_raw_mean = np.mean(img_region)
- # print(img_mean, img_raw_mean)
- cond = img_mean < PIX_VALUE_LOW or img_raw_mean > PIX_VALUE_HIGH
- if cond:
- in_xmin, in_ymin, in_xmax, in_ymax = bounds
- blank_box = {'class_name': 'blank',
- 'bounding_box': {'xmin': int(in_xmin), 'ymin': int(in_ymin),
- 'xmax': int(in_xmax), 'ymax': int(in_ymax)}}
- # sheet_dict_list.append(solve_box)
- infer_polygon.remove(poly)
- res.append(blank_box)
- # for poly in infer_polygon.copy(): # infer cloze_s
- # for cloze_polygon in all_cloze_polygon:
- # cond1 = cloze_polygon.within(poly) or cloze_polygon.contains(poly)
- # cond2 = False
- # cond3 = cloze_polygon.overlaps(poly)
- # if cond3:
- # intersection_poly = cloze_polygon.intersection(poly)
- # cond2 = intersection_poly.area / poly.area >= 0.8
- #
- # if cond1 or cond2:
- # in_xmin, in_ymin, in_xmax, in_ymax = poly.bounds
- # solve_box = {'class_name': 'cloze_s',
- # 'bounding_box': {'xmin': int(in_xmin), 'ymin': int(in_ymin),
- # 'xmax': int(in_xmax), 'ymax': int(in_ymax)}}
- #
- # sheet_dict_list.append(solve_box)
- # infer_polygon.remove(poly)
- # res.append(solve_box)
- # break
- for poly in infer_polygon.copy(): # infer solve
- in_xmin, in_ymin, in_xmax, in_ymax = poly.bounds
- w, h = in_xmax - in_xmin, in_ymax - in_ymin
- aspect_flag = max(w / h, h / w) < ASPECT_FLAG
- if aspect_flag:
- solve_box = {'class_name': 'solve_infer',
- 'bounding_box': {'xmin': int(in_xmin), 'ymin': int(in_ymin),
- 'xmax': int(in_xmax), 'ymax': int(in_ymax)}}
- else:
- solve_box = {'class_name': 'blank',
- 'bounding_box': {'xmin': int(in_xmin), 'ymin': int(in_ymin),
- 'xmax': int(in_xmax), 'ymax': int(in_ymax)}}
- sheet_dict_list.append(solve_box)
- infer_polygon.remove(poly)
- res.append(solve_box)
- if all_type_score_polygon:
- type_score_area = sum([ele.area for ele in all_type_score_polygon])
- mean_type_score_area = type_score_area/len(all_type_score_polygon)
- solve_filter = []
- for index, sheet_box in enumerate(sheet_dict_list.copy()):
- if sheet_box['class_name'] == 'solve_infer':
- w = sheet_box['bounding_box']['xmax'] - sheet_box['bounding_box']['xmin']
- h = sheet_box['bounding_box']['ymin'] - sheet_box['bounding_box']['ymin']
- if w * h < mean_type_score_area * 3:
- sheet_dict_list.remove(sheet_box)
- for ele in sheet_dict_list:
- if ele['class_name'] == 'solve_infer':
- ele.update({'class_name': 'solve'})
- return sheet_dict_list
- def box_infer_and_complete(image, sheet_region_dict, ocr=''):
- if len(image.shape) == 3:
- image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
- if len(image.shape) == 4:
- image = cv2.cvtColor(image, cv2.COLOR_RGBA2GRAY)
- exclude_classes = [
- 'cloze_s',
- 'exam_number_s',
- 'type_score',
- 'page',
- 'alarm_info',
- # 'score_collect',
- 'choice_m'
- 'choice_s',
- ] # 不找这些区域的延长线
- y, x = image.shape[0], image.shape[1]
- x1, x2 = subfield_answer_sheet(image, sheet_region_dict)
- # lon_split_line = []
- lon_split_line = [LineString([(px, 1), (px, y - 1)]) for px in [x1, x2] if px != 0]
- split_line_poly = [(px, 1, px + 1, y - 1) for px in [x1, x2] if px != 0]
- poly_list = infer_sheet_box(image, sheet_region_dict, lon_split_line, exclude_classes)
- image_cols = len(lon_split_line) + 1
- sheet_region_dict = infer_class(image, sheet_region_dict, poly_list, image_cols, ocr)
- return sheet_region_dict
- def infer_solve(sheet_dict_list, left, right, top, bottom, h, w, col_regions, col_split, subject='math'):
- col_split.insert(0, left)
- col_split.append(right)
- boundary_list = [(split, 1, col_split[i+1]-1, h-1) for i, split in enumerate(col_split[:-1])]
- infer_polygon = []
- for col_i, col in enumerate(col_regions):
- if not col:
- continue
- class_names = [ele['class_name'] for ele in col]
- # 确定是正面第一栏
- if ((subject != 'chinese')
- and
- (any([ele in class_names for ele in ['infer_title', 'bar_code', 'choice_m', 'cloze_s']]))):
- bottom_box = col[-1]
- bottom_box_ymax = bottom_box['bounding_box']['ymax']
- infer_loc = {'xmin': boundary_list[col_i][0], 'ymin': bottom_box_ymax+5,
- 'xmax': boundary_list[col_i][2], 'ymax': boundary_list[col_i][3]}
- box_polygon = Polygon([(infer_loc['xmin'], infer_loc['ymin']),
- (infer_loc['xmax'], infer_loc['ymin']),
- (infer_loc['xmax'], infer_loc['ymax']),
- (infer_loc['xmin'], infer_loc['ymax']),])
- infer_polygon.append(box_polygon)
- else:
- y_axis = np.zeros(h)
- y_axis[top:top+1] = 1
- y_axis[bottom-1:bottom] = 1
- for ele in col:
- ymin = ele['bounding_box']['ymin']
- ymax = ele['bounding_box']['ymax']
- y_axis[ymin:ymax+1] = 1
- split_index_list = []
- split_index = 0
- for i, ele in enumerate(y_axis):
- split_index = split_index % 2
- if ele == split_index:
- # print(i)
- split_index = split_index + 1
- split_index_list.append(i)
- for i in range(0, len(split_index_list)-1, 2):
- ymin = split_index_list[i] + 3
- ymax = split_index_list[i+1] - 3
- infer_loc = {'xmin': boundary_list[col_i][0], 'ymin': ymin,
- 'xmax': boundary_list[col_i][2], 'ymax': ymax}
- box_polygon = Polygon([(infer_loc['xmin'], infer_loc['ymin']),
- (infer_loc['xmax'], infer_loc['ymin']),
- (infer_loc['xmax'], infer_loc['ymax']),
- (infer_loc['xmin'], infer_loc['ymax']), ])
- infer_polygon.append(box_polygon)
- res = []
- all_type_score_polygon = []
- for region_box in sheet_dict_list:
- if region_box['class_name'] in ['type_score']:
- coordinates = region_box['bounding_box']
- xmin = coordinates['xmin']
- ymin = coordinates['ymin']
- xmax = coordinates['xmax']
- ymax = coordinates['ymax']
- box_polygon = Polygon([(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)])
- all_type_score_polygon.append(box_polygon)
- for poly in infer_polygon.copy(): # infer type_score solve
- p_xmin, p_ymin, p_xmax, p_ymax = poly.bounds
- type_score_num = 0
- type_score_ymin = []
- for type_score_polygon in all_type_score_polygon:
- cond1 = type_score_polygon.within(poly)
- cond2 = False
- cond3 = type_score_polygon.overlaps(poly)
- if cond3:
- intersection_poly = type_score_polygon.intersection(poly)
- d1 = intersection_poly.area / type_score_polygon.area >= TYPE_SCORE_MNS
- print('type_score:', intersection_poly.area / type_score_polygon.area)
- d2 = type_score_polygon.area < 0.5 * poly.area
- cond2 = d1 and d2
- if cond1 or cond2:
- t_xmin, t_ymin, t_xmax, t_ymax = type_score_polygon.bounds
- x_dif = t_xmin - p_xmin
- type_score_num += 1
- type_score_ymin.append(t_ymin)
- t_height = t_ymax - t_ymin
- if t_ymin - p_ymin > 3 * t_height:
- type_score_num += 1
- type_score_ymin.append(p_ymin)
- if type_score_num == 1:
- w, h = p_xmax-p_xmin, p_ymax-p_ymin
- w, h = max(w, 0.1), max(h, 0.1)
- aspect_flag = max(w / h, h / w) < ASPECT_FLAG
- if aspect_flag:
- solve_box = {'class_name': 'solve',
- 'bounding_box': {'xmin': int(p_xmin), 'ymin': int(p_ymin),
- 'xmax': int(p_xmax), 'ymax': int(p_ymax)}}
- infer_polygon.remove(poly)
- res.append(solve_box)
- else:
- solve_box = {'class_name': 'solve_with_type_score_without_aspect',
- 'bounding_box': {'xmin': int(p_xmin), 'ymin': int(p_ymin),
- 'xmax': int(p_xmax), 'ymax': int(p_ymax)}}
- infer_polygon.remove(poly)
- res.append(solve_box)
- if type_score_num > 1: # 多type_score
- type_score_ymin = sorted(type_score_ymin)
- type_score_ymin[0] = min(p_ymin, type_score_ymin[0])
- type_score_ymin.append(p_ymax)
- for col_i in range(0, len(type_score_ymin) - 1):
- w = p_xmax - p_xmin
- h = type_score_ymin[col_i + 1] - type_score_ymin[col_i]
- w, h = max(w, 0.1), max(h, 0.1)
- aspect_flag = max(w / h, h / w) < ASPECT_FLAG
- if aspect_flag:
- solve_box = {'class_name': 'solve',
- 'bounding_box': {'xmin': int(p_xmin), 'ymin': int(type_score_ymin[col_i]-5),
- 'xmax': int(p_xmax), 'ymax': int(type_score_ymin[col_i + 1])}}
- res.append(solve_box)
- infer_polygon.remove(poly)
- for poly in infer_polygon.copy(): # infer solve
- in_xmin, in_ymin, in_xmax, in_ymax = poly.bounds
- w, h = in_xmax - in_xmin, in_ymax - in_ymin
- w, h = max(w, 0.1), max(h, 0.1)
- aspect_flag = max(w / h, h / w) < ASPECT_FLAG
- if aspect_flag:
- solve_box = {'class_name': 'solve_without_type_score',
- 'bounding_box': {'xmin': int(in_xmin), 'ymin': int(in_ymin),
- 'xmax': int(in_xmax), 'ymax': int(in_ymax)}}
- else:
- solve_box = {'class_name': 'w_h_blank',
- 'bounding_box': {'xmin': int(in_xmin), 'ymin': int(in_ymin),
- 'xmax': int(in_xmax), 'ymax': int(in_ymax)}}
- infer_polygon.remove(poly)
- res.append(solve_box)
- return res
|