linux-基于tensorflow2.x的手写数字识别-基于MNIST数据集

数据集

数据集下载🔗MNIST
首先读取数据集, 并打印相关信息
包括

• 图像的数量, 形状
• 像素的最大, 最小值
• 以及看一下第一张图片

path = 'MNIST/mnist.npz'
with np.load(path, allow_pickle=True) as f:
    x_train, y_train = f['x_train'], f['y_train']
    x_test, y_test = f['x_test'], f['y_test']

print(f'dataset info: shape: {x_train.shape}, {y_train.shape}')
print(f'dataset info: max: {x_train.max()}')
print(f'dataset info: min: {x_train.min()}')

print("A sample:")
print("y_train: ", y_train[0])
# print("x_train: \n", x_train[0])
show_pic = x_train[0].copy()
show_pic = cv2.resize(show_pic, (28 * 10, 28 * 10))
cv2.imshow("A image sample", show_pic)
key = cv2.waitKey(0)
# 按 q 退出
if key == ord('q'):
    cv2.destroyAllWindows()
    print("show demo over")

转换为tf 数据集的格式, 并进行归一化

# convert to tf tensor
x_train = tf.convert_to_tensor(x_train, dtype=tf.float32) // 255.
x_test = tf.convert_to_tensor(x_test, dtype=tf.float32) // 255.
dataset_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset_train = dataset_train.batch(batch_size).repeat(class_num)

定义网络

在这里定义一个简单的全连接网络

def build_simple_net():
    net = Sequential([
        layers.Dense(256, activation='relu'),
        layers.Dense(256, activation='relu'),
        layers.Dense(256, activation='relu'),
        layers.Dense(class_num)
    ])
    net.build(input_shape=(None, 28 * 28))
    # net.summary()
    return net

训练

使用 SGD 优化器进行训练

def train(print_info_step=250):
    net = build_simple_net()
    # 优化器
    optimizer = optimizers.SGD(lr=0.01)
    # 计算准确率
    acc = metrics.Accuracy()

    for step, (x, y) in enumerate(dataset_train):
        with tf.GradientTape() as tape:
            # [b, 28, 28] => [b, 784]
            x = tf.reshape(x, (-1, 28 * 28))
            # [b, 784] => [b, 10]
            out = net(x)
            # [b] => [b, 10]
            y_onehot = tf.one_hot(y, depth=class_num)
            # [b, 10]
            loss = tf.square(out - y_onehot)
            # [b]
            loss = tf.reduce_sum(loss) / batch_size

        # 反向传播
        acc.update_state(tf.argmax(out, axis=1), y)
        grads = tape.gradient(loss, net.trainable_variables)
        optimizer.apply_gradients(zip(grads, net.trainable_variables))

        if acc.result() >= 0.90:
            net.save_weights(save_path)
            print(f'final acc: {acc.result()}, total step: {step}')
            break

        if step % print_info_step == 0:
            print(f'step: {step}, loss: {loss}, acc: {acc.result().numpy()}')
            acc.reset_states()

        if step % 500 == 0 and step != 0:
            print('save model')
            net.save_weights(save_path)

验证

验证在测试集的模型效果, 这里仅取出第一张进行验证

def test_dataset():
    net = build_simple_net()
    # 加载模型
    net.load_weights(save_path)
    # 拿到测试集第一张图片
    pred_image = x_test[0]
    pred_image = tf.reshape(pred_image, (-1, 28 * 28))
    pred = net.predict(pred_image)
    # print(pred)
    print(f'pred: {tf.argmax(pred, axis=1).numpy()}, label: {y_test[0]}')

应用

分割手写数字, 并进行逐一识别

• 先将图像二值化
• 找到轮廓
• 得到数字的坐标
• 转为模型的需要的输入格式, 并进行识别
• 显示

def split_number(img):
    result = []
    net = build_simple_net()
    # 加载模型
    net.load_weights(save_path)

    image = cv2.cvtColor(img.copy(), cv2.COLOR_RGB2GRAY)
    ret, thresh = cv2.threshold(image, 127, 255, 0)
    contours, hierarchy = cv2.findContours(thresh, 1, 2)
    for cnt in contours[:-1]:
        x, y, w, h = cv2.boundingRect(cnt)

        image = img[y:y+h, x:x+w]
        image = cv2.resize(image, (28, 28))

        pred_image = tf.convert_to_tensor(image, dtype=tf.float32) / 255.
        pred_image = tf.reshape(pred_image, (-1, 28 * 28))
        pred = net.predict(pred_image)
        out = tf.argmax(pred, axis=1).numpy()
        result = [out[0]] + result
        img = cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)

    cv2.imshow("demo", img)
    print(result)
    k = cv2.waitKey(0)
    # 按 q 退出
    if k == ord('q'):
        pass
    cv2.destroyAllWindows()

效果

单数字

多数字

附录

所有代码, 文件 tf2_mnist.py

import os
import cv2
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, Sequential, optimizers, metrics

# 屏蔽通知信息和警告信息
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

# 每批几张图片
batch_size = 2
# 类别数
class_num = 10
# 保存模型的路径
save_path = "./models/mnist.ckpt"
# 展示样例
show_demo = False
# 验证测试集
evaluate_dataset = False
# 是否训练
run_train = False
# 图片路径, 仅用于 detect_image(), 当为False时不识别
image_path = 'images/36.png'

path = 'MNIST/mnist.npz'
with np.load(path, allow_pickle=True) as f:
    x_train, y_train = f['x_train'], f['y_train']
    x_test, y_test = f['x_test'], f['y_test']

if show_demo:
    print(f'dataset info: shape: {x_train.shape}, {y_train.shape}')
    print(f'dataset info: max: {x_train.max()}')
    print(f'dataset info: min: {x_train.min()}')

    print("A sample:")
    print("y_train: ", y_train[0])
    # print("x_train: \n", x_train[0])
    show_pic = x_train[0].copy()
    show_pic = cv2.resize(show_pic, (28 * 10, 28 * 10))
    cv2.imshow("A image sample", show_pic)
    key = cv2.waitKey(0)
    if key == ord('q'):
        cv2.destroyAllWindows()
        print("show demo over")

# convert to tf tensor
x_train = tf.convert_to_tensor(x_train, dtype=tf.float32) // 255.
x_test = tf.convert_to_tensor(x_test, dtype=tf.float32) // 255.
dataset_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset_train = dataset_train.batch(batch_size).repeat(class_num)


def build_simple_net():
    net = Sequential([
        layers.Dense(256, activation='relu'),
        layers.Dense(256, activation='relu'),
        layers.Dense(256, activation='relu'),
        layers.Dense(class_num)
    ])
    net.build(input_shape=(None, 28 * 28))
    # net.summary()
    return net


def train(print_info_step=250):
    net = build_simple_net()
    # 优化器
    optimizer = optimizers.SGD(lr=0.01)
    # 计算准确率
    acc = metrics.Accuracy()

    for step, (x, y) in enumerate(dataset_train):
        with tf.GradientTape() as tape:
            # [b, 28, 28] => [b, 784]
            x = tf.reshape(x, (-1, 28 * 28))
            # [b, 784] => [b, 10]
            out = net(x)
            # [b] => [b, 10]
            y_onehot = tf.one_hot(y, depth=class_num)
            # [b, 10]
            loss = tf.square(out - y_onehot)
            # [b]
            loss = tf.reduce_sum(loss) / batch_size

        # 反向传播
        acc.update_state(tf.argmax(out, axis=1), y)
        grads = tape.gradient(loss, net.trainable_variables)
        optimizer.apply_gradients(zip(grads, net.trainable_variables))

        if acc.result() >= 0.90:
            net.save_weights(save_path)
            print(f'final acc: {acc.result()}, total step: {step}')
            break

        if step % print_info_step == 0:
            print(f'step: {step}, loss: {loss}, acc: {acc.result().numpy()}')
            acc.reset_states()

        if step % 500 == 0 and step != 0:
            print('save model')
            net.save_weights(save_path)


def test_dataset():
    net = build_simple_net()
    # 加载模型
    net.load_weights(save_path)
    # 拿到测试集第一张图片
    pred_image = x_test[0]
    pred_image = tf.reshape(pred_image, (-1, 28 * 28))
    pred = net.predict(pred_image)
    # print(pred)
    print(f'pred: {tf.argmax(pred, axis=1).numpy()}, label: {y_test[0]}')

def split_number(img):
    result = []
    net = build_simple_net()
    # 加载模型
    net.load_weights(save_path)

    image = cv2.cvtColor(img.copy(), cv2.COLOR_RGB2GRAY)
    ret, thresh = cv2.threshold(image, 127, 255, 0)
    contours, hierarchy = cv2.findContours(thresh, 1, 2)
    for cnt in contours[:-1]:
        x, y, w, h = cv2.boundingRect(cnt)

        image = img[y:y+h, x:x+w]
        image = cv2.resize(image, (28, 28))

        pred_image = tf.convert_to_tensor(image, dtype=tf.float32) / 255.
        pred_image = tf.reshape(pred_image, (-1, 28 * 28))
        pred = net.predict(pred_image)
        out = tf.argmax(pred, axis=1).numpy()
        result = [out[0]] + result
        img = cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)

    cv2.imshow("demo", img)
    print(result)
    k = cv2.waitKey(0)
    if k == ord('q'):
        pass
    cv2.destroyAllWindows()


if __name__ == '__main__':
    if run_train:
        train()
    elif evaluate_dataset:
        test_dataset()
    elif image_path:
        image = cv2.imread(image_path)
        # detect_image(image)
        split_number(image)