AI手写输入法 - pytorch从入门到入道(二)

本章承接上一篇的手写数字识别，利用训练好的模型，结合pyqt画板，实现简易手写输入法，为"hello world"例子增添乐趣。

pyqt是开发图形界面的框架，可以百度查找相关资料了解安装及基础方法，我搭建的环境是pycharm+pyqt5+qtdesigner，配置好之后的界面长这样：

在左边的项目中右键某个文件，也可以打开qt菜单

具体怎么画界面不展开了，直接看下代码：

 # coding: utf-8
 from PyQt5.QtWidgets import *
 from PyQt5.QtGui import *
 from PyQt5.QtCore import *
 import sys
 sys.path.append(r'../ml/torch')
 from digit_recog import Net
 import torch
 import os
 import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image

 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 net = Net().to(device)
 # 加载参数
 nn_state = torch.load(os.path.join('../ml/torch/model/', 'net.pth'))
 # 参数加载到指定模型
 net.load_state_dict(nn_state)
 net.eval()

 def predict(img):
     # 读取图片并重设尺寸
     image = Image.open(img).resize((28, 28))
     # 灰度图
     gray_image = image.convert('L')
     # plt.imshow(gray_image)
     # plt.show()
     # 图片数据处理
     im_data = np.array(gray_image)
     im_data = torch.from_numpy(im_data).float()
     im_data = im_data.view(1, 1, 28, 28)
     # 神经网络运算
     outputs = net(im_data)
     # 取最大预测值
     _, pred = torch.max(outputs, 1)
     return pred.item()

 class SimpleDrawingBoard(QWidget):
     win = ''
     wins = []

     @classmethod
     def showWin(cls):
         # 聚焦到已有窗口
         if not cls.win:
             cls.win = cls()
             cls.win.show()
         else:
             cls.win.activateWindow()

     def __init__(self, parent=None):
         super(SimpleDrawingBoard, self).__init__(parent)

         self.setWindowTitle(u"手写数字识别")
         self.setWindowFlags(Qt.WindowStaysOnTopHint)
         self.size = (400, 350)
         self.resize(*self.size)
         self.setWindowFlag(Qt.FramelessWindowHint)  # 隐藏边框
         # self.setWindowOpacity(0.9)  # 设置窗口透明度
         # self.setAttribute(Qt.WA_TranslucentBackground)  # 设置窗口背景透明

         self.canvasSize = (280, 350)
         self.sizeOffset = [a - b for a, b in zip(self.size, self.canvasSize)]
         self.canvas = QPixmap(*self.canvasSize)
         self.canvas.fill(Qt.black)
         self.tempCanvas = QPixmap()
         self.lastPoint = QPoint()
         self.endPoint = QPoint()
         self.isDrawing = False
         self.penSize = 15

         self.initUI()

     def initUI(self):
         self.penSizeLabel = QLabel(u'画笔粗细')
         self.penSizeSpinBox = QSpinBox()
         self.penSizeSpinBox.setValue(self.penSize)
         self.penSizeSpinBox.valueChanged.connect(self.penSizeSpinBox_valueChanged)
         self.penSizeSpinBox.setFixedWidth(80)

         self.clearButton = QPushButton(u'清空')
         self.clearButton.setFixedWidth(80)
         self.clearButton.clicked.connect(self.clearPainter)

         self.closeButton = QPushButton(u'关闭')
         self.closeButton.setFixedWidth(80)
         self.closeButton.clicked.connect(self.close)

         self.inputLabel = QLabel(self)
         self.inputLabel.setFixedSize(80, 200)
         self.inputLabel.setAutoFillBackground(True)
         self.inputLabel.setAlignment(Qt.AlignCenter)
         self.inputLabel.setStyleSheet('''QLabel{background:#F76677;border-radius:5px;font-size:60px;font-weight:bolder;}''')

         mainLayout = QVBoxLayout(self)

         toolbarLayout = QGridLayout()
         # toolbarLayout.setSpacing(20)
         toolbarLayout.addWidget(self.penSizeLabel, 0, 0, 1, 1)
         toolbarLayout.addWidget(self.penSizeSpinBox, 1, 0, 1, 1)
         toolbarLayout.addWidget(self.clearButton, 2, 0, 1, 1)
         toolbarLayout.addWidget(self.closeButton, 3, 0, 1, 1)
         toolbarLayout.addWidget(self.inputLabel, 4, 0, 1, 1)

         toolbarLayout.setAlignment(Qt.AlignLeft)

         mainLayout.addLayout(toolbarLayout)
         mainLayout.addStretch(1)

     def penSizeSpinBox_valueChanged(self):
         # 设置画笔粗细
         self.penSize = self.penSizeSpinBox.value()

     def paintEvent(self, event):
         pp = QPainter(self.canvas)
         pen = QPen(QColor(255, 255, 255), self.penSize)
         pp.setPen(pen)
         if self.lastPoint != self.endPoint:
             pp.drawLine(self.lastPoint - QPoint(*self.sizeOffset), self.endPoint - QPoint(*self.sizeOffset))
         painter = QPainter(self)
         painter.drawPixmap(self.sizeOffset[0], self.sizeOffset[1], self.canvas)
         self.lastPoint = self.endPoint

     def clearPainter(self):
         print('clear...')
         self.canvas.fill(Qt.black)
         painter = QPainter(self)
         painter.drawPixmap(self.sizeOffset[0], self.sizeOffset[1], self.canvas)
         self.lastPoint = self.endPoint
         self.update()
         self.inputLabel.clear()

     def mousePressEvent(self, event):
         # 按下左键
         if event.button() == Qt.LeftButton:
             self.lastPoint = event.pos()
             self.endPoint = self.lastPoint
             self.isDrawing = True

     def mouseMoveEvent(self, event):
         if self.isDrawing:
             self.update()
             self.endPoint = event.pos()

     def mouseReleaseEvent(self, event):
         if event.button() == Qt.LeftButton:
             self.isDrawing = False
             self.endPoint = event.pos()
             self.update()
             self.canvas.toImage().save('input.png')
             input = predict('input.png')
             self.inputLabel.setText(str(input))
             print('你输入的是{}'.format(input))

 if __name__ == '__main__':
     app = QApplication.instance()
     if not app:
         app = QApplication(sys.argv)
     SimpleDrawingBoard.showWin()
     app.exec_()

上面引入前一章训练好的模型，位于不同的文件夹内，需要加上这一行代码：

sys.path.append(r'../ml/torch')

看下运行效果：

上面写了两个数字，识别输出正确！

helloworld例子比较枯燥，通过动手参与与AI交互增强信心乐趣，信心是一步步建立起来的，而大的突破亦是如此，后面会持续围绕简单的例子，深入发掘AI的乐趣与应用场景。