import cv2
import numpy as np
import math
import os

def process_image(image_path):
    """
    读取灰度图像，找到灰度值大于阈值的亮线端点，测量端点连线与水平线的夹角，
    并将图片反向旋转该角度使得端点连线水平。

    """
    try:
        # 读取灰度图像
        gray_image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if gray_image is None:
            print(f"无法读取图像: {image_path}")
            return None

        # 找到灰度值大于90的像素点
        bright_pixels = np.where(gray_image > 90)
        if not bright_pixels[0].size:
            print(f"{image_path}: 未找到灰度值大于90的亮线。")
            return cv2.cvtColor(gray_image, cv2.COLOR_GRAY2BGR)  # 返回原始灰度图的BGR版本

        bright_y = bright_pixels[0]
        bright_x = bright_pixels[1]

        # 创建一个与原始灰度图大小相同的BGR图像，用于旋转
        color_image = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2BGR)

        # 找到亮线的边界框
        min_x = np.min(bright_x)
        max_x = np.max(bright_x)
        min_y = np.min(bright_y)
        max_y = np.max(bright_y)

        # 提取亮线区域的像素坐标
        bright_line_region_y = bright_y[(bright_x >= min_x) & (bright_x <= max_x) & (bright_y >= min_y) & (bright_y <= max_y)]
        bright_line_region_x = bright_x[(bright_x >= min_x) & (bright_x <= max_x) & (bright_y >= min_y) & (bright_y <= max_y)]

        
        if len(bright_line_region_x) > 0:
            # 寻找最左边和最右边的点作为可能的端点
            leftmost_index = np.argmin(bright_line_region_x)
            rightmost_index = np.argmax(bright_line_region_x)
            endpoint1 = (bright_line_region_x[leftmost_index], bright_line_region_y[leftmost_index])
            endpoint2 = (bright_line_region_x[rightmost_index], bright_line_region_y[rightmost_index])            

            # 计算端点连线与水平线的夹角
            delta_x = endpoint2[0] - endpoint1[0]
            delta_y = endpoint2[1] - endpoint1[1]
            angle_radians = math.atan2(delta_y, delta_x)
            angle_degrees = math.degrees(angle_radians)
            print ("倾斜角是",angle_degrees)

            # 获取图像的中心点
            height, width = gray_image.shape[:2]
            center = (width // 2, height // 2)

            # 构建旋转矩阵 (反向旋转)
            rotation_matrix = cv2.getRotationMatrix2D(center, angle_degrees, 1)

            # 执行图像旋转
            rotated_image = cv2.warpAffine(color_image, rotation_matrix, (width, height), borderMode=cv2.BORDER_REPLICATE)

            return rotated_image

        elif len(bright_x) == 1:
            print(f"{image_path}: 只找到一个亮像素点，无法进行旋转。")
            return color_image
        else:
            print(f"{image_path}: 未找到足够的亮像素点来确定端点，无法进行旋转。")
            return color_image

    except Exception as e:
        print(f"{image_path}: 处理图像时发生错误: {e}")
        return None

if __name__ == "__main__":
    image_dir = "."  # 设置为当前目录，你可以修改为你的图片所在的目录
    image_files = [f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff'))]

    if not image_files:
        print(f"在目录 '{image_dir}' 中未找到任何图片文件。")

    for image_file in image_files:
        image_path = os.path.join(image_dir, image_file)
        print(f"正在处理: {image_path}")
        rotated_image = process_image(image_path)

        if rotated_image is not None:
            name, ext = os.path.splitext(image_file)
            rotated_filename = f"{name}-rotated{ext}"
            rotated_path = os.path.join(image_dir, rotated_filename)
            cv2.imwrite(rotated_path, rotated_image)
            print(f"已保存旋转后的图像: {rotated_path}")

    print("所有图片处理完成。")