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

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

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

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

在左边的项目中右键某个文件,也可以打开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')

看下运行效果:

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

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

上面写了两个数字,识别输出正确!

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

上一篇:pyqt5和qt-designer结合


下一篇:CodeForces 719A. Vitya in the Countryside