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的樂趣與應用場景。