-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathtrain.py
61 lines (44 loc) · 1.69 KB
/
train.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
import os
import numpy as np
from keras.models import load_model
from load_data import load_dataset
from embedding import get_embedded_data
from sklearn.preprocessing import Normalizer, LabelEncoder
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
import joblib
def normalize(trainX):
norm_encoder = Normalizer(norm='l2')
trainX = norm_encoder.transform(trainX)
return trainX
def get_svm_model(trainX, trainy):
model = SVC(kernel='linear', probability=True)
model.fit(trainX, trainy)
return model
def train_svm():
input_dir = 'data'
train_dir = os.path.join(input_dir,'processed_data', 'train')
validation_dir = os.path.join(input_dir,'processed_data', 'validation')
trainX, trainy = load_dataset(train_dir)
testX, testy = load_dataset(validation_dir)
model = load_model(os.path.join('model','facenet_keras.h5'))
trainX = get_embedded_data(model, trainX)
testX = get_embedded_data(model, testX)
trainX = normalize(trainX)
testX = normalize(testX)
label_encode = LabelEncoder()
label_encode.fit(trainy)
trainy = label_encode.transform(trainy)
testy = label_encode.transform(testy)
np.save(os.path.join('model','classes.npy'), label_encode.classes_)
model = get_svm_model(trainX, trainy)
filename = os.path.join('model', 'svm_model.sav')
joblib.dump(model, filename)
print("SVM model saved!")
pred_train = model.predict(trainX)
pred_test = model.predict(testX)
score_train = accuracy_score(trainy, pred_train)
score_test = accuracy_score(testy, pred_test)
print("Accuracy\nTrain : ",score_train,"\n","Test : ", score_test)
if __name__ == "__main__":
train_svm()