jigsaw/angle2.py

import cv2
import numpy as np
import itertools
import os
import time

# 设置输入输出文件夹路径
input_folder = "output10"  # 输入文件夹
output_folder = "output12"  # 输出文件夹

# 创建输出文件夹（如果不存在）
os.makedirs(output_folder, exist_ok=True)
start_time = time.time()

def order_points(points):
    """
    将四边形的四个点按顺时针顺序排列。
    :param points: 四边形的四个点 [(x1, y1), (x2, y2), (x3, y3), (x4, y4)]
    :return: 排序后的点
    """
    # 计算质心
    center = np.mean(points, axis=0)

    # 计算每个点相对于质心的角度
    angles = [np.arctan2(point[1] - center[1], point[0] - center[0]) for point in points]

    # 按角度排序
    sorted_points = [point for _, point in sorted(zip(angles, points))]
    return sorted_points

def order_counters(contours):
    counterclockwise_contours = []
    for contour in contours:
        # 计算轮廓的签名面积
        area = cv2.contourArea(contour, oriented=True)

        # 如果面积是负值，说明轮廓是逆时针
        if area > 0:
            # 将轮廓反转为逆时针
            contour = contour[::-1]

        # 添加到新的列表中
        counterclockwise_contours.append(contour)
    return counterclockwise_contours

def calculate_weight(points, non_transparent_area):
    """
    计算四边形的权重，用于选择最优四边形。
    :param points: 四边形的四个点 [(x1, y1), (x2, y2), (x3, y3), (x4, y4)]
    :param non_transparent_area: 非透明区域的面积
    :return: 四边形的权重（-1 表示无效四边形）
    """
    contour = np.array(points, dtype=np.int32)
    area = cv2.contourArea(contour)

    # 筛选条件 1: 四边形的面积比例
    if area / non_transparent_area < 0.8:
        return -1

    # 计算边长
    edges = [np.linalg.norm(np.array(points[i]) - np.array(points[(i + 1) % 4])) for i in range(4)]
    max_edge, min_edge = max(edges), min(edges)

    # 筛选条件 2: 边长相似性
    if max_edge - min_edge > max_edge / 7:
        return -1

    edge_similarity = 1 - (max_edge - min_edge) / max_edge

    # 计算角度
    angles = []
    for i in range(4):
        v1 = np.array(points[(i + 1) % 4]) - np.array(points[i])
        v2 = np.array(points[(i + 3) % 4]) - np.array(points[i])
        cosine_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
        angle = np.degrees(np.arccos(np.clip(cosine_angle, -1.0, 1.0)))
        angles.append(angle)

    # 筛选条件 3: 角度相似性
    if any(angle < 80 or angle > 100 for angle in angles):
        return -1

    angle_similarity = 1 - sum(abs(angle - 90) for angle in angles) / 360

    # 计算最终权重
    square_similarity = edge_similarity * angle_similarity
    return square_similarity * 0.5 + (area / non_transparent_area) * 0.5  # 使用非透明面积归一化

def find_best_quadrilateral(corner_points, non_transparent_area):
    """
    从角点中筛选最优四边形。
    :param corner_points: 检测到的角点 [(x1, y1), (x2, y2), ...]
    :param non_transparent_area: 非透明区域的面积
    :return: 最优四边形的点集和最大权重
    """
    quadrilaterals = list(itertools.combinations(corner_points, 4))
    best_quad = None
    max_weight = -1

    for quad in quadrilaterals:
        # 对四边形的点进行顺时针排序
        quad = order_points(quad)

        weight = calculate_weight(quad, non_transparent_area)
        if weight > max_weight:
            max_weight = weight
            best_quad = quad

    return best_quad, max_weight

def process_image(image_path, save_path):
    # 加载图片
    image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    if image is None:
        print(f"无法加载图片: {image_path}")
        return

    img_height, img_width = image.shape[:2]
    alpha_channel = image[:, :, 3]

    non_transparent_area = np.count_nonzero(alpha_channel > 0)
    if non_transparent_area == 0:  # 如果图片完全透明，跳过
        print(f"图片完全透明: {image_path}")
        return

    # 找到边框上的轮廓
    contours, _ = cv2.findContours(alpha_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 创建一个空白的边框图片
    border_image = np.zeros_like(alpha_channel)
    cv2.drawContours(border_image, contours, -1, color=255, thickness=1)

    # 角点检测 (最多 16 个角点)
    max_corners = 14
    corners = cv2.goodFeaturesToTrack(border_image, maxCorners=max_corners, qualityLevel=0.01, minDistance=15, blockSize=5)
    if corners is not None:
        corners = np.intp(corners)
        corner_points = [tuple(corner.ravel()) for corner in corners]

        # 使用独立函数筛选最优四边形
        best_quad, max_weight = find_best_quadrilateral(corner_points, non_transparent_area)

        if best_quad is None:
            print(f"未找到有效四边形: {image_path}")

            output_image = cv2.cvtColor(image[:, :, :3], cv2.COLOR_BGR2RGB)
            for point in corner_points:
                cv2.circle(output_image, point, radius=3, color=(255, 0, 0), thickness=-1)

            # 用蓝色标记边框
            cv2.drawContours(output_image, contours, -1, color=(0, 0, 255), thickness=1)  # 蓝色 (BGR: 255, 0, 0)

            cv2.imwrite(save_path, cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR))
            return

        area = cv2.contourArea(np.array(best_quad, dtype=np.int32))
        print(image_path+"\t\t"+str(area / non_transparent_area))

        # 绘制最优四边形和角点
        output_image = cv2.cvtColor(image[:, :, :3], cv2.COLOR_BGR2RGB)
        for point in corner_points:
            cv2.circle(output_image, point, radius=3, color=(255, 0, 0), thickness=-1)

        cv2.polylines(output_image, [np.array(best_quad, dtype=np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)

        # 用蓝色标记边框
        cv2.drawContours(output_image, contours, -1, color=(0, 0, 255), thickness=1)  # 蓝色 (BGR: 255, 0, 0)

        cv2.imwrite(save_path, cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR))
    else:
        print(f"未检测到角点: {image_path}")

# 遍历输入文件夹中的图片
for filename in os.listdir(input_folder):
    if filename.lower().endswith((".png", ".jpg", ".jpeg")):  # 支持常见图片格式
        input_path = os.path.join(input_folder, filename)
        output_path = os.path.join(output_folder, filename)
        process_image(input_path, output_path)

end_time = time.time()
print(f"处理图片的时间: {end_time - start_time:.4f} 秒")
print(f"所有图片已处理完成，结果保存在: {output_folder}")