110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
|
import tensorflow as tf
|
|||
|
|
import numpy as np
|
|||
|
|
from tensorflow.keras import layers
|
|||
|
|
from tensorflow.keras.models import Sequential
|
|||
|
|
import pathlib
|
|||
|
|
import cv2
|
|||
|
|
|
|||
|
|
|
|||
|
|
# %% 构建模型
|
|||
|
|
def create_model():
|
|||
|
|
model = Sequential([
|
|||
|
|
layers.experimental.preprocessing.Rescaling(1. / 255, input_shape=(24, 24, 1)),
|
|||
|
|
layers.Conv2D(24, 3, padding='same', activation='relu'),
|
|||
|
|
layers.MaxPooling2D(),
|
|||
|
|
layers.Conv2D(32, 3, padding='same', activation='relu'),
|
|||
|
|
layers.MaxPooling2D(),
|
|||
|
|
layers.Conv2D(64, 3, padding='same', activation='relu'),
|
|||
|
|
layers.MaxPooling2D(),
|
|||
|
|
layers.Dropout(0.2),
|
|||
|
|
layers.Flatten(),
|
|||
|
|
layers.Dense(96, activation='relu'),
|
|||
|
|
layers.Dense(15)]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
model.compile(optimizer='adam',
|
|||
|
|
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
|
|||
|
|
metrics=['accuracy'])
|
|||
|
|
|
|||
|
|
return model
|
|||
|
|
|
|||
|
|
|
|||
|
|
# %% 训练数据
|
|||
|
|
def train():
|
|||
|
|
# 统计文件夹下的所有图片数量
|
|||
|
|
data_dir = pathlib.Path('dataset')
|
|||
|
|
batch_size = 64
|
|||
|
|
img_width = 24
|
|||
|
|
img_height = 24
|
|||
|
|
|
|||
|
|
# 从文件夹下读取图片，生成数据集
|
|||
|
|
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
|
|||
|
|
data_dir,
|
|||
|
|
validation_split=0.2,
|
|||
|
|
subset='training',
|
|||
|
|
seed=123,
|
|||
|
|
color_mode="grayscale",
|
|||
|
|
image_size=(img_height, img_width),
|
|||
|
|
batch_size=batch_size
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
val_ds = tf.keras.preprocessing.image_dataset_from_directory(
|
|||
|
|
data_dir,
|
|||
|
|
validation_split=0.2,
|
|||
|
|
subset="validation",
|
|||
|
|
seed=123,
|
|||
|
|
color_mode="grayscale",
|
|||
|
|
image_size=(img_height, img_width),
|
|||
|
|
batch_size=batch_size
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 数据集的分类，对应dataset文件夹下有多少图片分类
|
|||
|
|
class_names = train_ds.class_names
|
|||
|
|
# 保存数据集分类
|
|||
|
|
np.save("checkpoint/class_name.npy", class_names)
|
|||
|
|
|
|||
|
|
# 数据集缓存处理
|
|||
|
|
AUTOTUNE = tf.data.experimental.AUTOTUNE
|
|||
|
|
train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
|
|||
|
|
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
|
|||
|
|
# 创建模型
|
|||
|
|
model = create_model()
|
|||
|
|
# 训练模型，epochs=10，所有数据集训练10遍
|
|||
|
|
model.fit(train_ds, validation_data=val_ds, epochs=20)
|
|||
|
|
# 保存训练后的权重
|
|||
|
|
model.save_weights('checkpoint/char_checkpoint')
|
|||
|
|
|
|||
|
|
|
|||
|
|
# %% 预测
|
|||
|
|
def predict(model, imgs, class_name):
|
|||
|
|
label_dict = {0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: '=', 11: '+',
|
|||
|
|
12: '-', 13: '×', 14: '÷'}
|
|||
|
|
# 预测图片，获取预测值
|
|||
|
|
predicts = model.predict(imgs)
|
|||
|
|
results = [] # 保存结果的数组
|
|||
|
|
for predict in predicts: # 遍历每一个预测结果
|
|||
|
|
index = np.argmax(predict) # 寻找最大值
|
|||
|
|
result = class_name[index] # 取出字符
|
|||
|
|
results.append(label_dict[int(result)])
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
|
|||
|
|
# %%
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
train()
|
|||
|
|
|
|||
|
|
# model = create_model()
|
|||
|
|
# # 加载前期训练好的权重
|
|||
|
|
# model.load_weights('checkpoint/char_checkpoint')
|
|||
|
|
# # 读出图片分类
|
|||
|
|
# class_name = np.load('checkpoint/class_name.npy')
|
|||
|
|
# print(class_name)
|
|||
|
|
# img1=cv2.imread('img1.png',0)
|
|||
|
|
# img2=cv2.imread('img2.png',0)
|
|||
|
|
# img3=cv2.imread('img3.png',0)
|
|||
|
|
# img4=cv2.imread('img4.png',0)
|
|||
|
|
# img5=cv2.imread('img5.png',0)
|
|||
|
|
# img6=cv2.imread('img6.png',0)
|
|||
|
|
# imgs = np.array([img1,img2,img3,img4,img5,img6])
|
|||
|
|
# results = predict(model, imgs, class_name)
|
|||
|
|
# print(results)
|