程序员求职经验分享与学习资料整理平台

网站首页 > 文章精选 正文

Pytorch利用CNN卷积神经网络进行多数字(0-9)识别

balukai 2025-01-01 17:56:40 文章精选 8 ℃

往期的文章中,我们分享了CNN卷积神经网络的数字识别,字母识别,本期我们分享一下如何使用CNN卷积神经网络进行多数字的识别,毕竟我们使用 过程中,主要是用来多数字识别的,比如手机号码识别,身份证识别,发票识别等等

pytorch利用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的数据图片是一维 28*28的图片,所以第一层的输入(1,28,28),高度为1,设置输出16通道,使用5*5的卷积核对图片进行卷积运算,每步移动一格,为了避免图片尺寸变化,设置pading为2,则经过第一层卷积就输出(16,28,28)数据格式

再经过relu与maxpooling (使用2*2卷积核)数据输出(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()全连接层进行数据的全连接数据结构(32*7*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)

    return Borderlist

getContours函数主要是进行图片中数字区域的区分,把每个数字的坐标检测出来,这样就可以 把每个数字进行CNN卷积神经网络的识别,进而实现多个数字识别的目的

step 4:图片中数字的识别

经过以上2个子函数,我们便可以得到图片中每个数字的具体坐标位置,我们从图片中扣出此数字送进神经网络,这样多此遍历后就可以得到多个数字的识别了

Borderlist = []  # 不同的轮廓图像及坐标
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]

        cv2.rectangle(imgCopy, (x, y), (x + s, y + s), color=(0, 255, 0), thickness=1)
        cv2.putText(imgCopy, result, (x + s // 2 - 5, y + s // 2 - 5), cv2.FONT_HERSHEY_COMPLEX, 1.5, (0, 0, 255), 2)
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
最近发表
标签列表