fix a bug with classifier data path

Te-k · Jul 16, 2016 · 9f2b222 · 9f2b222
1 parent 0424b3a
commit 9f2b222
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/checkmanalyzer.py b/checkmanalyzer.py
@@ -3,6 +3,7 @@
 import pickle
 import requests
 import sys
+import os
 from sklearn.externals import joblib
 
 MACHINE_TYPES = {
@@ -186,8 +187,15 @@ def feature_extraction(data):
     args = parser.parse_args()
 
     # Load classifier
-    clf = joblib.load('classifier/classifier.pkl')
-    features = pickle.loads(open('classifier/features.pkl', 'r').read())
+    clf = joblib.load(os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        'classifier/classifier.pkl'
+    ))
+    features = pickle.loads(open(os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        'classifier/features.pkl'),
+        'r').read()
+    )
 
     if 'manalyzer.org' not in args.URL:
         print('This is not a manalyzer url')

diff --git a/checkpe.py b/checkpe.py
@@ -175,8 +175,15 @@ def extract_infos(fpath):
     parser.add_argument('FILE', help='File to be tested')
     args = parser.parse_args()
     # Load classifier
-    clf = joblib.load('classifier/classifier.pkl')
-    features = pickle.loads(open('classifier/features.pkl', 'r').read())
+    clf = joblib.load(os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        'classifier/classifier.pkl'
+    ))
+    features = pickle.loads(open(os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        'classifier/features.pkl'),
+        'r').read()
+    )
 
     data = extract_infos(args.FILE)