Fixed saving the names of features

The feature names were not actually saved originally.

train_model.py

@@ -358,17 +358,19 @@ if __name__ == "__main__":
                                                                          sw_hw_cleaned["Sample weight"],
                                                                          stratify=labels,
                                                                          test_size=0.4)
-    
+
     #%%
 
     # Fitting without the sample weights gives better 
     # Accuracy than with it, not sure why.
-    
+
     classifier = RandomForestClassifier()
 
     classifier.fit(X_train, y_train, w_train)
     # We just add it to the class
-    classifier.feature_names = poly_transformer.get_feature_names()
+
+    classifier.feature_names = poly_transformer \
+        .get_feature_names(features_extended.columns.values)
 
     print("Accuracy: {0}".format(classifier.score(X_test, y_test, w_test)))
     print("Classifier trained")