126 lines
3.9 KiB
Python
126 lines
3.9 KiB
Python
# -*- coding:utf-8 -*-
|
|
# @Author len
|
|
# @Create 2023/10/30 15:14
|
|
|
|
|
|
import math
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
def detect_shape(contour, epsilon_factor=0.04):
|
|
epsilon = epsilon_factor * cv2.arcLength(contour, True)
|
|
approx = cv2.approxPolyDP(contour, epsilon, True)
|
|
|
|
if len(approx) == 3:
|
|
return "三角形"
|
|
elif len(approx) == 4:
|
|
# # 计算四边形的内角
|
|
# angles = []
|
|
# for i in range(4):
|
|
# pt1 = approx[i]
|
|
# pt2 = approx[(i + 1) % 4]
|
|
# pt3 = approx[(i + 2) % 4]
|
|
# angle = np.degrees(np.arctan2(pt3[0][1] - pt2[0][1], pt3[0][0] - pt2[0][0]) -
|
|
# np.arctan2(pt1[0][1] - pt2[0][1], pt1[0][0] - pt2[0][0]))
|
|
# angles.append(angle)
|
|
# for angle in angles:
|
|
# print(angle)
|
|
# # 检查角度是否接近90度
|
|
# if all(85 <= abs(angle) <= 95 for angle in angles):
|
|
return "矩形"
|
|
# else:
|
|
# return "菱形"
|
|
elif len(approx) == 10:
|
|
return "五角星"
|
|
else:
|
|
area = cv2.contourArea(contour)
|
|
perimeter = cv2.arcLength(contour, True)
|
|
radius = perimeter / (2 * np.pi)
|
|
circle_area = np.pi * (radius ** 2)
|
|
|
|
if abs(circle_area - area) < area * 0.2:
|
|
return "圆形"
|
|
else:
|
|
return "未知形状"
|
|
|
|
def process_image(image_path, epsilon_factor=0.04):
|
|
img = cv2.imread(image_path)
|
|
# img = equalize_image(img)
|
|
h, w, _ = img.shape
|
|
|
|
center_color = img[h // 2, w // 2]
|
|
|
|
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
|
center_color_hsv = cv2.cvtColor(np.uint8([[center_color]]), cv2.COLOR_BGR2HSV)[0][0]
|
|
|
|
lower_bound = np.array([center_color_hsv[0] - 10, 50, 50])
|
|
upper_bound = np.array([center_color_hsv[0] + 10, 255, 255])
|
|
mask = cv2.inRange(hsv, lower_bound, upper_bound)
|
|
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
if not contours:
|
|
return img, "No contours found!"
|
|
|
|
largest_contour = max(contours, key=cv2.contourArea)
|
|
|
|
shape = detect_shape(largest_contour, epsilon_factor)
|
|
cv2.drawContours(img, [largest_contour], -1, (0, 0, 0), 2)
|
|
|
|
return img, shape
|
|
|
|
# 均衡化图像
|
|
def equalize_image(img):
|
|
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
|
hsv[:,:,2] = cv2.equalizeHist(hsv[:,:,2])
|
|
img_equalized = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
|
|
return img_equalized
|
|
|
|
|
|
def process_directory(directory_path, epsilon_factor=0.04):
|
|
filenames = []
|
|
images = [] # 用于存放处理过的图片
|
|
labels = [] # 用于存放识别结果
|
|
|
|
for file_name in os.listdir(directory_path):
|
|
file_path = os.path.join(directory_path, file_name)
|
|
if os.path.isfile(file_path) and file_name.lower().endswith(('.png', '.jpg', '.jpeg')):
|
|
img, shape = process_image(file_path, epsilon_factor)
|
|
images.append(img)
|
|
labels.append(shape)
|
|
|
|
# 保存文件名
|
|
filenames.append(file_name)
|
|
|
|
num_images = len(images)
|
|
num_cols = 6 # 每行显示3张图片
|
|
num_rows = math.ceil(num_images / num_cols) # 计算需要的行数
|
|
|
|
# 使用matplotlib显示图片
|
|
plt.figure(figsize=(15, 10))
|
|
|
|
# 设置字体以显示中文
|
|
plt.rcParams['font.sans-serif'] = ['SimHei']
|
|
plt.rcParams['axes.unicode_minus'] = False
|
|
|
|
for i, (img, label, filename) in enumerate(zip(images, labels, filenames)):
|
|
plt.subplot(num_rows, num_cols, i + 1)
|
|
plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
|
|
# plt.title(f"{filename}:{label}")
|
|
plt.title(f"{label}")
|
|
plt.axis('off')
|
|
|
|
# 添加大标题
|
|
# plt.suptitle(directory_path, fontsize=16, y=1.15)
|
|
|
|
plt.tight_layout()
|
|
plt.show()
|
|
|
|
|
|
# 使用方法
|
|
directory_path = r'D:\Code\Python\Embedded_game\polygon\exp16\crops\polygon' # 替换为你的图片文件夹路径
|
|
process_directory(directory_path)
|