import cv2
import numpy as np


def deskew_image(image_array):
    # 转换为灰度图像
    gray = cv2.cvtColor(image_array, cv2.COLOR_BGR2GRAY)

    # 使用边缘检测找到倾斜角度
    edges = cv2.Canny(gray, 150, 255, apertureSize=3)
    cv2.imshow("edges", edges)
    lines = cv2.HoughLines(edges, 1, np.pi / 180, 200)
    print(lines.shape)
    print(lines)
    theta = lines[0][0][1] if lines is not None else 0

    # 旋转图像以校正倾斜
    (h, w) = image_array.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, np.degrees(theta), 1.0)
    rotated = cv2.warpAffine(image_array, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)

    # 裁剪校正后的图像
    gray = cv2.cvtColor(rotated, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    coords = np.column_stack(np.where(thresh > 0))
    angle = cv2.minAreaRect(coords)[-1]
    if angle < -45:
        angle = -(90 + angle)
    else:
        angle = -angle
    (h, w) = rotated.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(rotated, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)

    # 校准透视变换
    gray = cv2.cvtColor(rotated, cv2.COLOR_BGR2GRAY)
    thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cnt = contours[0]
    rect = cv2.minAreaRect(cnt)
    box = cv2.boxPoints(rect)
    box = np.intp(box)
    width = int(rect[1][0])
    height = int(rect[1][1])
    src_pts = box.astype("float32")
    dst_pts = np.array([[0, height - 1],
                        [0, 0],
                        [width - 1, 0],
                        [width - 1, height - 1]], dtype="float32")
    M = cv2.getPerspectiveTransform(src_pts, dst_pts)
    warped = cv2.warpPerspective(rotated, M, (width, height))

    return warped


def main():
    # 示例用法
    path = "img/18.jpg"
    img = cv2.imread(path)
    corrected_array = deskew_image(img)
    cv2.imshow("name", corrected_array)
    cv2.waitKey(0)


if __name__ == '__main__':
    main()