FORMAT_CONVERSION_MAXCOMPUTE/format_convert/table_correct.py

import math
import cmath
import traceback

import numpy as np
import cv2


# 图片旋转
from format_convert.judge_platform import get_platform


def rotate_bound(image, angle):
    try:
        # 获取宽高
        (h, w) = image.shape[:2]
        (cX, cY) = (w // 2, h // 2)

        # 提取旋转矩阵 sin cos
        M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
        cos = np.abs(M[0, 0])
        sin = np.abs(M[0, 1])

        # 计算图像的新边界尺寸
        nW = int((h * sin) + (w * cos))
        #     nH = int((h * cos) + (w * sin))
        nH = h

        # 调整旋转矩阵
        M[0, 2] += (nW / 2) - cX
        M[1, 2] += (nH / 2) - cY

        return cv2.warpAffine(image, M, (nW, nH),flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
    except Exception as e:
        print("rotate_bound", e)


# 获取图片旋转角度
def get_minAreaRect(image):
    try:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        ret, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        if get_platform() == "Windows":
            cv2.imshow("binary", binary)
            cv2.waitKey(0)
        contours, hier = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # 绘制矩形
        cv2.drawContours(image, contours, 0, (255, 0, 255), 3)
        if get_platform() == "Windows":
            cv2.imshow("image", image)
            cv2.waitKey(0)

        for c in contours:  #遍历轮廓
            rect = cv2.minAreaRect(c)  #生成最小外接矩形
            box_ = cv2.boxPoints(rect)
            h = abs(box_[3, 1] - box_[1, 1])
            w = abs(box_[3, 0] - box_[1, 0])
            print("宽，高", w, h)
            # 只保留需要的轮廓
            if h > 3000 or w > 2200:
                continue
            if h < 2500 or w < 1500:
                continue
            # 计算最小面积矩形的坐标
            box = cv2.boxPoints(rect)
            # 将坐标规范化为整数
            box = np.int0(box)
            # 获取矩形相对于水平面的角度
            angle = rect[2]
            if angle > 0:
                if abs(angle) > 45:
                    angle = 90 - abs(angle)
            else:
                if abs(angle) > 45:
                    angle = (90 - abs(angle))

        print("轮廓数量", len(contours))
        return image, box, angle

    except Exception as e:
        print("get_minAreaRect", traceback.print_exc())
        return [-1], [-1], [-1]


def get_table_line(image):
    # 二值化
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    binary = cv2.adaptiveThreshold(~gray, 255,
                                   cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 203, -10)
    # cv2.imshow("cell", binary)
    # cv2.waitKey(0)

    # 轮廓
    kernel_row = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], np.float32)
    kernel_col = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], np.float32)
    # Sobel
    # kernel_row = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], np.float32)
    # kernel_col = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], np.float32)

    # binary = cv2.filter2D(binary, -1, kernel=kernel)
    binary_row = cv2.filter2D(binary, -1, kernel=kernel_row)
    binary_col = cv2.filter2D(binary, -1, kernel=kernel_col)
    # cv2.imshow("custom_blur_demo", binary)
    # cv2.waitKey(0)

    # rows, cols = binary.shape
    # scale = 20
    # # 识别横线
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (cols//scale, 1))
    # erodedcol = cv2.erode(binary_row, kernel, iterations=1)
    # cv2.imshow("Eroded Image", erodedcol)
    # cv2.waitKey(0)
    # dilatedcol = cv2.dilate(erodedcol, kernel, iterations=3)
    # cv2.imshow("Dilated Image", dilatedcol)
    # cv2.waitKey(0)
    #
    # # 识别竖线
    # scale = 20
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, rows//scale))
    # erodedrow = cv2.erode(binary_col, kernel, iterations=1)
    # cv2.imshow("Eroded Image", erodedrow)
    # cv2.waitKey(0)
    # dilatedrow = cv2.dilate(erodedrow, kernel, iterations=3)
    # cv2.imshow("Dilated Image", dilatedrow)
    # cv2.waitKey(0)
    #
    # # 标识表格
    # merge = cv2.add(dilatedcol, dilatedrow)
    # cv2.imshow("add Image", merge)
    # cv2.imwrite('table_outline.jpg', merge)
    # cv2.waitKey(0)

    return binary_row


def detect_line(img):
    try:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        edges = cv2.Canny(gray, 100, 1000)
        # cv2.imshow("edges", edges)
        # cv2.waitKey(0)

        lines = cv2.HoughLines(edges, 1, np.pi/180, 200)
        if lines is None or len(lines) == 0:
            return 0

        angle_list = []
        for i in range(len(lines)):
            # 第一个元素是距离rho
            rho = lines[i][0][0]
            # 第二个元素是角度theta
            theta = lines[i][0][1]
            a = np.cos(theta)
            b = np.sin(theta)
            x0 = a*rho
            y0 = b*rho
            x1 = float(x0 + 1000*(-b))
            y1 = float(y0 + 1000*a)
            x2 = float(x0 - 1000*(-b))
            y2 = float(y0 - 1000*a)

            if x2-x1 == 0 or y2-y1 == 0 or x2-x1 <= 0.01 or y2-y1 <= 0.01:
                continue

            k = -(y2-y1) / (x2-x1)
            if abs(k) <= 5:
                # h = math.atan(k)
                # angle = math.degrees(h)
                angle = np.arctan(k) * (180/math.pi)
                # 调整
                angle_list.append(angle)
                # print(angle)
            # cv2.line(img, (x1, y1), (x2, y2), (0, 0, 255), 2)

        angle_sum = 0
        for a in angle_list:
            angle_sum += a
        angle_avg = angle_sum/len(angle_list)
        print("angle_avg", angle_avg)

        return angle_avg
    except Exception as e:
        print("detect_line", e)


def get_rotated_image_old(image, output_path):
    try:
        angle = detect_line(image)
        if angle != 0:
            rotated = rotate_bound(image, angle)
            # cv2.putText(rotated, "angle: {:.3f} ".format(angle),
            #             (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

            # show the output image
            # cv2.imshow("imput", image)
            # cv2.waitKey(0)
            # cv2.imshow("output", rotated)
            # cv2.waitKey(0)
            print("[INFO] angle: {:.3f}".format(angle))
            cv2.imwrite(output_path, rotated)
            return True
        else:
            print("angle", angle)
            return False
    except Exception as e:
        print("get_rotated_image", e)
        return False


def rotate_bound(image, angle):
    # 获取宽高
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)

    # 提取旋转矩阵 sin cos
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    # 计算图像的新边界尺寸
    nW = int((h * sin) + (w * cos))
    #     nH = int((h * cos) + (w * sin))
    nH = h

    # 调整旋转矩阵
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY

    return cv2.warpAffine(image, M, (nW, nH),flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)


# 获取图片旋转角度
def get_minAreaRect(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    gray = cv2.bitwise_not(gray)
    thresh = cv2.threshold(gray, 0, 255,
                           cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
    coords = np.column_stack(np.where(thresh > 0))
    return cv2.minAreaRect(coords)


def get_rotated_image(image, output_path):
    try:
        image_temp, box, angle = get_minAreaRect(image)
        if angle == [-1]:
            return [-1]

        # angle = get_minAreaRect(image)[-1]
        if abs(angle) >= 15:
            angle = 0
        rotated = rotate_bound(image, angle)

        # cv2.putText(rotated, "angle: {:.2f} ".format(angle),
        #             (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        # show the output image
        # print("[INFO] angle: {:.3f}".format(angle))
        if not angle:
            cv2.imwrite(output_path, rotated)
        # cv2.imshow("input", image)
        # cv2.waitKey(0)
        # cv2.imshow("output", rotated)
        # cv2.waitKey(0)
        return True
    except cv2.error:
        traceback.print_exc()
        return [-3]
    except Exception as e:
        traceback.print_exc()
        return [-1]


if __name__ == '__main__':
    temp_path = "temp/complex/8.png"
    # temp_path = "temp/92d885dcb77411eb914300163e0ae709/92d8ef5eb77411eb914300163e0ae709_pdf_page1.png"
    image = cv2.imread(temp_path)
    get_rotated_image(image, temp_path)

    # test()