added pycaffe wrapper. Preparing to clean the structure

diff --git a/Makefile b/Makefile
index d7ec5820edfcb117752edf99e84e854e81b8325a..b8481f37f65102313e79536d956be40e4251bb8b 100644 (file)
--- a/Makefile
+++ b/Makefile
@@ -57,9 +57,14 @@ LDFLAGS += $(foreach librarydir,$(LIBRARY_DIRS),-L$(librarydir)) \
 
 NVCC = nvcc $(NVCCFLAGS) $(CPPFLAGS) $(CUDA_ARCH)
 
-.PHONY: all test clean distclean linecount examples distribute
+.PHONY: all test clean distclean linecount examples pycaffe distribute
 
-all: $(NAME) $(STATIC_NAME) test examples
+all: $(NAME) $(STATIC_NAME) test examples pycaffe
+
+pycaffe: $(STATIC_NAME) src/pycaffe/pycaffe.cpp
+       $(CXX) -o pycaffe.so -I/usr/include/python2.7 -shared \
+                       src/pycaffe/pycaffe.cpp $(STATIC_NAME) $(CXXFLAGS) $(LDFLAGS) \
+                       $(WARNING) -lboost_python -lpython2.7
 
 linecount: clean
        cloc --read-lang-def=caffe.cloc src/caffe/

diff --git a/src/caffe/net.cpp b/src/caffe/net.cpp
index 3266064d7c70b647dde8d303e4fde35b4fa24b48..5bc67a8562c4b916a5bfcd6c6d80848957f9a0a2 100644 (file)
--- 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 (file)
index 0000000..fb47595
--- /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