图片处理
代码如下:
ImgProcess.py
import cv2
from cv2 import Mat
IMG_SPLIT_SIZE = 4
def img_split(input: Mat, size: int) -> list[Mat]:
"""
将图片分割成N份
:param input: 输入图片
:param size: 切合的份数
:return: 图片列表
"""
result = []
h, w, _ = input.shape
interval = int(w / size)
lx = 0
ly = h
for index in range(0, size):
split = input[int(ly * 0.2):int(ly * 0.7), # 去除上方20%，下方30%，每个字符左边去除10%，右边去除35%
lx + interval * index + int(interval * 0.1): lx + interval * index + int(interval * 0.65)]
result.append(split)
return result
def cal_average(mat: Mat) -> float:
"""
计算二值化后或者单通道图片像素的平均值
:param mat:
:return:
"""
h, w = mat.shape
total = 0
count = 0
for i in range(h):
for j in range(w):
total += mat[i, j]
count += 1
return total / count
def calculate_min_average(img: Mat) -> (Mat, str, float):
"""
获取图片所有通道中像素均值最小的通道
:param img: 输入图像
:return: 通道图像，通道名称，均值
"""
B, G, R = cv2.split(img)
avg_b = cal_average(B)
avg_g = cal_average(G)
avg_r = cal_average(R)
min_img = B
min_avg = avg_b
img_name = 'B'
if avg_g list[Mat]:
"""
处理图片
:param img: 输入图片
:return: 图片列表
"""
result = []
img_splits = img_split(img, IMG_SPLIT_SIZE)
for index, split in enumerate(img_splits):
one_img, _, _ = calculate_min_average(split)
one_img_blur = cv2.medianBlur(one_img, 3)
_, one_img_binary = cv2.threshold(one_img_blur, 156, 255, cv2.THRESH_BINARY_INV)
result.append(one_img_binary)
return result
模型训练
TrainModel.py
import cv2.ml
import numpy as np
from ImgProcess import process_check_code
import os
def train_and_save_model(check_code_dir: str, model_save_path: str):
"""
训练模型并保存
:param check_code_dir: 保存有验证码图片的目录，图片以验证码上的数字作为文件名。例如：图片中的数字为0541,那么文件名称就是0541.jpg
:param model_save_path: 模型保存地址
:return:
"""
samples = []
labels = []
for file in os.listdir(check_code_dir):
file_path = os.path.join(check_code_dir, file)
label_str, _ = os.path.splitext(file)
img = cv2.imread(file_path, cv2.IMREAD_COLOR)
mats = process_check_code(img)
for mat in mats:
samples.append(mat.flatten())
for label in label_str:
labels.append(int(label))
model = cv2.ml.KNearest.create()
model.train(np.array(samples, np.float32), cv2.ml.ROW_SAMPLE, np.array(labels, np.int32))
model.save(model_save_path)
图片识别
Recognize.py
import os
import cv2
import cv2.ml
from ImgProcess import process_check_code
from TrainModel import train_and_save_model
import numpy as np
def recognize(model: cv2.ml.KNearest, mat: cv2.Mat) -> str:
buffer = []
mats = process_check_code(mat)
for i, mat in enumerate(mats):
sample = np.array([mat.flatten()], np.float32)
ret, results, neighbours, distances = model.findNearest(sample, k=1)
buffer.append(str(int(ret)))
return "".join(buffer)
if __name__ == '__main__':
check_code_dir = "check_code"
model_path = "model"
if not os.path.exists(model_path):
train_and_save_model(check_code_dir, model_path)
model = cv2.ml.KNearest.load(model_path)
mat = cv2.imread("test.jpg")
code = recognize(model, mat)
print(f"Code = {code}")
附：
普通处理和按照通道处理对比图。右侧为均值最小的通道
