added pycaffe wrapper. Preparing to clean the structure

2 files changed, 53 insertions, 8 deletions

diff --git a/src/caffe/net.cpp b/src/caffe/net.cpp
index 3266064d..5bc67a85 100644
--- a/src/caffe/net.cpp
+++ b/src/caffe/net.cpp
@@ -5,8 +5,6 @@
 #include <string>
 #include <vector>
 
-#include <boost/python.hpp>
-
 #include "caffe/proto/caffe.pb.h"
 #include "caffe/layer.hpp"
 #include "caffe/net.hpp"
@@ -240,17 +238,19 @@ const vector<Blob<Dtype>*>& Net<Dtype>::Forward(
 template <typename Dtype>
 string Net<Dtype>::Forward(const string& input_blob_protos) {
   BlobProtoVector blob_proto_vec;
-  blob_proto_vec.ParseFromString(input_blob_protos);
-  CHECK_EQ(blob_proto_vec.blobs_size(), net_input_blob_indices_.size())
-      << "Incorrect input size.";
-  for (int i = 0; i < blob_proto_vec.blobs_size(); ++i) {
-    blobs_[net_input_blob_indices_[i]]->FromProto(blob_proto_vec.blobs(i));
+  if (net_input_blob_indices_.size()) {
+    blob_proto_vec.ParseFromString(input_blob_protos);
+    CHECK_EQ(blob_proto_vec.blobs_size(), net_input_blob_indices_.size())
+        << "Incorrect input size.";
+    for (int i = 0; i < blob_proto_vec.blobs_size(); ++i) {
+      blobs_[net_input_blob_indices_[i]]->FromProto(blob_proto_vec.blobs(i));
+    }
   }
   for (int i = 0; i < layers_.size(); ++i) {
     layers_[i]->Forward(bottom_vecs_[i], &top_vecs_[i]);
   }
   blob_proto_vec.Clear();
-  for (int i = 0; i < layers_.size(); ++i) {
+  for (int i = 0; i < net_output_blobs_.size(); ++i) {
     net_output_blobs_[i]->ToProto(blob_proto_vec.add_blobs());
   }
   string output;
diff --git a/src/pycaffe/pycaffe.cpp b/src/pycaffe/pycaffe.cpp
new file mode 100644
index 00000000..fb475951
--- /dev/null
+++ b/src/pycaffe/pycaffe.cpp
@@ -0,0 +1,45 @@
+#include <boost/python.hpp>
+#include "caffe/caffe.hpp"
+
+using namespace caffe;
+using namespace boost::python;
+
+// For python, we will simply use float. This is wrapped in a
+#define PTYPE float
+
+// A simple wrapper over CaffeNet that runs the forward process.
+struct CaffeNet
+{
+  CaffeNet(string param_file, string pretrained_param_file,
+      boost::python::list bottom) {
+    vector<int> bottom_vec;
+    for (int i = 0; i < boost::python::len(bottom); ++i) {
+      bottom_vec.push_back(boost::python::extract<int>(bottom[i]));
+    }
+    net_.reset(new Net<float>(param_file, bottom_vec));
+    net_->CopyTrainedLayersFrom(pretrained_param_file);
+  }
+
+  string Forward(const string& input_blobs) {
+    return net_->Forward(input_blobs);
+  }
+  
+  void set_mode_cpu() { Caffe::set_mode(Caffe::CPU); }
+  void set_mode_gpu() { Caffe::set_mode(Caffe::GPU); }
+  
+  void set_phase_train() { Caffe::set_phase(Caffe::TRAIN); }
+  void set_phase_test() { Caffe::set_phase(Caffe::TEST); }
+
+	shared_ptr<Net<float> > net_;
+};
+
+BOOST_PYTHON_MODULE(pycaffe)
+{
+  class_<CaffeNet>("CaffeNet", init<string, string, boost::python::list>())
+      .def("Forward", &CaffeNet::Forward)
+      .def("set_mode_cpu", &CaffeNet::set_mode_cpu)
+      .def("set_mode_gpu", &CaffeNet::set_mode_gpu)
+      .def("set_phase_train", &CaffeNet::set_phase_train)
+      .def("set_phase_test", &CaffeNet::set_phase_test)
+  ;
+}
\ No newline at end of file