import os
|
import time
|
|
from PIL import Image
|
import torch
|
from torchvision import transforms
|
import cnn
|
import matplotlib.pyplot as plt
|
|
|
# 分隔图片
|
def __split_img_for_cnn(img):
|
# 切割图片
|
pass
|
|
|
if __name__ == '__main__':
|
path = 'C:/Users/Administrator/Desktop/ocr/codes/'
|
imgs = []
|
labels = []
|
for name in sorted(os.listdir(path)):
|
if name.find(".png") < 0:
|
continue
|
img = Image.open(path + name).convert('L')
|
img = transforms.ToTensor()(img)
|
imgs.append(img)
|
labels.append(int(name[0]))
|
imgs = torch.stack(imgs, 0)
|
# %% 加载模型
|
model = cnn.CNN()
|
model.load_state_dict(torch.load('model.pt', map_location=torch.device('cpu')))
|
model.eval()
|
for i in range(0, 10):
|
start_time = time.time()
|
with torch.no_grad():
|
output = model(imgs)
|
# %% 打印结果
|
pred = output.argmax(1)
|
true = torch.LongTensor(labels)
|
print(pred)
|
print(true)
|
print("识别时间：", (time.time() - start_time) * 1000)
|
|
# %% 结果显示
|
# plt.figure(figsize=(10, 4))
|
# for i in range(len(imgs)):
|
# plt.subplot(2, 5, i + 1)
|
# plt.title(f'pred {pred[i]} | true {true[i]}')
|
# plt.axis('off')
|
# plt.imshow(imgs[i].squeeze(0), cmap='gray')
|
# # plt.savefig('test.png')
|
# plt.show()
|
