-
Notifications
You must be signed in to change notification settings - Fork 10
/
test_cnn.py
49 lines (40 loc) · 1.17 KB
/
test_cnn.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
#!/usr/bin/env python
import sys, os
import numpy as np
from keras.models import load_model
from keras.utils import plot_model
def ensure_dir(file_path):
directory = os.path.dirname(file_path)
if len(directory) == 0: return
if not os.path.exists(directory):
os.makedirs(directory)
# Parameter
height = width = 48
num_classes = 7
input_shape = (height, width, 1)
model_name = 'pre6.h5'
# Read the test data
with open(sys.argv[1], "r+") as f:
line = f.read().strip().replace(',', ' ').split('\n')[1:]
raw_data = ' '.join(line)
length = width*height+1 #1 is for label
data = np.array(raw_data.split()).astype('float').reshape(-1, length)
X = data[:, 1:]
X /= 255
# Load model
model = load_model(model_name)
# Plot model
# plot_model(model,to_file='cnn_model.png')
# Predict the test data
X = X.reshape(X.shape[0], height, width, 1)
ans = model.predict_classes(X)
ans = list(ans)
# Write prediction
## check the folder of out.csv is exist; otherwise, make it
ensure_dir(sys.argv[2])
result = []
for index, value in enumerate(ans):
result.append("{0},{1}".format(index, value))
with open(sys.argv[2], "w+") as f:
f.write("id,label\n")
f.write("\n".join(result))