163 lines
5.8 KiB
Python
163 lines
5.8 KiB
Python
import cv2
|
|
import numpy as np
|
|
import itertools
|
|
import os
|
|
import time # 引入 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 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
|
|
|
|
# 双边滤波平滑处理
|
|
#smooth_mask = cv2.bilateralFilter(alpha_channel, d=9, sigmaColor=75, sigmaSpace=75)
|
|
|
|
# 高斯模糊平滑处理
|
|
#blurred_mask = cv2.GaussianBlur(smooth_mask, (3, 3), sigmaX=0, sigmaY=0)
|
|
|
|
# 二值化处理：转换为边框 mask
|
|
_, border_mask = cv2.threshold(alpha_channel, 1, 255, cv2.THRESH_BINARY)
|
|
|
|
# 找到边框上的轮廓
|
|
contours, _ = cv2.findContours(border_mask, 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]
|
|
|
|
quadrilaterals = list(itertools.combinations(corner_points, 4))
|
|
|
|
def calculate_weight(points):
|
|
|
|
contour = np.array(points, dtype=np.int32)
|
|
area = cv2.contourArea(contour)
|
|
|
|
if area/non_transparent_area < 0.8:
|
|
return -1
|
|
|
|
if not cv2.isContourConvex(contour):
|
|
return -1
|
|
|
|
# print(area/non_transparent_area)
|
|
|
|
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)
|
|
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)
|
|
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.4 + (area / non_transparent_area) * 0.6 # 使用非透明面积归一化
|
|
|
|
best_quad = None
|
|
max_weight = -1
|
|
for quad in quadrilaterals:
|
|
# 对最优四边形的点进行排序
|
|
quad = order_points(quad)
|
|
|
|
weight = calculate_weight(quad)
|
|
if weight > max_weight:
|
|
max_weight = weight
|
|
best_quad = quad
|
|
|
|
|
|
|
|
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}")
|