123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441 |
- # @Author : lightXu
- # @File : analysis_sheet.py
- import time
- import os
- import traceback
- import numpy as np
- import cv2
- from django.conf import settings
- import segment.logging_config as logging
- from segment.sheet_resolve.lib.model.test import im_detect
- from segment.sheet_resolve.lib.model.nms_wrapper import nms
- from segment.sheet_resolve.lib.utils.timer import Timer
- from segment.sheet_resolve.tools import utils
- from segment.sheet_resolve.analysis.solve.optional_solve import resolve_optional_choice
- logger = logging.getLogger(settings.LOGGING_TYPE)
- def analysis_single_image_with_regions(analysis_type, classes,
- sess, net,
- im_raw, conf_thresh, mns_thresh,
- coordinate_bias_dict):
- """Detect object classes in an image using pre-computed object proposals."""
- size = im_raw.shape
- # Detect all object classes and regress object bounds
- timer = Timer()
- timer.tic()
- if '_blank' in analysis_type:
- analysis_type = analysis_type.replace('_blank', '')
- if analysis_type in ['unknown_subject', 'math', 'math_zxhx', 'english', 'chinese',
- 'physics', 'chemistry', 'biology', 'politics', 'history',
- 'geography', 'science_comprehensive', 'arts_comprehensive'
- ]:
- analysis_type = 'sheet'
- # im, ratio = utils.img_resize(analysis_type, im_raw)
- im, ratio = utils.resize_faster_rcnn(analysis_type, im_raw)
- scores, boxes = im_detect(analysis_type, sess, net, im)
- timer.toc()
- print('Detection took {:.3f}s for {:d} object proposals'.format(timer.total_time, boxes.shape[0]))
- content_list = []
- analysis_cls_list = []
- qr_code_info = 'Nan'
- for cls_ind, cls in enumerate(classes[1:]): # classes
- cls_ind += 1 # because we skipped background
- cls_boxes = boxes[:, 4 * cls_ind:4 * (cls_ind + 1)]
- cls_scores = scores[:, cls_ind]
- dets = np.hstack((cls_boxes,
- cls_scores[:, np.newaxis])).astype(np.float32)
- keep = nms(dets, mns_thresh)
- dets = dets[keep, :]
- inds = np.where(dets[:, -1] >= conf_thresh)[0]
- if len(inds) > 0:
- if cls in list(coordinate_bias_dict.keys()):
- xmin_bias = coordinate_bias_dict[cls]['xmin_bias']
- ymin_bias = coordinate_bias_dict[cls]['ymin_bias']
- xmax_bias = coordinate_bias_dict[cls]['xmax_bias']
- ymax_bias = coordinate_bias_dict[cls]['ymax_bias']
- else:
- xmin_bias = 0
- ymin_bias = 0
- xmax_bias = 0
- ymax_bias = 0
- for i in inds:
- bbox = dets[i, :4]
- score = '{:.4f}'.format(dets[i, -1])
- xmin = int(int(bbox[0]) * ratio[0]) + xmin_bias
- ymin = int(int(bbox[1]) * ratio[1]) + ymin_bias
- xmax = int(int(bbox[2]) * ratio[0]) + xmax_bias
- ymax = int(int(bbox[3]) * ratio[1]) + ymax_bias
- xmin = (xmin if (xmin > 0) else 1)
- ymin = (ymin if (ymin > 0) else 1)
- xmax = (xmax if (xmax < size[1]) else size[1] - 1)
- ymax = (ymax if (ymax < size[0]) else size[0] - 1)
- if cls in ['solve0', 'composition0']:
- cls = cls.replace('0', '')
- bbox_dict = {"xmin": xmin, "ymin": ymin, "xmax": xmax, "ymax": ymax}
- class_dict = {"class_name": cls, "bounding_box": bbox_dict, "score": score}
- # if cls == 'qr_code':
- # qr_img = utils.crop_region(im_raw, bbox_dict)
- # qr_path = r'./qr_code.jpg'
- # cv2.imwrite(qr_path, qr_img)
- # qr_code_info = utils.check_qr_code_with_region_img(qr_path)
- # os.remove(qr_path)
- content_list.append(class_dict)
- return content_list, analysis_cls_list, qr_code_info
- def get_single_image_sheet_regions(analysis_type, img_path, img, classes,
- sess, net, conf_thresh, mns_thresh,
- coordinate_bias_dict):
- start_time = time.time()
- print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
- print('analysis for JPG {}'.format(img_path))
- content, cls, qr_code_info = \
- analysis_single_image_with_regions(analysis_type, classes, sess, net,
- img, conf_thresh, mns_thresh,
- coordinate_bias_dict)
- analysis_type.replace('_blank', '')
- img_dict = {"img_name": img_path,
- # 'qr_code': qr_code_info,
- 'subject': analysis_type,
- "regions": content,
- }
- end_time = time.time()
- print(end_time - start_time)
- return img_dict
- def question_number_format(init_number, crt_numbers, regions):
- """
- 将重复或者是-1的题号改为501
- :param init_number: 初始替换的题号
- :param crt_numbers: 目前已经有的题号
- :param regions: 答题卡各区域
- :return:
- """
- logger.info('regions: {}'.format(regions))
- for region in regions:
- logger.info('region: {}'.format(region))
- if region['class_name'] == 'optional_choice':
- continue
- while init_number in crt_numbers:
- init_number += 1
- numbers = region.get("number")
- if numbers and (isinstance(numbers, int) or isinstance(numbers, float)):
- if numbers <= 0 or numbers in crt_numbers or numbers >= 1000:
- if not region.get("span"):
- numbers = init_number
- crt_numbers.append(numbers)
- init_number += 1
- region.update({"number": numbers})
- crt_numbers.append(numbers)
- if numbers and isinstance(numbers, list):
- for i, num in enumerate(numbers):
- if num <= 0 or num in crt_numbers or num >= 1000:
- if not region.get("span"):
- numbers[i] = init_number
- crt_numbers.append(init_number)
- init_number += 1
- region.update({"number": numbers})
- crt_numbers.extend(numbers)
- return regions, init_number, crt_numbers
- def resolve_select_s(image, bbox):
- box_region = utils.crop_region(image, bbox)
- left = bbox['xmin']
- top = bbox['ymin']
- right = bbox['xmax']
- bottom = bbox['ymax']
- if (right - left) >= (bottom - top):
- direction = 180
- else:
- direction = 90
- try:
- res = resolve_optional_choice(left, top, direction, box_region)
- except Exception as e:
- res = {'class_name': 'optional_choice',
- 'rows': 1, 'cols': 2,
- 'number': [501, 502],
- 'single_width': right - left,
- 'single_height': bottom - top,
- 'bounding_box': {'xmin': left,
- 'ymin': top,
- 'xmax': right,
- 'ymax': bottom}}
- return res
- def box_region_format(sheet_dict, image, subject, shrink=True):
- include_class = ['anchor_point',
- 'bar_code',
- 'choice_m',
- 'cloze',
- 'cloze_s',
- 'exam_number_col_row',
- 'optional_choice',
- 'optional_solve',
- # 'qr_code',
- 'solve',
- 'optional_solve',
- 'composition',
- # 'correction'
- ]
- default_points_dict = {'choice_m': 5, "cloze": 5, 'solve': 12, 'optional_solve': 10, 'cloze_s': 5,
- "composition": 60}
- if subject in ["english", 'physics', 'chemistry', 'biology', 'science_comprehensive']:
- default_points_dict = {'choice_m': 2, "cloze": 2, 'solve': 10, 'optional_solve': 10, 'cloze_s': 2,
- "composition": 25}
- sheet_regions = sheet_dict['regions']
- select_s_list = []
- for i in range(len(sheet_regions) - 1, -1, -1):
- if subject == "math":
- if sheet_regions[i]['class_name'] == 'cloze':
- sheet_regions[i]['class_name'] = 'cloze_big' # math exclude cloze big
- if sheet_regions[i]['class_name'] == 'cloze_s':
- sheet_regions[i]['class_name'] = 'cloze' # math exclude cloze big
- if subject == "english":
- if sheet_regions[i]['class_name'] == 'cloze':
- sheet_regions[i]['class_name'] = 'solve'
- if sheet_regions[i]['class_name'] == 'correction':
- sheet_regions[i]['class_name'] = 'solve'
- if sheet_regions[i]['class_name'] in ['solve0']:
- sheet_regions[i]['class_name'] = 'solve'
- if sheet_regions[i]['class_name'] in ['composition0']:
- sheet_regions[i]['class_name'] = 'composition'
- if sheet_regions[i]['class_name'] == 'select_s':
- select_s_list.append(sheet_regions[i])
- if shrink:
- if sheet_regions[i]['class_name'] not in include_class:
- sheet_regions.pop(i)
- # 去重
- sheet_tmp = sheet_regions.copy()
- remove_index = []
- for i, region in enumerate(sheet_tmp):
- if i not in remove_index:
- box = region['bounding_box']
- name = region['class_name']
- for j, region_in in enumerate(sheet_tmp):
- box_in = region_in['bounding_box']
- name_in = region_in['class_name']
- iou = utils.cal_iou(box, box_in)
- if name == name_in and (iou[0] > 0.75 or iou[1] > 0.85 or iou[2] > 0.85) and i != j:
- box_area = (box['xmax'] - box['xmin']) * (box['ymax'] - box['ymin'])
- box_in_area = (box_in['xmax'] - box_in['xmin']) * (box_in['ymax'] - box_in['ymin'])
- if box_area >= box_in_area:
- sheet_regions.remove(region_in)
- remove_index.append(j)
- else:
- sheet_regions.remove(region)
- remove_index.append(i)
- break
- # 合并select_s
- optional_choice_tmp = []
- select_s_list_copy = select_s_list.copy()
- if len(select_s_list) > 0:
- for ele in sheet_regions:
- if ele['class_name'] == 'solve':
- solve_box = (ele['bounding_box']['xmin'], ele['bounding_box']['ymin'],
- ele['bounding_box']['xmax'], ele['bounding_box']['ymax'])
- xn, yn, xm, ym = 9999, 9999, 0, 0
- merge = False
- for select_s in select_s_list:
- select_s_box = (select_s['bounding_box']['xmin'], select_s['bounding_box']['ymin'],
- select_s['bounding_box']['xmax'], select_s['bounding_box']['ymax'])
- if utils.decide_coordinate_contains(select_s_box, solve_box):
- merge = True
- xn = min(xn, select_s_box[0])
- yn = min(yn, select_s_box[1])
- xm = max(xm, select_s_box[2])
- ym = max(ym, select_s_box[3])
- select_s_list_copy.remove(select_s)
- if merge:
- new_box = {'xmin': xn, 'ymin': yn, 'xmax': xm, 'ymax': ym}
- optional_choice_info = resolve_select_s(image, new_box)
- optional_choice_tmp.append(optional_choice_info)
- for ele in select_s_list_copy:
- box = ele['bounding_box']
- optional_choice_info = resolve_select_s(image, box)
- optional_choice_tmp.append(optional_choice_info)
- optional_choice_tmp_ = optional_choice_tmp.copy()
- for ele in sheet_regions:
- if len(optional_choice_tmp) > 0:
- if ele['class_name'] == 'solve':
- solve_box = (ele['bounding_box']['xmin'], ele['bounding_box']['ymin'],
- ele['bounding_box']['xmax'], ele['bounding_box']['ymax'])
- for optional_choice in optional_choice_tmp_:
- optional_choice_box = (optional_choice['bounding_box']['xmin'], optional_choice['bounding_box']['ymin'],
- optional_choice['bounding_box']['xmax'], optional_choice['bounding_box']['ymax'])
- if utils.decide_coordinate_contains(optional_choice_box, solve_box):
- optional_choice_tmp.remove(optional_choice)
- ele['class_name'] = 'optional_solve'
- choice_numbers = optional_choice['number']
- solve_numbers = ele['number']
- if choice_numbers[0] < 500:
- ele['number'] = choice_numbers
- ele['default_points'] = [ele['default_points']] * len(choice_numbers)
- else:
- tmp = [solve_numbers] * len(choice_numbers)
- for i, num in enumerate(tmp):
- tmp[i] = num + i
- ele['number'] = tmp
- optional_choice['number'] = tmp
- ele['default_points'] = [ele['default_points']] * len(choice_numbers)
- break
- else:
- continue
- if ele['class_name'] in ["choice_m", "cloze", "cloze_s", "solve", "optional_solve", "composition"]:
- if isinstance(ele['default_points'], list):
- for i, dp in enumerate(ele['default_points']):
- if dp < 1: # 小于一分
- ele['default_points'][i] = default_points_dict[ele['class_name']]
- if isinstance(ele['default_points'], int) or isinstance(ele['default_points'], float):
- if ele['default_points'] < 1: # 小于一分
- ele['default_points'] = default_points_dict[ele['class_name']]
- # select_s 在解答区域外侧
- if len(optional_choice_tmp) > 0:
- for oc in optional_choice_tmp:
- optional_choice_box = (oc['bounding_box']['xmin'], oc['bounding_box']['ymin'],
- oc['bounding_box']['xmax'], oc['bounding_box']['ymax'],
- oc['bounding_box']['xmin']
- + (oc['bounding_box']['xmax'] - oc['bounding_box']['xmin']) // 2)
- for sr in sheet_regions:
- if sr['class_name'] == 'solve':
- solve_box = (sr['bounding_box']['xmin'], sr['bounding_box']['ymin'],
- sr['bounding_box']['xmax'], sr['bounding_box']['ymax'])
- if (optional_choice_box[1] <= solve_box[1] and
- solve_box[0] < optional_choice_box[4] < solve_box[2] and
- abs(optional_choice_box[1] - solve_box[1]) < solve_box[3] - solve_box[1]):
- sr['class_name'] = 'optional_solve'
- choice_numbers = oc['number']
- solve_numbers = sr['number']
- if choice_numbers[0] < 500:
- sr['number'] = choice_numbers
- sr['default_points'] = [sr['default_points']] * len(choice_numbers)
- else:
- tmp = [solve_numbers] * len(choice_numbers)
- for i, num in enumerate(tmp):
- tmp[i] = num + i
- sr['number'] = tmp
- oc['number'] = tmp
- sr['default_points'] = [sr['default_points']] * len(choice_numbers)
- break
- if len(optional_choice_tmp_):
- sheet_regions.extend(optional_choice_tmp_)
- # 去重
- sheet_tmp = sheet_regions.copy()
- remove_index = []
- for i, region in enumerate(sheet_tmp):
- if i not in remove_index:
- box = region['bounding_box']
- name = region['class_name']
- for j, region_in in enumerate(sheet_tmp):
- box_in = region_in['bounding_box']
- name_in = region_in['class_name']
- iou = utils.cal_iou(box, box_in)
- if name == name_in and (iou[0] > 0.75 or iou[1] > 0.85 or iou[2] > 0.85) and i != j:
- box_area = (box['xmax'] - box['xmin']) * (box['ymax'] - box['ymin'])
- box_in_area = (box_in['xmax'] - box_in['xmin']) * (box_in['ymax'] - box_in['ymin'])
- if box_area >= box_in_area:
- sheet_regions.remove(region_in)
- remove_index.append(j)
- else:
- sheet_regions.remove(region)
- remove_index.append(i)
- break
- sheet_dict.update({'regions': sheet_regions})
- return sheet_dict
- def merge_span_boxes(col_sheets):
- if len(col_sheets) <= 1:
- return col_sheets
- for i, cur_col in enumerate(col_sheets[:-1]):
- next_col = col_sheets[i + 1]
- if not cur_col or not next_col:
- continue
- current_bottom_box = cur_col[-1] # 当前栏的最后一个,bottom
- next_col_top_box = next_col[0] # 下一栏的第一个,top
- b_name = current_bottom_box['class_name']
- t_name = next_col_top_box['class_name']
- if b_name == t_name == 'solve':
- b_number = current_bottom_box['number']
- t_number = next_col_top_box['number']
- if b_number >= 500 or t_number >= 500 or b_number == t_number:
- numbers = min(b_number, t_number)
- crt_points = current_bottom_box['default_points']
- next_points = next_col_top_box['default_points']
- # default_points = max(current_bottom_box['default_points'], next_col_top_box['default_points'])
- default_points = crt_points
- current_bottom_box.update({'number': numbers, 'default_points': default_points, "span": True})
- next_col_top_box.update({'number': numbers, 'default_points': default_points,
- "span": True, "span_id": current_bottom_box["span_id"] + 1})
- elif b_name == t_name == 'composition':
- b_number = current_bottom_box['number']
- t_number = next_col_top_box['number']
- numbers = min(b_number, t_number)
- default_points = max(current_bottom_box['default_points'], next_col_top_box['default_points'])
- current_bottom_box.update({'number': numbers, 'default_points': default_points, "span": True})
- next_col_top_box.update({'number': numbers, 'default_points': default_points,
- "span": True, "span_id": current_bottom_box["span_id"] + 1})
- else:
- continue
- return col_sheets
|