Python使用OpenCV的KNN算法进行字母数字识别

发表于 2024-09-05 更新于 2025-03-25 分类于 opencv

前言

上文我们使用Microsoft.ML做了字母和数字的识别，效果还可以。

但是OpenCV的KNN在识别图像上效果会更好，但是C#上不支持KNN模型的训练，所以本文使用Python来训练。

阿里针对教育的OCR识别

https://ai.aliyun.com/ocr/edu?spm=5176.21213303.J_qCOwPWspKEuWcmp8qiZNQ.30.25ed2f3dcX4P0c&scm=20140722.S_product@@%E4%BA%91%E4%BA%A7%E5%93%81@@82233._.ID_product@@%E4%BA%91%E4%BA%A7%E5%93%81@@82233-RL_%E7%AD%94%E9%A2%98%E5%8D%A1%E8%AF%86%E5%88%AB-LOC_llm-OR_ser-V_3-RE_new3@@cardOld-P0_0

环境配置参考

https://www.psvmc.cn/article/2022-06-13-python-development-env.html

IDEA

IDEA中报错配置

idea no python interpreter configured for the module

File => Project Structure => PlatForm Settings => SDKs

点击 + => Add Python SDK...

我这里使用的虚拟环境，选择已存在的环境，可以根据名字判断环境。

Modules中也选择上面对应的SDK

添加依赖

项目目录下初始化

1	pipenv install

安装OpenCV

1	pipenv install opencv-python==4.5.4.60

KNN

步骤

准备数据：准备特征和标签数据，并训练 KNN 模型。
保存和加载模型：将模型保存到文件中，然后从文件中加载模型。
测试数据：定义需要进行预测的测试数据。
进行预测：使用 findNearest() 方法进行预测。

该方法返回四个值：
- ret：返回值（通常是 KNN 结果的准确性）
- results：预测结果的标签
- neighbors：最近邻的标签
- dist：每个测试样本到其最近邻的距离
获取预测标签：results 是一个二维数组，你可以通过 flatten() 方法将其转化为一维数组，以便直接获取预测的标签值。

模型训练

import cv2
import numpy as np

label_map = {'cat': 0, 'dog': 1}
# 反向映射，用于将整数转换回标签
inverse_label_map = {v: k for k, v in label_map.items()}
# 模型文件
model_filename = 'knn_model.xml'

# 示例数据
features = np.array([[1, 2], [2, 3], [3, 4], [4, 5]], dtype=np.float32)
labels = np.array([0, 1, 0, 1], dtype=np.int32)

# 创建 KNN 模型并训练
knn = cv2.ml.KNearest_create()
knn.train(features, cv2.ml.ROW_SAMPLE, labels)

# 保存模型
knn.save(model_filename)

注意

在 OpenCV 的 KNN 模型中：

特征数据必须是 numpy.ndarray 类型。数据的类型应该是 np.float32。这是因为 KNN 算法要求特征数据要以浮点数格式表示。

labels 参数必须是整数类型 (int32)。这意味着标签不能是字符串或者其他非整数类型。

模型预测

# 加载模型
loaded_knn = cv2.ml.KNearest_load(model_filename)

# 测试数据
test_data = np.array([[2.5, 3.5], [2.5, 3.5]], dtype=np.float32)

# 进行预测
ret, results, neighbors, dist = loaded_knn.findNearest(test_data, k=3)

# 获取预测的标签
predicted_labels = results.flatten()

print("预测结果的标签:", predicted_labels)
# 将整数标签转换回原始字符串标签
predicted_labels_str = [inverse_label_map[label] for label in predicted_labels]
print("预测结果的标签（字符串形式）:", predicted_labels_str)

结果是这样的

预测结果的标签: [0. 0.]
预测结果的标签（字符串形式）: [‘cat’, ‘cat’]

图片识别工具类

我们可以把参与训练的图片都转为20*20的大小，这样输入的数据都是400长度的数组。

import csv
import json
import os
import cv2
import numpy as np
from pathlib import Path
from utils.cv_common_utils import CvCommonUtils

script_path = Path(__file__).resolve()
project_path = str(script_path.parent.parent)
model_filepath = os.path.join(project_path, "ml-resource", "knn_model.xml")
json_filepath = os.path.join(project_path, "ml-resource", 'knn_label_map.json')


class KnnUtils:

    @staticmethod
    def model_train(tsv_path):
        path_list = []
        word_list = []
        # 打开 TSV 文件
        with open(tsv_path, 'r', newline='', encoding='utf-8') as file:
            reader = csv.reader(file, delimiter='\t')
            # 读取数据
            for row in reader:
                if len(row) >= 2:
                    path_list.append(row[0])
                    word_list.append(row[1])
        word_set = set(word_list)
        label_map = {}

        for index, word in enumerate(word_set):
            label_map[word] = index

        label_list = []
        for word in word_list:
            label_list.append(label_map[word])
        # 反向映射，用于将整数转换回标签
        inverse_label_map = {v: k for k, v in label_map.items()}
        # 将字典写入 JSON 文件
        with open(json_filepath, 'w', encoding='utf-8') as file:
            json.dump(inverse_label_map, file, ensure_ascii=False, indent=4)
        img_list = []

        for path in path_list:
            img = CvCommonUtils.read_img(path, cv2.THRESH_BINARY)
            img2 = CvCommonUtils.get_mat32(img)
            img_list.append(img2.flatten())

        # 示例数据
        features = np.array(img_list, dtype=np.float32)
        labels = np.array(label_list, dtype=np.int32)

        # 创建 KNN 模型并训练
        knn = cv2.ml.KNearest_create()
        knn.train(features, cv2.ml.ROW_SAMPLE, labels)

        # 保存模型
        knn.save(model_filepath)
        print("KNN训练完成")

    loaded_knn = None
    inverse_label_map = None

    @staticmethod
    def model_load():
        if KnnUtils.loaded_knn is None:
            # 加载模型
            KnnUtils.loaded_knn = cv2.ml.KNearest_load(model_filepath)
            # 从 JSON 文件中读取字典
            with open(json_filepath, 'r', encoding='utf-8') as file:
                KnnUtils.inverse_label_map = json.load(file)

    @staticmethod
    def recognition_img(img):
        KnnUtils.model_load()

        if KnnUtils.loaded_knn is None:
            print("模型加载失败")
            return
        if not CvCommonUtils.is_binary_image(img):
            img = CvCommonUtils.binary(img)
            print("非二值化图片")
        img2 = CvCommonUtils.get_mat32(img)

        img_data = img2.flatten()
        img_data2 = list(map(float, img_data))
        # print(img_data)

        img_list = []
        img_list.append(img_data2)

        # 测试数据
        test_data = np.array(img_list, dtype=np.float32)
        # 进行预测
        ret, results, neighbors, dist = KnnUtils.loaded_knn.findNearest(test_data, k=3)
        # 获取预测的标签
        predicted_labels = results.flatten()
        if len(predicted_labels) > 0:
            predicted_label = predicted_labels[0]
            # 将整数标签转换回原始字符串标签
            key = str(int(predicted_label))
            return KnnUtils.inverse_label_map[key]
        return ""

读取TSV文件

import csv

tsv_path = "D:\\Project\\csharp\\z-exam-card-recognize\\ml-resource\\ml\\Word.tsv"

path_list = []
word_list = []
# 打开 TSV 文件
with open(tsv_path, 'r', newline='', encoding='utf-8') as file:
    reader = csv.reader(file, delimiter='\t')

    # 读取数据
    for row in reader:
        if (len(row) >= 2):
            path_list.append(row[0])
            word_list.append(row[1])

print(path_list)
print(word_list)