exam-segment-django/segment/sheet_resolve/analysis/sheet/analysis_sheet.py

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