Fix nan issue in features
[trackerpp.git] / python / predictor.py
1 from sklearn.linear_model import LogisticRegression
2 try:
3     import cPickle as pickle
4 except:
5     import pickle
6
7
8 predictors = None
9
10 def init(fname = './model.pkl'):
11     global predictors
12     f = open(fname, 'rb')
13     predictors = pickle.load(f)
14     f.close()
15     return True;
16
17 def dump():
18     global predictors
19     ss = '\n'
20     for i in predictors:
21         ss += str(i.__dict__)
22         ss += '\n'
23     return ss
24
25
26 def predict(index, features):
27     pp = predictors[index]
28     true_class = int(pp.classes_[1] == 1)
29     prob = pp.predict_proba([features])[0, true_class]
30     print prob
31     return prob
32
33
34 if __name__ == '__main__':
35     init('./model.pkl')
36     dump()
37     feature=[1,1,1,1,1,1,1,1,1,1,  1,1,1,1,1,1,1,1,1,1,  2,2,2,2,2,2,2,2,2,2, 2,2,2,2,2,2,2,2,2,2]
38     print predict(len(predictors) - 1, feature)
39