5050L-MTF/line.py

import cv2
import numpy as np
import sys
from utils_io import imread_unicode, imwrite_unicode

def correct_rotation_hough(image: np.ndarray, edge_method: str = 'canny', edge_save_path: str = None) -> np.ndarray:
    """
    基于霍夫变换的旋转矫正。
    步骤：灰度化->边缘检测->霍夫直线检测->计算角度->仿射旋转校正。
    参数：
        image: 输入图像（BGR或灰度）
        edge_method: 边缘检测方法（'canny'/'sobel'）
    返回：
        旋转校正后的图像
    """
    # 灰度化
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image.copy()

    # 边缘检测
    if edge_method == 'canny':
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    elif edge_method == 'sobel':
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        mag = cv2.magnitude(sobelx, sobely)
        edges = cv2.convertScaleAbs(mag)
        _, edges = cv2.threshold(edges, 50, 255, cv2.THRESH_BINARY)
    else:
        raise ValueError('Unsupported edge method')
    # 保存边缘检测图片
    if edge_save_path:
        imwrite_unicode(edge_save_path, edges)

    # 霍夫直线检测
    lines = cv2.HoughLines(edges, 1, np.pi / 180, 100)
    if lines is None or len(lines) == 0:
        # 没有检测到直线，返回原图
        print("Detected angle for rotation correction:")
        return image.copy()

    # 计算所有直线的角度（以垂直为0度，逆时针为正）
    angles = []
    for line in lines:
        rho, theta = line[0]
        angle = theta * 180 / np.pi
        # 只考虑接近水平的直线（排除垂直线）
        if 45 < angle < 135:
            angles.append(angle - 90)  # 以90度为基准
    if len(angles) == 0:
        # 没有合适的直线，返回原图
        return image.copy()
    # 取中位数角度作为旋转角度
    rotate_angle = float(np.median(angles))
    print("Detected angle for rotation correction:", rotate_angle)
    print(rotate_angle)
    # 仿射旋转校正
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, float(rotate_angle), 1.0)
    rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
    return rotated


def correct_rotation_fourier(image: np.ndarray, edge_save_path: str = None) -> np.ndarray:
    """
    基于傅里叶变换的旋转矫正。
    步骤：傅里叶变换->频谱分析->计算角度->仿射旋转校正。
    参数：
        image: 输入图像（BGR或灰度）
    返回：
        旋转校正后的图像
    """
    # 灰度化
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image.copy()
    # 傅里叶变换
    f = np.fft.fft2(gray)
    fshift = np.fft.fftshift(f)
    magnitude_spectrum = np.log(np.abs(fshift) + 1)
    # 频谱归一化到0-255
    mag_img = cv2.normalize(magnitude_spectrum, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    # 对频谱图做二值化
    _, mag_bin = cv2.threshold(mag_img, 180, 255, cv2.THRESH_BINARY)
    # 保存频谱二值图
    if edge_save_path:
        imwrite_unicode(edge_save_path, mag_bin)
    # 霍夫变换检测主方向
    lines = cv2.HoughLines(mag_bin, 1, np.pi / 180, 100)
    angles = []
    if lines is not None:
        for line in lines:
            rho, theta = line[0]
            angle = theta * 180 / np.pi
            # 只考虑接近水平的方向
            if 45 < angle < 135:
                angles.append(angle - 90)
    if len(angles) == 0:
        # 没有检测到主方向，返回原图
        print("Detected angle for rotation correction (Fourier): None")
        return image.copy()
    rotate_angle = float(np.median(angles))
    print("Detected angle for rotation correction (Fourier):", rotate_angle)
    # 仿射旋转校正
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, float(rotate_angle), 1.0)
    rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
    return rotated


def correct_rotation_contour(image: np.ndarray, edge_method: str = 'canny', edge_save_path: str = None) -> np.ndarray:
    """
    基于最小外接矩形的轮廓旋转矫正。
    步骤：灰度化->边缘检测->轮廓检测->最小外接矩形->仿射旋转校正。
    参数：
        image: 输入图像（BGR或灰度）
        edge_method: 边缘检测方法（'canny'/'sobel'）
        edge_save_path: 边缘图保存路径
    返回：
        旋转校正后的图像
    """
    # 灰度化
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image.copy()
    # 边缘检测
    if edge_method == 'canny':
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    elif edge_method == 'sobel':
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        mag = cv2.magnitude(sobelx, sobely)
        edges = cv2.convertScaleAbs(mag)
        _, edges = cv2.threshold(edges, 50, 255, cv2.THRESH_BINARY)
    else:
        raise ValueError('Unsupported edge method')
    # 保存边缘检测图片
    if edge_save_path:
        imwrite_unicode(edge_save_path, edges)
    # 轮廓检测
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        print("未检测到轮廓，返回原图")
        return image.copy()
    # 找最大轮廓
    largest_contour = max(contours, key=cv2.contourArea)
    # 最小外接矩形
    rect = cv2.minAreaRect(largest_contour)
    box = cv2.boxPoints(rect)
    box = box.astype(int)
    angle = rect[2]
    # OpenCV的minAreaRect角度定义：
    # 水平为0，逆时针为负，范围[-90,0)。如果宽高互换，角度会有跳变。
    # 统一为正角度（逆时针为正）
    base_angle = angle if rect[1][0] < rect[1][1] else angle + 90
    rotate_angle = base_angle % 90
    if rotate_angle > 45:
        rotate_angle -= 90
    rotate_angle = float(rotate_angle) + 180
    print("Detected angle for rotation correction (Contour):", rotate_angle-180)
    # 仿射旋转校正
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, float(rotate_angle), 1.0)
    rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
    return rotated


def interface_corr(input):
    import argparse
    import os
    parser = argparse.ArgumentParser(description='基于霍夫变换/傅里叶变换/轮廓最小外接矩形的旋转矫正')
    parser.add_argument('--input', '-i', type=str, default=r"input\25mm\1.bmp", help='输入图像路径')
    parser.add_argument('--output', '-o', type=str, default=r"output\result_hough.bmp", help='输出图像路径')
    parser.add_argument('--edge', type=str, default='canny', choices=['canny', 'sobel'], help='边缘检测方法（仅霍夫法/轮廓法用）')
    parser.add_argument('--method', type=str, default='contour', choices=['hough', 'fourier', 'contour'], help='旋转矫正方法')
    args = parser.parse_args()

    img = imread_unicode(input)
    if img is None:
        print(f'无法读取图像: {input}')
        sys.exit(1)
    # 边缘图保存路径
    out_dir = os.path.dirname(args.output)
    if out_dir and not os.path.exists(out_dir):
        os.makedirs(out_dir)
    if args.method == 'hough':
        edge_save_path = os.path.join(out_dir, 'result_edges_hough.bmp')
        result = correct_rotation_hough(img, edge_method=args.edge, edge_save_path=edge_save_path)
    elif args.method == 'fourier':
        edge_save_path = os.path.join(out_dir, 'result_edges_fourier.bmp')
        result = correct_rotation_fourier(img, edge_save_path=edge_save_path)
    elif args.method == 'contour':
        edge_save_path = os.path.join(out_dir, 'result_edges_contour.bmp')
        result = correct_rotation_contour(img, edge_method=args.edge, edge_save_path=edge_save_path)
    else:
        print('未知方法')
        sys.exit(1)
    imwrite_unicode(args.output, result)
    print(f'旋转校正完成，结果已保存到: {args.output}')
    return args.output
    # interface_pro(r"D:\WJY\step1demo\code\FOC\output\result_hough.bmp")

if __name__ == '__main__':
    input = r"output\res_roi.png"
    interface_corr(input)