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

前言

上文我们使用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...

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

image-20240905112544119

Modules中也选择上面对应的SDK

image-20240905113330645

添加依赖

项目目录下初始化

1
pipenv install

安装OpenCV

1
pipenv install opencv-python==4.5.4.60

KNN

步骤

  1. 准备数据:准备特征和标签数据,并训练 KNN 模型。

  2. 保存和加载模型:将模型保存到文件中,然后从文件中加载模型。

  3. 测试数据:定义需要进行预测的测试数据。

  4. 进行预测:使用 findNearest() 方法进行预测。

    该方法返回四个值:

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

模型训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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)。这意味着标签不能是字符串或者其他非整数类型。

模型预测

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 加载模型
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长度的数组。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
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文件

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
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)