exam-segment-django/segment/sheet_resolve/analysis/solve/optional_solve.py

# @Author  : lightXu
# @File    : optional_solve.py
# @Time    : 2019/9/17 0017 下午 13:18
import cv2
import re
from segment.sheet_resolve.tools.brain_api import get_ocr_text_and_coordinate
from segment.sheet_resolve.analysis.sheet.choice_infer import find_digital


def rgb2binary(im):
    gray_img = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    _ret, thresh_img = cv2.threshold(gray_img, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    return thresh_img


def find_contours(left, top, image, ex_x=30, ex_y=1):
    threshed = rgb2binary(image)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ex_x, ex_y))  # 膨胀系数
    # morphed = cv2.morphologyEx(threshed, cv2.MORPH_CLOSE, kernel)
    morphed = cv2.dilate(threshed, kernel, iterations=1)

    (major, minor, _) = cv2.__version__.split(".")
    contours = cv2.findContours(morphed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cnts = contours[0] if int(major) > 3 else contours[1]

    cnt = sorted(cnts, key=cv2.contourArea)
    boxes = []
    for ele in cnt:
        x, y, w, h = cv2.boundingRect(ele)
        x_mid = x + w//2
        y_mid = y + h//2
        xmax = x+w
        ymax = y+h
        boxes.append((x, y, xmax, ymax, x_mid,y_mid))

    return boxes


def point_in_polygon(point, polygon):
    xmin, ymin, xmax, ymax = polygon[0], polygon[1], polygon[2], polygon[3]
    if xmin <= point[0] <= xmax and ymin <= point[1] <= ymax:
        return True
    else:
        return False


def resolve_optional_choice(l_, t_, direction, image):
    ocr_res = get_ocr_text_and_coordinate(image)
    digital_list, chars_list, d_mean_height, d_mean_width = find_digital(ocr_res, 0, 0,)
    if not digital_list:
        numbers = [501, 502]
        h, w = image.shape
        optional_choice_dict = {'class_name': 'optional_choice',
                                'rows': 1, 'cols': 2,
                                'single_width': w,
                                'single_height': h,
                                'direction': direction,
                                'bounding_box': {'xmin': l_,
                                                 'ymin': t_,
                                                 'xmax': l_ + w,
                                                 'ymax': t_ + h},
                                'number': numbers}
    else:
        numbers = sorted([ele['digital'] for ele in digital_list])

        contours = find_contours(l_, t_, image, d_mean_width*2, d_mean_height)

        res_region = []
        for contour in contours:
            for num in digital_list:
                loc = num['loc']
                num_center = (loc[-2], loc[-1])
                if point_in_polygon(num_center, contour):
                    if num_center[0] < contour[-2] - d_mean_width//3:
                        # 数字在box的左侧
                        region = (loc[2]+1, loc[1]-1, contour[2]-3, loc[3]+1)
                    else:
                        # 数字在box中间, 数字在box左侧
                        region = (contour[0], loc[1] - 1, contour[2], loc[3] + 1)
                    res_region.append(region)

        l, t, r, b = 9999, 9999, 0, 0
        sum_w, sum_h = 0, 0
        for region in res_region:
            l, t, r, b = min(l, region[0]), min(t, region[1]), max(r, region[2]), max(b, region[3])

            width, height = region[2] - region[0], region[3] - region[1]
            sum_w += width
            sum_h += height

        mean_w, mean_h = sum_w//len(numbers), sum_h//len(numbers)

        if direction == 180:
            rows, cols = 1, max(2, len(numbers))
            if len(numbers) < 2:
                numbers.append(501)
        else:
            rows, cols = max(2, len(numbers)), 1
            if len(numbers) < 2:
                numbers.append(501)

        optional_choice_dict = {'class_name': 'optional_choice',
                                'rows': rows, 'cols': cols,
                                'single_width': mean_w,
                                'single_height': mean_h,
                                'direction': direction,
                                'bounding_box': {'xmin': l_ + l,
                                                 'ymin': t_ + t,
                                                 'xmax': l_ + r,
                                                 'ymax': t_ + b},
                                'number': numbers}

    # print(optional_choice_dict)
    return optional_choice_dict