| import cv2 |
| import torch |
| import torchvision.transforms as transforms |
| import time |
| from Model import MNIST |
|
|
|
|
| def images2tensor(image): |
| img = cv2.imread(image) |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
| transf = transforms.ToTensor() |
| img_tensor = torch.unsqueeze(transf(img), dim=0) |
| return img_tensor |
|
|
|
|
| if __name__ == "__main__": |
| start = time.time() |
| device = torch.device('cpu') |
| model = MNIST().to(device) |
| model.load_state_dict(torch.load('mnist.pkl')) |
| input_data = images2tensor("0.png") |
| res = model(input_data) |
| end = time.time() |
| print("手写数字图片检测的结果为：", res.argmax()) |
| print("infer time: ", end - start) |
|
|
