代码收藏家技术教程 2024-09-24

Python OpenCV实现基于颜色提取的印章精准分割技术

利用Python实现了一个图像处理功能，即批量提取图像中的印章区域；使用了颜色聚类的方法来提取颜色。

本代码也发布到了github，欢迎大家试用（如果帮助，请star一下）：

GitHub – AICVHub/seal_seg_opencv: pyhon opencv seal segpyhon opencv seal seg. Contribute to AICVHub/seal_seg_opencv development by creating an account on GitHub.https://github.com/AICVHub/seal_seg_opencv.git

1. 方法

2. 代码

3. 运行结果

1. 方法

K-means 聚类 (kmeans_color_quantization 函数): 通过将图像从BGR颜色空间转换到HSV颜色空间，并使用K-means聚类算法来找到图像中的主要颜色。K-means算法是一种无监督学习算法，用于将数据点划分为K个簇，使得每个数据点与其所属簇的中心距离之和最小。

颜色阈值 (extract_seal_with_kmeans 函数): 根据聚类得到的颜色中心，创建HSV颜色空间中的上下界限，以确定与印章颜色相似的像素范围。通过调整hue_threshold参数，可以控制颜色相似度的敏感度。

掩码生成: 使用cv2.inRange函数根据颜色阈值生成掩码，这些掩码用于识别图像中可能包含印章的区域。

掩码膨胀 (cv2.dilate 函数): 对生成的掩码进行膨胀操作，以填补印章区域中的小空洞或连接断裂的部分。

图像合成: 使用掩码提取印章区域，并将其与原图结合在一个透明的背景上。这里使用了PIL库（Python Imaging Library）来处理图像的透明度和合成。

批量处理 (batch_extract_seals_with_dominant_color 函数): 遍历指定文件夹中的所有图像，对每个图像执行印章提取和合成操作，并将结果保存到另一个文件夹。

2. 代码

import cv2
import numpy as np
import os
from sklearn.cluster import KMeans
from PIL import Image


def kmeans_color_quantization(hsv_image, k=2):
    """
    使用 K-means 聚类找到图像中的主要颜色。
    """
    # 重新调整图像大小以加速聚类过程
    resized_image = cv2.resize(hsv_image, (0, 0), fx=0.5, fy=0.5)

    # 将图像数据转换为二维数组，每行是一个像素的 HSV 值
    reshaped_image = resized_image.reshape(-1, 3)

    # 应用 K-means 聚类
    kmeans = KMeans(n_clusters=k)
    kmeans.fit(reshaped_image)

    # 返回聚类中心
    return kmeans.cluster_centers_


def extract_seal_with_kmeans(image_path, output_path, k=4, hue_threshold=15):
    image = cv2.imread(image_path)
    if image is None:
        print(f"Image not found at {image_path}")
        return

    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    dominant_colors = kmeans_color_quantization(hsv_image, k=k)

    masks = []
    for color in dominant_colors:
        lower_color = np.array([color[0] - hue_threshold, 100, 100])
        upper_color = np.array([color[0] + hue_threshold, 255, 255])
        mask = cv2.inRange(hsv_image, lower_color, upper_color)
        masks.append(mask)
    combined_mask = cv2.bitwise_or(masks[0], masks[1]) if len(masks) > 1 else masks[0]

    # 膨胀掩码
    kernel = np.ones((3, 3), np.uint8)
    final_mask = cv2.dilate(combined_mask, kernel, iterations=1)

    # 提取印章区域
    seal = cv2.bitwise_and(image, image, mask=final_mask)

    # 将OpenCV图像转换为PIL图像
    original_image = Image.open(image_path).convert('RGBA')
    seal_pil = Image.fromarray(cv2.cvtColor(seal, cv2.COLOR_BGR2RGBA)).convert('RGBA')

    # 计算新图像的宽度，原图宽度加上印章区域宽度
    original_width, original_height = original_image.size
    seal_width, seal_height = seal_pil.size
    new_width = original_width + seal_width
    new_height = max(original_height, seal_height)

    # 创建一个全透明的图像，用于左右拼接
    transparent_background = Image.new('RGBA', (new_width, new_height), (0, 0, 0, 0))

    # 将原图粘贴到透明背景上
    transparent_background.paste(original_image, (0, 0))

    # 将印章区域粘贴到透明背景的右侧
    transparent_background.paste(seal_pil, (original_width, 0), seal_pil)

    # 保存结果
    transparent_background.save(output_path)

    print(f"Original image and extracted seal combined with transparent background saved to {output_path}")


def batch_extract_seals_with_dominant_color(input_dir, output_dir):
    """
    批量处理文件夹中的所有图像，使用主要颜色进行印章提取。

    参数:
    - input_dir: 包含输入图像的文件夹路径。
    - output_dir: 输出图像的保存文件夹路径。
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for filename in os.listdir(input_dir):
        image_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, os.path.splitext(filename)[0]+'.png')
        extract_seal_with_kmeans(image_path, output_path)


# 使用示例
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Batch extract seals with dominant color from images.")
    parser.add_argument("--input_dir", default='/data/projects/Matting/modnet_demo/seals',
                        help="Path to the folder containing input images.")
    parser.add_argument("--output_dir", default='/data/projects/Matting/modnet_demo/output_seals_01',
                        help="Path to the folder for saving output images.")
    args = parser.parse_args()

    batch_extract_seals_with_dominant_color(args.input_dir, args.output_dir)

关于其中一些方法的说明：

KMeans：

KMeans 是一个在机器学习库 scikit-learn 中实现的 K-means 聚类算法的类。K-means 聚类是一种无监督学习算法，它的目标是将数据点划分为 K 个簇，使得每个数据点与其所属簇的中心（即聚类中心）之间的距离之和最小。这种算法在许多领域都有应用，包括图像处理、市场细分、天文数据分析等。

以下是 KMeans 类的一些关键点：

初始化 (__init__ 方法): 当你创建一个 KMeans 实例时，你需要指定 n_clusters 参数，它代表你想要划分的簇的数量。在这个例子中，k 就是这个参数的值。

拟合数据 (fit 方法): 调用 fit 方法时，你需要传递一个数据集，通常是二维数组的形式，其中每行代表一个数据点，每列代表一个特征。KMeans 会根据这些数据点初始化簇中心，并迭代地调整簇中心的位置，直到满足停止条件（如达到一定的迭代次数或簇中心的变化小于某个阈值）。

聚类中心 (cluster_centers_ 属性): 聚类完成后，KMeans 实例的 cluster_centers_ 属性会包含 K 个聚类中心的坐标。每个聚类中心是一个向量，其维度与输入数据的特征维度相同。

标签分配 (labels_ 属性): 除了聚类中心，KMeans 还会为每个输入数据点分配一个标签，表示它属于哪个簇。这个标签分配是通过计算数据点与各个聚类中心之间的距离来决定的，数据点会被分配到最近的聚类中心所代表的簇。

优化: K-means 算法在初始化簇中心时可能会陷入局部最优，因此可能需要多次运行以获得更好的结果。KMeans 类提供了 init 参数来指定初始化方法，例如 init='k-means++'，这是一种更智能的初始化方法，可以提高找到全局最优解的概率。

灵活性: KMeans 类还提供了其他参数来调整算法的行为，比如 max_iter（最大迭代次数）、tol（收敛阈值）等。

cv2.inRange：

在OpenCV中，cv2.inRange 函数用于创建一个掩码（mask），该掩码将图像中与指定颜色范围匹配的像素设置为白色（即像素值为255），而将不匹配的像素设置为黑色（即像素值为0）。这个函数在图像处理中非常有用，特别是在需要根据颜色来分割图像时。