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的乐趣与应用场景。