105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
# @Time : 2024/2/18 14:45
|
||
|
|
# @Author : len
|
||
|
|
# @File : find_qr_block_2.py
|
||
|
|
# @Software: PyCharm
|
||
|
|
# @Comment :
|
||
|
|
import os
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
from sklearn.cluster import KMeans
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
|
||
|
|
def qr_KMeans(image):
|
||
|
|
|
||
|
|
# 计算中心区域的坐标
|
||
|
|
height, width = image.shape[:2]
|
||
|
|
center_region = image[height // 4: 3 * height // 4, width // 4: 3 * width // 4]
|
||
|
|
|
||
|
|
# 将中心区域的颜色空间从BGR转换到HSV
|
||
|
|
center_hsv = cv2.cvtColor(center_region, cv2.COLOR_BGR2HSV)
|
||
|
|
|
||
|
|
# 将中心区域图像数据重塑为二维数组 pixels (center_height*center_width, num_channels)
|
||
|
|
center_pixels = center_hsv.reshape((-1, 3))
|
||
|
|
|
||
|
|
# 应用K-means聚类来找到中心区域的两种主要颜色
|
||
|
|
kmeans = KMeans(n_clusters=2, random_state=0, n_init=10).fit(center_pixels)
|
||
|
|
|
||
|
|
dominant_colors = kmeans.cluster_centers_.astype(int)
|
||
|
|
# for color in dominant_colors:
|
||
|
|
# print(color)
|
||
|
|
|
||
|
|
# 确定前景色的颜色范围
|
||
|
|
# 这里我们选择距离黑色(或白色)最远的颜色作为前景色
|
||
|
|
foreground_color = max(dominant_colors, key=lambda c: cv2.norm(c, cv2.NORM_L2))
|
||
|
|
|
||
|
|
# 转换整个图像到HSV空间
|
||
|
|
hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||
|
|
|
||
|
|
fig, axs = plt.subplots(1, len(dominant_colors) + 2, figsize=(20, 6))
|
||
|
|
# 在第一个位置显示原图
|
||
|
|
axs[0].imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||
|
|
axs[0].set_title('Original Image with Contours')
|
||
|
|
axs[0].axis('off')
|
||
|
|
axs[0].set_xticks([])
|
||
|
|
axs[0].set_yticks([])
|
||
|
|
|
||
|
|
for idx, color in enumerate(dominant_colors):
|
||
|
|
# 创建HSV颜色范围
|
||
|
|
lower_val = np.clip(color - np.array([10, 30, 30]), 0, 255)
|
||
|
|
upper_val = np.clip(color + np.array([10, 255, 255]), 0, 255)
|
||
|
|
|
||
|
|
# 创建颜色范围内的掩码
|
||
|
|
mask = cv2.inRange(hsv_image, lower_val, upper_val)
|
||
|
|
|
||
|
|
# 寻找轮廓
|
||
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
|
|
||
|
|
# 在原图上画出轮廓
|
||
|
|
if idx == 0:
|
||
|
|
cv2.drawContours(image, contours, -1, (255, 0, 0), 2)
|
||
|
|
else:
|
||
|
|
cv2.drawContours(image, contours, -1, (0, 255, 0), 2)
|
||
|
|
mask_inv = cv2.bitwise_not(mask) # 黑白取反
|
||
|
|
|
||
|
|
# kernel = np.ones((5, 5), np.uint8)
|
||
|
|
# erosion = cv2.erode(mask_inv, kernel, iterations=1)
|
||
|
|
# dilation = cv2.dilate(erosion, kernel, iterations=1)
|
||
|
|
# erosion = cv2.erode(dilation, kernel, iterations=1)
|
||
|
|
|
||
|
|
# 显示掩码
|
||
|
|
idx += 2
|
||
|
|
axs[idx].imshow(mask_inv, cmap='gray')
|
||
|
|
axs[idx].set_title('Mask of Color: {}'.format(color))
|
||
|
|
axs[idx].axis('off')
|
||
|
|
axs[idx].set_xticks([])
|
||
|
|
axs[idx].set_yticks([])
|
||
|
|
# 在第一个位置显示原图
|
||
|
|
axs[1].imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||
|
|
axs[1].set_title('Original Image with Contours')
|
||
|
|
axs[1].axis('off')
|
||
|
|
axs[1].set_xticks([])
|
||
|
|
axs[1].set_yticks([])
|
||
|
|
# 显示结果
|
||
|
|
plt.tight_layout()
|
||
|
|
plt.show()
|
||
|
|
|
||
|
|
# input_path = r"./android_img/qr_kmeans_img"
|
||
|
|
# image_files = [f for f in os.listdir(input_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
|
||
|
|
#
|
||
|
|
# for image_path in tqdm(image_files):
|
||
|
|
# image_full_path = os.path.join(input_path, image_path)
|
||
|
|
# image = cv2.imread(image_full_path)
|
||
|
|
# if image is not None:
|
||
|
|
# qr_KMeans(image)
|
||
|
|
# else:
|
||
|
|
# print(f"Failed to load image: {image_path}")
|
||
|
|
|
||
|
|
|
||
|
|
input_path = r"./data/img_1.png"
|
||
|
|
# input_path = r"./data/img.png"
|
||
|
|
image = cv2.imread(input_path)
|
||
|
|
qr_KMeans(image)
|