-
Notifications
You must be signed in to change notification settings - Fork 0
/
predict_emnist.py
93 lines (66 loc) · 2.89 KB
/
predict_emnist.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
# -*- coding: utf-8 -*-
"""
@date: 2023/10/10 上午9:59
@file: predict.py
@author: zj
@description:
Usage: Predict EMNIST:
$ python predict_emnist.py crnn_tiny-emnist-b512-e100.pth ../datasets/emnist/ ./runs/predict/emnist/
$ python predict_emnist.py crnn-emnist-b512-e100.pth ../datasets/emnist/ ./runs/predict/emnist/ --not-tiny
"""
import argparse
import os
import numpy as np
from itertools import groupby
import matplotlib.pyplot as plt
import torch
from utils.general import load_crnn
from utils.dataset.emnist import EMNISTDataset, DIGITS_CHARS
def parse_opt():
parser = argparse.ArgumentParser(description='Predict CRNN with EMNIST')
parser.add_argument('pretrained', metavar='PRETRAINED', type=str, help='path to pretrained model')
parser.add_argument('val_root', metavar='DIR', type=str, help='path to val dataset')
parser.add_argument('save_dir', metavar='DST', type=str, help='path to save dir')
parser.add_argument('--use-lstm', action='store_true', help='use nn.LSTM instead of nn.GRU')
parser.add_argument('--not-tiny', action='store_true', help='Use this flag to specify non-tiny mode')
args = parser.parse_args()
print(f"args: {args}")
return args
@torch.no_grad()
def predict(args, val_root, pretrained, save_dir):
img_h = 32
digits_per_sequence = 5
model, device = load_crnn(pretrained=pretrained, shape=(1, 1, img_h, digits_per_sequence * img_h),
num_classes=len(DIGITS_CHARS), not_tiny=args.not_tiny, use_lstm=args.use_lstm)
val_dataset = EMNISTDataset(val_root, is_train=False, num_of_sequences=50000,
digits_per_sequence=digits_per_sequence, img_h=img_h)
plt.figure(figsize=(10, 6))
blank_label = len(DIGITS_CHARS) - 1
for i in range(1, 7):
random_index = np.random.randint(len(val_dataset))
sequence, emnist_labels, transformed_images = val_dataset.__getitem__(random_index, return_tf=True)
# [1, H, N*W] -> [1, 1, H, N*W]
images = sequence.unsqueeze(0).to(device)
with torch.no_grad():
output = model(images).cpu()[0]
_, max_index = torch.max(output, dim=1)
raw_pred = list(max_index.numpy())
pred = torch.IntTensor([c for c, _ in groupby(raw_pred) if c != blank_label])
pred = pred.numpy()
emnist_labels = emnist_labels.numpy()
# print(pred, emnist_labels)
np_images = transformed_images
plt.subplot(3, 2, i)
title = f"Label: {str(emnist_labels)} Pred: {str(pred)}"
print(title)
plt.title(title)
plt.imshow(np_images, cmap='gray')
plt.axis('off')
if not os.path.exists(save_dir):
os.makedirs(save_dir)
plt.savefig(os.path.join(save_dir, "predict_emnist.jpg"))
def main():
args = parse_opt()
predict(args, args.val_root, args.pretrained, args.save_dir)
if __name__ == '__main__':
main()