Pytorch利用CNN卷积神经网络进行多数字（0-9）识别
往期的文章中，我们分享了CNN卷积神经网络的数字识别，字母识别，本期我们分享一下如何使用CNN卷积神经网络进行多数字的识别，毕竟我们使用 过程中，主要是用来多数字识别的，比如手机号码识别，身份证识别，发票识别等等
利用pytorch CNN手写字母识别神经网络模型识别手写字母
step 1:神经网络的搭建
首先按照上期的代码搭建一下我们的CNN 卷积神经网络
import torch
import torch.nn as nn
from PIL import Image # 导入图片处理工具
import PIL.ImageOps
import numpy as np
from torchvision import transforms
import cv2
import matplotlib.pyplot as plt
# 定义神经网络
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Sequential( # input shape (1, 28, 28)
nn.Conv2d(
in_channels=1, # 输入通道数
out_channels=16, # 输出通道数
kernel_size=5, # 卷积核大小
stride=1, #卷积步数
padding=2, # 如果想要 con2d 出来的图片长宽没有变化,
# padding=(kernel_size-1)/2 当 stride=1
), # output shape (16, 28, 28)
nn.ReLU(), # activation
nn.MaxPool2d(kernel_size=2), # 在 2x2 空间里向下采样, output shape (16, 14, 14) )
self.conv2 = nn.Sequential( # input shape (16, 14, 14)
nn.Conv2d(16, 32, 5, 1, 2), # output shape (32, 14, 14)
nn.ReLU(), # activation
nn.MaxPool2d(2), # output shape (32, 7, 7) )
self.out = nn.Linear(32 * 7 * 7, 10) # 全连接层，0-9一共10个类
# 前向反馈
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size(0), -1) # 展平多维的卷积图成 (batch_size, 32 * 7 * 7)
output = self.out(x)
return output第一层，我们输入minist的数据集，minist的数据图片是一维
2828的图片，所以第一层的输入（1，28，28），高度为1，设置输出16通道，使用55的卷积核对图片进行卷积运算，每步移动一格，为了避免图片尺寸变化，设置pading为2，则经过第一层卷积就输出（16，28，28）数据格式
再经过relu与maxpooling （使用22卷积核）数据输出（16，14，14）
第二层卷积层是简化写法nn.Conv2d(16,
32, 5, 1, 2)的第一个参数为输入通道数in_channels=16，其第二个参数是输出通道数out_channels=32, #
n_filters（输出通道数），第三个参数为卷积核大小，第四个参数为卷积步数，最后一个为pading,此参数为保证输入输出图片的尺寸大小一致
self.conv2 = nn.Sequential( # input shape (16, 14, 14)
nn.Conv2d(16, 32, 5, 1, 2), # output shape (32, 14, 14)
nn.ReLU(), # activation
nn.MaxPool2d(2), # output shape (32, 7, 7)
)全连接层，最后使用nn.linear()全连接层进行数据的全连接数据结构（327*7,10）以上便是整个卷积神经网络的结构，
大致为：input-卷积-Relu-pooling-卷积
-Relu-pooling-linear-output
卷积神经网络建完后，使用forward（）前向传播神经网络进行输入图片的识别
step 2:图片预处理
# 预处理函数
def preProccessing(img):
imgGray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
imgBlur = cv2.GaussianBlur(imgGray, (5, 5), 1)
imgCanny = cv2.Canny(imgBlur, 200, 200)
imgDial = cv2.dilate(imgCanny, np.ones((5, 5)), iterations=2) # 膨胀操作
imgThres = cv2.erode(imgDial, np.ones((5, 5)), iterations=1) # 腐蚀操作
return imgThres
这里我们使用腐蚀，膨胀操作对图片进行一下预处理操作，方便神经网络的识别，当然，我们往期的字母数字识别也可以添加此预处理操作，方便神经网络进行预测，提高精度
step 3:图片轮廓检测获取每个数字的坐标位置
def getContours(img):
x, y, w, h, xx, yy, ss = 0, 0, 10, 10, 20, 20, 10 # 因为图像大小不能为0
imgGet = np.array([[], []]) # 不能为空
contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) # 检索外部轮廓
for cnt in contours:
area = cv2.contourArea(cnt)
if area > 800: # 面积大于800像素为封闭图形
cv2.drawContours(imgCopy, cnt, -1, (255, 0, 0), 3)
peri = cv2.arcLength(cnt, True) # 计算周长
approx = cv2.approxPolyDP(cnt, 0.02 * peri, True) # 计算有多少个拐角
x, y, w, h = cv2.boundingRect(approx) # 得到外接矩形的大小
a = (w + h) // 2
dd = abs((w - h) // 2) # 边框的差值
imgGet = imgProcess[y:y + h, x:x + w]
if w <= h: # 得到一个正方形框，边界往外扩充20像素,黑色边框
imgGet = cv2.copyMakeBorder(imgGet, 20, 20, 20 + dd, 20 + dd, cv2.BORDER_CONSTANT, value=[0, 0, 0])
xx = x - dd - 10
yy = y - 10
ss = h + 20
cv2.rectangle(imgCopy, (x - dd - 10, y - 10), (x + a + 10, y + h + 10), (0, 255, 0),
2) # 看看框选的效果，在imgCopy中
print(a + dd, h)
else: # 边界往外扩充20像素值
imgGet = cv2.copyMakeBorder(imgGet, 20 + dd, 20 + dd, 20, 20, cv2.BORDER_CONSTANT, value=[0, 0, 0])
xx = x - 10
yy = y - dd - 10
ss = w + 20
cv2.rectangle(imgCopy, (x - 10, y - dd - 10), (x + w + 10, y + a + 10), (0, 255, 0), 2)
print(a + dd, w)
Temptuple = (imgGet, xx, yy, ss) # 将图像及其坐标放在一个元组里面，然后再放进一个列表里面就可以访问了
Borderlist.append(Temptuple)
<span class="hljs-string">return</span> <span class="hljs-string">Borderlist</span></code></pre><p data-track="20">getContours函数主要是进行图片中数字区域的区分，把每个数字的坐标检测出来，这样就可以 把每个数字进行CNN卷积神经网络的识别，进而实现多个数字识别的目的</p><h1 class="pgc-h-arrow-right" data-track="23">step 4:图片中数字的识别</h1><p data-track="24">经过以上2个子函数，我们便可以得到图片中每个数字的具体坐标位置，我们从图片中扣出此数字送进神经网络，这样多此遍历后就可以得到多个数字的识别了</p><pre class="syl-page-code hljs nginx"><code><span class="hljs-attribute">Borderlist</span> = [] <span class="hljs-comment"># 不同的轮廓图像及坐标</span>
Resultlist = [] # 识别结果
img = cv2.imread('5.png')
imgCopy = img.copy()
imgProcess = preProccessing(img)
Borderlist = getContours(imgProcess)
train_transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Grayscale(),
transforms.Resize((28, 28)),
transforms.ToTensor(),
])
model = CNN()
model.load_state_dict(torch.load('./model/CNN_NO1.pk', map_location='cpu'))
model.eval()
index_to_class = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
首先我们imread读取一张需要识别的图片
图片中包括了多个数字，且是手写数字，然后我们使用其子函数preProccessing进行图片的预处理操作，然后使用getContours函数获取每个数字的坐标位置
train_transform是pytorch需要进行图片识别的一个转换器，里面包含了很多可以转换图片的操作，比如转换到torch变量，灰度变化，图片尺寸变化等等都可以使用transforms进行转换
然后我们加载神经网络数字识别模型
接下来，就可以遍历我们前面图片检测到的每个数字的坐标，里面坐标获取每个数字的图片
if len(Borderlist) != 0: # 不能为空
for (imgRes, x, y, s) in Borderlist:
cv2.imshow('imgCopy', imgRes)
cv2.waitKey(0)
img = train_transform(imgRes)
img = torch.unsqueeze(img, dim=0)
with torch.no_grad():
pre = model(img)
output = torch.squeeze(pre)
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print(index_to_class[predict_cla], predict[predict_cla].numpy())
result = index_to_class[predict_cla]
<span class="hljs-string">cv2.rectangle(imgCopy,</span> <span class="hljs-string">(x,</span> <span class="hljs-string">y),</span> <span class="hljs-string">(x</span> <span class="hljs-string">+</span> <span class="hljs-string">s,</span> <span class="hljs-string">y</span> <span class="hljs-string">+</span> <span class="hljs-string">s),</span> <span class="hljs-string">color=(0,</span> <span class="hljs-number">255</span><span class="hljs-string">,</span> <span class="hljs-number">0</span><span class="hljs-string">),</span> <span class="hljs-string">thickness=1)</span>
<span class="hljs-string">cv2.putText(imgCopy,</span> <span class="hljs-string">result,</span> <span class="hljs-string">(x</span> <span class="hljs-string">+</span> <span class="hljs-string">s</span> <span class="hljs-string">//</span> <span class="hljs-number">2</span> <span class="hljs-bullet">-</span> <span class="hljs-number">5</span><span class="hljs-string">,</span> <span class="hljs-string">y</span> <span class="hljs-string">+</span> <span class="hljs-string">s</span> <span class="hljs-string">//</span> <span class="hljs-number">2</span> <span class="hljs-bullet">-</span> <span class="hljs-number">5</span><span class="hljs-string">),</span> <span class="hljs-string">cv2.FONT_HERSHEY_COMPLEX,</span> <span class="hljs-number">1.5</span><span class="hljs-string">,</span> <span class="hljs-string">(0,</span> <span class="hljs-number">0</span><span class="hljs-string">,</span> <span class="hljs-number">255</span><span class="hljs-string">),</span> <span class="hljs-number">2</span><span class="hljs-string">)</span>
cv2.imshow('imgCopy', imgCopy)
cv2.waitKey(0)
这里我们可以imshow一下每个数字的图片，得到图片后，把图片送进CNN卷积神经网络进行图片数字的识别，得到结果后，我们可以利用数字的坐标把数字在原始图片上画出来，并添加上图片预测的结果
运行代码，我们可以 得到每个数字的坐标位置，并可以看到每个数字预测的结果与概率
当遍历完成所有图片数字后，我们就可以得到所有数字识别的结果了
#预测结果
0 1.0
2 1.0
1 0.99997175
3 1.0
8 0.99945825
4 1.0
9 0.9999397
5 1.0
6 0.9981457
7 0.9547404