exam-segment-django/segment/sheet_resolve/analysis/choice/choice_box.py

# @Author  : lightXu
# @File    : choice_box.py
# @Time    : 2018/11/22 0022 下午 16:01
import re
import time
import xml.etree.cElementTree as ET

import cv2
import numpy as np

from segment.sheet_resolve.analysis.choice.choice_m_row_column import get_choice_m_row_and_col
from segment.sheet_resolve.tools import utils
from segment.sheet_resolve.tools.brain_api import get_ocr_text_and_coordinate


def get_interval(word_result_list):
    all_char_str = ''
    location = []
    for i, chars_dict in enumerate(word_result_list):
        chars_list = chars_dict['chars']
        for ele in chars_list:
            all_char_str = all_char_str + ele['char']
            location.append(ele['location'])

    pattern1 = re.compile(r"\]\[")
    pattern2 = re.compile(r"\[[ABCD]")

    def intervel(pattern):
        group_list = []
        for i in pattern.finditer(all_char_str):
            # print(i.group() + str(i.span()))
            group_list.append(list(i.span()))
        # print(group_list)

        sum_intervel = 0
        size = 0
        for group in group_list:
            left_x, right_x = location[group[0]]['left'] \
                              + location[group[0]]['width'], location[group[1] - 1]['left']
            if abs(location[group[0]]['top'] - location[group[1]]['top']) < location[group[0]]['height']:
                if right_x - left_x > 0:
                    sum_intervel = sum_intervel + right_x - left_x
                    size += 1

        # print(sum_intervel // size)
        return sum_intervel // size

    intervel_width1 = intervel(pattern1)
    intervel_width2 = intervel(pattern2)

    return (intervel_width1 + intervel_width2) * 2 // 3


def preprocess(image0, xe, ye):
    scale = 0
    dilate = 1
    blur = 5
    # 预处理图像
    img = image0

    # rescale the image
    if scale != 0:
        img = cv2.resize(img, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)

    # Convert to gray
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # # Apply dilation and erosion to remove some noise
    # if dilate != 0:
    #     kernel = np.ones((dilate, dilate), np.uint8)
    #     img = cv2.dilate(img, kernel, iterations=1)
    #     img = cv2.erode(img, kernel, iterations=1)

    # Apply blur to smooth out the edges
    # if blur != 0:
    #     img = cv2.GaussianBlur(img, (blur, blur), 0)

    # Apply threshold to get image with only b&w (binarization)
    img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]

    # cv2.namedWindow('image', cv2.WINDOW_NORMAL)
    # cv2.imshow('image', img)
    # if cv2.waitKey(0) == 27:
    #     cv2.destroyAllWindows()
    # cv2.imwrite('otsu.jpg', img)

    kernel = np.ones((ye, xe), np.uint8)  # y轴膨胀, x轴膨胀

    dst = cv2.dilate(img, kernel, iterations=1)
    # cv2.imshow('dilate', dst)
    # if cv2.waitKey(0) == 27:
    #     cv2.destroyAllWindows()

    return dst


def contours(image):
    _, cnts, hierarchy = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    bboxes = []
    for cnt_id, cnt in enumerate(reversed(cnts)):
        x, y, w, h = cv2.boundingRect(cnt)
        bboxes.append((x, y, x + w, y + h))

    return bboxes


def box_coordinates(img):
    img_arr = np.asarray(img)

    def axix_break_point(img, tolerance_number, axis):
        sum_x_axis = img.sum(axis=axis)
        sum_x_axis[sum_x_axis > 255 * tolerance_number] = 1  # 白色有字
        sum_x_axis[sum_x_axis != 1] = 0  # 黑色无字
        sum_x_axis_list = list(sum_x_axis)
        sum_x_axis_list.append(0)  # 最后几行到结束有字时，使索引值增加最后一位

        split_x_index = []
        num = 1
        for index, ele in enumerate(sum_x_axis_list):
            num = num % 2
            if ele == num:
                # print(i)
                num = num + 1
                split_x_index.append(index)
        # print('length: ', len(split_x_index), split_x_index)
        return split_x_index

    y_break_points_list = axix_break_point(img_arr, 1, axis=1)
    x_break_points_list = axix_break_point(img_arr, 1, axis=0)

    all_coordinates = []
    for i in range(0, len(y_break_points_list), 2):  # y轴分组
        ymin = y_break_points_list[i]
        ymax = y_break_points_list[i + 1]
        for j in range(0, len(x_break_points_list), 2):
            xmin = x_break_points_list[j]
            xmax = x_break_points_list[j + 1]
            all_coordinates.append([xmin, ymin, xmax, ymax])

    return all_coordinates


def get_choice_box_coordinate(word_result_list, choice_img, cv_box_list, choice_bbox_list):
    shape = choice_img.shape
    y, x = shape[0], shape[1]

    # cv2.imshow('ocr_region', ocr_region)
    # if cv2.waitKey(0) == 27:
    #     cv2.destroyAllWindows()

    all_digital_list = []
    digital_model = re.compile(r'\d')
    for i, chars_dict in enumerate(word_result_list):
        chars_list = chars_dict['chars']
        for ele in chars_list:
            if digital_model.search(ele['char']):
                all_digital_list.append(ele)

    new_all_digital_list = []
    i = 1
    while i <= len(all_digital_list):
        pre_one = all_digital_list[i - 1]
        if i == len(all_digital_list):
            new_all_digital_list.append(pre_one)
            break
        rear_one = all_digital_list[i]
        condition1 = abs(pre_one['location']['top'] - rear_one['location']['top']) < pre_one['location'][
            'height']  # 两字高度差小于一字高度
        condition2 = pre_one['location']['left'] + 2 * pre_one['location']['width'] > rear_one['location'][
            'left']  # 某字宽度的2倍大于两字间间隔
        if condition1:
            if condition2:
                new_char = pre_one['char'] + rear_one['char']
                new_location = {'left': pre_one['location']['left'],
                                'top': min(pre_one['location']['top'], rear_one['location']['top']),
                                'width': rear_one['location']['left'] + rear_one['location']['width'] -
                                         pre_one['location']['left'],
                                'height': max(pre_one['location']['height'], rear_one['location']['height'])}
                new_all_digital_list.append({'char': new_char, 'location': new_location})
                i = i + 1 + 1
            else:
                new_all_digital_list.append(pre_one)
                i = i + 1
        else:
            new_all_digital_list.append(pre_one)  # 遇到字符y轴相差过大就结束
            i = i + 1

    content_list = list()
    for index, box in enumerate(choice_bbox_list['regions']):  # rcnn识别的框匹配题号
        box = box['bounding_box']
        box_coordinate = (box['xmin'], box['ymin'], box['xmax'], box['ymax'])
        horizontal = box['xmax'] - box['xmin'] >= box['ymax'] - box['ymin']
        vertical = box['xmax'] - box['xmin'] < box['ymax'] - box['ymin']
        choice_number = {'number': 99, 'location': box_coordinate}
        content_list.insert(index, choice_number)
        for digital in new_all_digital_list:
            digital_coordiante = (digital['location']['left'], digital['location']['top'],
                                  digital['location']['left'] + digital['location']['width'],
                                  digital['location']['top'] + digital['location']['height'])

            if utils.decide_coordinate_contains(digital_coordiante, box_coordinate):
                if horizontal:
                    box['xmin'] = digital['location']['left'] + digital['location']['width'] + 1  # 从数字处截取
                if vertical:
                    box['ymin'] = digital['location']['top'] + digital['location']['height'] + 1

                box_coordinate = (box['xmin'], box['ymin'], box['xmax'], box['ymax'])
                content_list[index]['number'] = digital['char']
                content_list[index]['location'] = box_coordinate
                break

    for box in content_list:
        box_coordinate = (box['location'][0], box['location'][1], box['location'][2], box['location'][3])
        mtx = []
        for cv_box in cv_box_list:
            if utils.decide_coordinate_contains(cv_box, box_coordinate):  # 若fasterrcnn未识别到选项框，单独的ABCD也舍去
                mtx.append(cv_box)

        matrix = np.asarray(sorted(mtx))
        dif = matrix[1:, 0] - matrix[:-1, 2]  # 后一个char的left与起一个char的right的差
        dif[dif < 0] = 0
        dif_length = np.mean(dif)  # 小于平均间隔的合并
        block_list = utils.box_by_x_intervel(matrix, dif_length)
        # block_list = utils.box_by_x_intervel(matrix, 5)
        box['abcd'] = block_list

    return content_list


def choice(left, top, image, choice_bbox_list, xml_path):
    a_z = '_ABCDEFGHIJKLMTUNOPQRSVWXYZ'
    t1 = time.time()
    word_result_list0 = get_ocr_text_and_coordinate(image, ocr_accuracy='accurate', language_type='ENG')
    t2 = time.time()
    print('choice ocr time cost: ', t2 - t1)
    # print(word_result_list0)

    # try:
    #     intervel_x = get_interval(word_result_list0)
    # except Exception:
    #     intervel_x = 15
    intervel_x = 3
    img = preprocess(image, intervel_x, 3)
    cv_box_list0 = box_coordinates(img)

    content_list = get_choice_box_coordinate(word_result_list0, image, cv_box_list0, choice_bbox_list)

    tree = ET.parse(xml_path)  # xml tree

    w = content_list[0]['location'][2] - content_list[0]['location'][0]
    h = content_list[0]['location'][3] - content_list[0]['location'][1]

    def xml(xml_tree, sorted_abcd_list, bias=0):
        ii = 0
        for i, choice_bbox in enumerate(sorted_abcd_list):
            area = (choice_bbox[2] - choice_bbox[0]) * (choice_bbox[3] - choice_bbox[1])
            if area > 400:
                name = '{:02d}_{}'.format(int(choice['number']), a_z[ii + bias])
                xml_tree = utils.create_xml(name, xml_tree,
                                            choice_bbox[0] + left, choice_bbox[1] + top, choice_bbox[2] + left,
                                            choice_bbox[3] + top)
                ii += 1
        return xml_tree

    def get_json(ajson_list, sorted_abcd_list, bias=0):
        ii = 0
        for i, choice_bbox in enumerate(sorted_abcd_list):
            area = (choice_bbox[2] - choice_bbox[0]) * (choice_bbox[3] - choice_bbox[1])
            if area > 400:
                name = '{:02d}_{}'.format(int(choice['number']), a_z[ii + bias])
                region = [choice_bbox[0] + left, choice_bbox[1] + top, choice_bbox[2] + left, choice_bbox[3] + top]
                ajson_list.append({'number': name, 'region': region})
                ii += 1
        return ajson_list

    json_list = []
    for index_num, choice in enumerate(content_list):
        abcd = choice['abcd']
        if int(choice['number']) == 99:
            if w >= h:
                tree = xml(tree, sorted(abcd))
                json_list = get_json(json_list, sorted(abcd))

            else:
                tree = xml(tree, sorted(abcd, key=lambda x: (x[1], x[0])))
                json_list = get_json(json_list, sorted(abcd, key=lambda x: (x[1], x[0])))

        else:
            if w >= h:
                tree = xml(tree, sorted(abcd), bias=1)
                json_list = get_json(json_list, sorted(abcd), bias=1)

            else:
                tree = xml(tree, sorted(abcd, key=lambda x: (x[1], x[0])), bias=1)
                json_list = get_json(json_list, sorted(abcd, key=lambda x: (x[1], x[0])), bias=1)

    tree.write(xml_path)
    return json_list


def get_number_by_enlarge_choice_m(image, choice_m_region_list, xml_path):
    a_z = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

    choice_m_dict_list = []  # choice_m region with same index
    choice_m_enlarge = []
    left, top, right, bottom = 9999, 9999, 0, 0
    for _, box in enumerate(choice_m_region_list):
        box = box['bounding_box']
        m_left, m_top = box['xmin'], box['ymin'],
        width, height = box['xmax'] - box['xmin'], box['ymax'] - box['ymin']

        box_coordinate = (m_left, m_top, box['xmax'], box['ymax'])
        single_choice_m = utils.crop_region_direct(image, box_coordinate)
        row_col_dict = get_choice_m_row_and_col(m_left, m_top, single_choice_m)
        choice_m_dict_list.append(row_col_dict)

        box_coordinate_enlarge = (
            m_left - int(width / 2), m_top - int(height / 2), box['xmax'], box['ymax'])  # 扩大的choice_m， 多个分散choice_m
        choice_m_enlarge.append(box_coordinate_enlarge)
        left = min(left, box_coordinate_enlarge[0])
        top = min(top, box_coordinate_enlarge[1])
        right = max(right, box_coordinate_enlarge[2])
        bottom = max(bottom, box_coordinate_enlarge[3])

    choice_whole_region = utils.crop_region_direct(image, (left, top, right, bottom))
    # cv2.imwrite(r'C:\Users\Administrator\Desktop\test\sheet\choice_enlarge.jpg', choice_whole_region)
    # cv2.imshow('img', choice_whole_region)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    choice_region_text = get_ocr_text_and_coordinate(choice_whole_region)
    all_digital_list = []
    pattern = re.compile(r'\d')
    for i, chars_dict in enumerate(choice_region_text):
        chars_list = chars_dict['chars']
        for ele in chars_list:
            if pattern.search(ele['char']):
                all_digital_list.append(ele)

    combined_digital_list = utils.combine_char(all_digital_list)
    direction_list = []
    for index, enlarge_box in enumerate(choice_m_enlarge):
        digital_list = []
        xmin, ymin, xmax, ymax = 9999, 9999, 0, 0

        choice_m_dict = choice_m_dict_list[index]
        choice_m_dict_box = (choice_m_dict['bounding_box']['xmin'], choice_m_dict['bounding_box']['ymin'],
                             choice_m_dict['bounding_box']['xmax'], choice_m_dict['bounding_box']['ymax'],)

        for jndex, digital_box in enumerate(combined_digital_list):
            digital_coordinate = (digital_box['location']['left'] + left,
                                  digital_box['location']['top'] + top,
                                  digital_box['location']['left'] + digital_box['location']['width'] + left,
                                  digital_box['location']['top'] + digital_box['location']['height'] + top)
            digital_box.update({'coordinate': digital_coordinate})
            if (utils.decide_coordinate_contains(digital_coordinate, enlarge_box)) and not \
                    (utils.decide_coordinate_contains(digital_coordinate, choice_m_dict_box)):
                digital_list.append(digital_box)
                xmin = min(xmin, digital_box['coordinate'][0])
                ymin = min(ymin, digital_box['coordinate'][1])
                xmax = max(xmax, digital_box['coordinate'][2])
                ymax = max(ymax, digital_box['coordinate'][3])

        digital_list_coordinate = (xmin, ymin, xmax, ymax)

        direction = utils.decide_choice_m_left_top(digital_list_coordinate, choice_m_dict_box)
        if int(direction):
            choice_m_dict['direction'] = direction
            direction_list.append(direction)
            if direction == '180':  # 数字垂直排列
                std_num_length = choice_m_dict['rows']
                choice_option = a_z[:choice_m_dict['cols']].replace('', ',')[1:-1]
                default_points = [-1] * std_num_length
                choice_m_dict.update({'option': choice_option, 'default_points': default_points})

                sorted(digital_list, key=lambda k: k.get('coordinate')[1])
                choice_ymin = choice_m_dict['bounding_box']['ymin']
                single_height = choice_m_dict['single_height']
                mean_interval = ((choice_m_dict['bounding_box']['ymax'] - choice_m_dict['bounding_box']['ymin'])
                                 - single_height * std_num_length) / (std_num_length - 1)
                spilt_index = [choice_ymin - mean_interval / 2 + (single_height + mean_interval) * ele for ele in
                               range(std_num_length + 1)]

                number_list = [-1] * std_num_length
                number_location = [(-1, -1, -1, -1)] * std_num_length
                for i in range(0, len(spilt_index) - 1):
                    start = spilt_index[i]
                    end = spilt_index[i + 1]
                    number_location[i] = (xmin, start, xmax, end)
                    for digital_coordinate in digital_list:
                        middle_y = (digital_coordinate['coordinate'][3] - digital_coordinate['coordinate'][1]) / 2 + \
                                   digital_coordinate['coordinate'][1]
                        middle_x = (digital_coordinate['coordinate'][2] - digital_coordinate['coordinate'][0]) / 2 + \
                                   digital_coordinate['coordinate'][0]
                        if (start <= middle_y <= end
                                and
                                middle_x < choice_m_dict['bounding_box']['xmin']):  # 数字在choice_m外侧
                            number_list[i] = int(digital_coordinate['char'])
                            number_location[i] = digital_coordinate['coordinate']

                number_list = _infer_number(number_list)
                choice_m_dict['number'] = _infer_number(number_list)
                # choice_m_dict['number'] = [{'number': number,
                #                             'location': {'xmin': xi, 'ymin': yi, 'xmax': xm, 'ymax': ym}}
                #                            for number in number_list
                #                            for (xi, yi, xm, ym) in number_location]

            if direction == '90':  # 数字水平排列
                std_num_length = choice_m_dict['cols']
                choice_option = a_z[:std_num_length].replace('', ',')[1:-1]
                default_points = [-1] * std_num_length
                choice_m_dict.update({'option': choice_option, 'default_points': default_points})

                sorted(digital_list, key=lambda k: k.get('coordinate')[0])
                choice_xmin = choice_m_dict['bounding_box']['ymin']
                single_width = choice_m_dict['single_width']
                mean_interval = ((choice_m_dict['bounding_box']['xmax'] - choice_m_dict['bounding_box']['xmin'])
                                 - single_width * std_num_length) / (std_num_length - 1)
                spilt_index = [choice_xmin - mean_interval / 2 + (single_width + mean_interval) * ele for ele in
                               range(std_num_length)]
                number_list = [-1] * std_num_length
                number_location = [(-1, -1, -1, -1)] * std_num_length
                for i in range(0, len(spilt_index) - 1):
                    start = spilt_index[i]
                    end = spilt_index[i + 1]
                    number_location[i] = (start, ymin, end, ymax)
                    for digital_coordinate in digital_list:
                        middle_y = (digital_coordinate['coordinate'][3] - digital_coordinate['coordinate'][1]) / 2 + \
                                   digital_coordinate['coordinate'][1]
                        middle_x = (digital_coordinate['coordinate'][2] - digital_coordinate['coordinate'][0]) / 2 + \
                                   digital_coordinate['coordinate'][0]
                        if start <= middle_x <= end and middle_y < choice_m_dict['bounding_box']['ymin']:
                            number_list[i] = int(digital_coordinate['char'])
                            number_location[i] = digital_coordinate['coordinate']

                number_list = _infer_number(number_list)
                choice_m_dict['number'] = _infer_number(number_list)

                # choice_m_dict['number'] = [{'number': number,
                #                             'location': {'xmin': xi, 'ymin': yi, 'xmax': xm, 'ymax': ym}}
                #                            for number in number_list
                #                            for (xi, yi, xm, ym) in number_location]

        else:
            choice_m_dict['direction'] = '0'
            choice_m_dict['number'] = [-1]
            choice_m_dict['default_points'] = [-1]

    count180 = ','.join(direction_list).count('180')
    count90 = ','.join(direction_list).count('90')

    infer_direction = ['180', '90'][[count180, count90].index(max(count180, count90))]
    for ele in choice_m_dict_list:
        if ele['direction'] != '0':
            ele.update({'direction': infer_direction})

    # tree = ET.parse(xml_path)  # xml tree
    # for index_num, choice_box in enumerate(choice_m_dict_list):
    #     if len(choice_box['bounding_box']) > 0:
    #         abcd = choice_box['bounding_box']
    #         number = str(choice_box['number'])
    #         name = '{}_{}*{}_{}_{}'.format('choice_m', choice_box['rows'],
    #                                        choice_box['cols'], choice_box['direction'],
    #                                        number)
    #         tree = utils.create_xml(name, tree,
    #                                 abcd['xmin'], abcd['ymin'],
    #                                 abcd['xmax'], abcd['ymax'])
    #
    # tree.write(xml_path)
    return choice_m_dict_list


def _infer_number(number_list):
    if -1 not in number_list or sum(number_list) == -1 * len(number_list):
        return number_list
    else:
        for n_index in range(0, len(number_list) - 1):
            if n_index == 0:
                if number_list[n_index] != -1:

                    if len(number_list) > 1 and number_list[n_index + 1] == -1:
                        number_list[n_index + 1] = number_list[n_index] + 1

            if number_list[n_index] != -1:
                if number_list[n_index - 1] == -1:
                    number_list[n_index - 1] = number_list[n_index] - 1
                if number_list[n_index + 1] == -1:
                    number_list[n_index + 1] = number_list[n_index] + 1
        return _infer_number(number_list)