1 files changed, 12 insertions, 7 deletions

diff --git a/main.cpp b/main.cpp
index a517d88..d5c6f69 100644
--- a/main.cpp
+++ b/main.cpp
@@ -14,10 +14,11 @@ namespace po = boost::program_options;
 
 po::variables_map vm;
 
-int train() {
+typedef int (*mainFunc)(DataProvider &dp);
+
+int train(DataProvider &dp) {
     RSVM *rsvm;
     rsvm = RSVM::loadModel(vm["model"].as<std::string>());
-    FileDP dp(vm["feature"].as<std::string>());
 
     dp.open();
     DataList D;
@@ -39,10 +40,9 @@ int train() {
     return 0;
 }
 
-int predict() {
+int predict(DataProvider &dp) {
     RSVM *rsvm;
     rsvm = RSVM::loadModel(vm["model"].as<std::string>().c_str());
-    FileDP dp(vm["feature"].as<std::string>().c_str());
 
     dp.open();
     DataList D;
@@ -99,12 +99,17 @@ int main(int argc, char **argv) {
         std::cout << desc;
         return 0;
     }
-
+    mainFunc mainf;
     if (vm.count("train")) {
-        train();
+        mainf = &train;
     }
     else if (vm.count("validate")||vm.count("predict")) {
-        predict();
+        mainf = &predict;
     }
+    else return 0;
+    DataProvider* dp;
+    dp = new FileDP(vm["feature"].as<std::string>());
+    mainf(*dp);
+    delete dp;
     return 0;
 }