make RunInferenceOnImage accept parameters

heigeo · May 9, 2017 · d1a98ab · d1a98ab
1 parent 3ae8a18
commit d1a98ab
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 84 deletions.
diff --git a/src/ios/tf_libs/tensorflow_utils.h b/src/ios/tf_libs/tensorflow_utils.h
@@ -49,4 +49,11 @@ void GetTopN(const Eigen::TensorMap<Eigen::Tensor<float, 1, Eigen::RowMajor>,
              const int num_results, const float threshold,
              std::vector<std::pair<float, int> >* top_results);
 
+struct Result {
+   NSString* label;
+   NSNumber* confidence;
+};
+
+tensorflow::Status RunInferenceOnImage(NSString* image, int input_size, float input_mean, float input_std, std::string input_layer, std::string output_layer, std::unique_ptr<tensorflow::Session>* session, std::vector<std::string>* labels, std::vector<Result>* results);
+
 #endif  // TENSORFLOW_CONTRIB_IOS_EXAMPLES_CAMERA_TENSORFLOW_UTILS_H_
diff --git a/src/ios/tf_libs/tensorflow_utils.mm b/src/ios/tf_libs/tensorflow_utils.mm
@@ -15,6 +15,7 @@
 #import <Foundation/Foundation.h>
 
 #include "tensorflow_utils.h"
+#include "ios_image_load.h"
 
 #include <pthread.h>
 #include <unistd.h>
@@ -204,56 +205,17 @@ bool PortableReadFileToProto(const std::string& file_name,
   return tensorflow::Status::OK();
 }
 
-NSString* RunInferenceOnImage() {
-  tensorflow::SessionOptions options;
-
-  tensorflow::Session* session_pointer = nullptr;
-  tensorflow::Status session_status = tensorflow::NewSession(options, &session_pointer);
-  if (!session_status.ok()) {
-    std::string status_string = session_status.ToString();
-    return [NSString stringWithFormat: @"Session create failed - %s",
-	status_string.c_str()];
-  }
-  std::unique_ptr<tensorflow::Session> session(session_pointer);
-  LOG(INFO) << "Session created.";
-
-  tensorflow::GraphDef tensorflow_graph;
-  LOG(INFO) << "Graph created.";
+tensorflow::Status RunInferenceOnImage(NSString* image, int input_size, float input_mean, float input_std, std::string input_layer, std::string output_layer, std::unique_ptr<tensorflow::Session>* session, std::vector<std::string>* labels, std::vector<Result>* results) {
 
-  NSString* network_path = FilePathForResourceName(@"tensorflow_inception_graph", @"pb");
-  PortableReadFileToProto([network_path UTF8String], &tensorflow_graph);
-
-  LOG(INFO) << "Creating session.";
-  tensorflow::Status s = session->Create(tensorflow_graph);
-  if (!s.ok()) {
-    LOG(ERROR) << "Could not create TensorFlow Graph: " << s;
-    return @"";
-  }
-
-  // Read the label list
-  NSString* labels_path = FilePathForResourceName(@"imagenet_comp_graph_label_strings", @"txt");
-  std::vector<std::string> label_strings;
-  std::ifstream t;
-  t.open([labels_path UTF8String]);
-  std::string line;
-  while(t){
-    std::getline(t, line);
-    label_strings.push_back(line);
-  }
-  t.close();
-
-  // Read the Grace Hopper image.
-  NSString* image_path = FilePathForResourceName(@"grace_hopper", @"jpg");
   int image_width;
   int image_height;
   int image_channels;
-  std::vector<tensorflow::uint8> image_data = LoadImageFromFile(
-	[image_path UTF8String], &image_width, &image_height, &image_channels);
-  const int wanted_width = 224;
-  const int wanted_height = 224;
+  std::vector<tensorflow::uint8> image_data = LoadImageFromBase64(
+        image, &image_width, &image_height, &image_channels);
+  const int wanted_width = input_size;
+  const int wanted_height = input_size;
   const int wanted_channels = 3;
-  const float input_mean = 117.0f;
-  const float input_std = 1.0f;
+
   assert(image_channels >= wanted_channels);
   tensorflow::Tensor image_tensor(
       tensorflow::DT_FLOAT,
@@ -277,56 +239,26 @@ bool PortableReadFileToProto(const std::string& file_name,
     }
   }
 
-  NSString* result = [network_path stringByAppendingString: @" - loaded!"];
-  result = [NSString stringWithFormat: @"%@ - %d, %s - %dx%d", result,
-	label_strings.size(), label_strings[0].c_str(), image_width, image_height];
-
-  std::string input_layer = "input";
-  std::string output_layer = "output";
   std::vector<tensorflow::Tensor> outputs;
-  tensorflow::Status run_status = session->Run({{input_layer, image_tensor}},
-				               {output_layer}, {}, &outputs);
+  tensorflow::Status run_status = (*session)->Run({{input_layer, image_tensor}},
+                                               {output_layer}, {}, &outputs);
   if (!run_status.ok()) {
     LOG(ERROR) << "Running model failed: " << run_status;
-    tensorflow::LogAllRegisteredKernels();
-    result = @"Error running model";
-    return result;
+    return run_status;
   }
-  tensorflow::string status_string = run_status.ToString();
-  result = [NSString stringWithFormat: @"%@ - %s", result,
-	status_string.c_str()];
 
   tensorflow::Tensor* output = &outputs[0];
   const int kNumResults = 5;
   const float kThreshold = 0.1f;
   std::vector<std::pair<float, int> > top_results;
   GetTopN(output->flat<float>(), kNumResults, kThreshold, &top_results);
-
-  std::stringstream ss;
-  ss.precision(3);
   for (const auto& result : top_results) {
-    const float confidence = result.first;
-    const int index = result.second;
-
-    ss << index << " " << confidence << "  ";
-
-    // Write out the result as a string
-    if (index < label_strings.size()) {
-      // just for safety: theoretically, the output is under 1000 unless there
-      // is some numerical issues leading to a wrong prediction.
-      ss << label_strings[index];
-    } else {
-      ss << "Prediction: " << index;
-    }
-
-    ss << "\n";
+    struct Result res;
+    std::string label = labels->at(result.second % labels->size());
+    res.label = [NSString stringWithUTF8String:label.c_str()];
+    res.confidence = [NSNumber numberWithFloat:result.first];
+    results->push_back(res);
   }
 
-  LOG(INFO) << "Predictions: " << ss.str();
-
-  tensorflow::string predictions = ss.str();
-  result = [NSString stringWithFormat: @"%@ - %s", result,
-	predictions.c_str()];
-
-  return result;
+  return run_status;
 }