BVLC · shelhamer · Mar 18, 2014 · Mar 1, 2014 · Mar 11, 2014 · Mar 11, 2014
diff --git a/Makefile b/Makefile
@@ -45,8 +45,8 @@ NONGEN_CXX_SRCS := $(shell find \
 LINT_REPORT := $(BUILD_DIR)/cpp_lint.log
 FAILED_LINT_REPORT := $(BUILD_DIR)/cpp_lint.error_log
 # PY$(PROJECT)_SRC is the python wrapper for $(PROJECT)
-PY$(PROJECT)_SRC := python/$(PROJECT)/py$(PROJECT).cpp
-PY$(PROJECT)_SO := python/$(PROJECT)/py$(PROJECT).so
+PY$(PROJECT)_SRC := python/$(PROJECT)/_$(PROJECT).cpp
+PY$(PROJECT)_SO := python/$(PROJECT)/_$(PROJECT).so
 # MAT$(PROJECT)_SRC is the matlab wrapper for $(PROJECT)
 MAT$(PROJECT)_SRC := matlab/$(PROJECT)/mat$(PROJECT).cpp
 MAT$(PROJECT)_SO := matlab/$(PROJECT)/$(PROJECT)

diff --git a/python/caffe/__init__.py b/python/caffe/__init__.py
@@ -1 +1 @@
-from .pycaffe import *
+from .pycaffe import Net
diff --git a/python/caffe/pycaffe.cpp → python/caffe/_caffe.cpp b/python/caffe/pycaffe.cpp → python/caffe/_caffe.cpp
@@ -38,19 +38,14 @@ class CaffeBlob {
   CaffeBlob(const shared_ptr<Blob<float> > &blob, const string& name)
       : blob_(blob), name_(name) {}
 
-  explicit CaffeBlob(const shared_ptr<Blob<float> > &blob)
-      : blob_(blob) {}
-
-  CaffeBlob()
-  {}
-
   string name() const { return name_; }
   int num() const { return blob_->num(); }
   int channels() const { return blob_->channels(); }
   int height() const { return blob_->height(); }
   int width() const { return blob_->width(); }
   int count() const { return blob_->count(); }
 
+  // this is here only to satisfy boost's vector_indexing_suite
   bool operator == (const CaffeBlob &other) {
       return this->blob_ == other.blob_;
   }
@@ -66,9 +61,6 @@ class CaffeBlob {
 //  is not freed while still being used in Python
 class CaffeBlobWrap : public CaffeBlob {
  public:
-  CaffeBlobWrap(PyObject *p, const shared_ptr<Blob<float> > &blob)
-      : CaffeBlob(blob), self_(p) {}
-
   CaffeBlobWrap(PyObject *p, const CaffeBlob &blob)
       : CaffeBlob(blob), self_(p) {}
 
@@ -101,6 +93,30 @@ class CaffeBlobWrap : public CaffeBlob {
 };
 
 
+class CaffeLayer {
+ public:
+  CaffeLayer(const shared_ptr<Layer<float> > &layer, const string &name)
+    : layer_(layer), name_(name) {}
+
+  string name() const { return name_; }
+  vector<CaffeBlob> blobs() {
+    vector<CaffeBlob> result;
+    for (int i = 0; i < layer_->blobs().size(); ++i) {
+      result.push_back(CaffeBlob(layer_->blobs()[i], name_));
+    }
+    return result;
+  }
+
+  // this is here only to satisfy boost's vector_indexing_suite
+  bool operator == (const CaffeLayer &other) {
+      return this->layer_ == other.layer_;
+  }
+
+ protected:
+  shared_ptr<Layer<float> > layer_;
+  string name_;
+};
+
 
 // A simple wrapper over CaffeNet that runs the forward process.
 struct CaffeNet {
@@ -217,7 +233,7 @@ struct CaffeNet {
   }
 
   void ForwardPrefilled() {
-      net_->ForwardPrefilled();
+    net_->ForwardPrefilled();
   }
 
   // The caffe::Caffe utility functions.
@@ -228,24 +244,19 @@ struct CaffeNet {
   void set_device(int device_id) { Caffe::SetDevice(device_id); }
 
   vector<CaffeBlob> blobs() {
-      vector<CaffeBlob> result;
-      for (int i = 0; i < net_->blobs().size(); ++i) {
-        result.push_back(CaffeBlob(net_->blobs()[i], net_->blob_names()[i]));
-      }
-      return result;
+    vector<CaffeBlob> result;
+    for (int i = 0; i < net_->blobs().size(); ++i) {
+      result.push_back(CaffeBlob(net_->blobs()[i], net_->blob_names()[i]));
+    }
+    return result;
   }
 
-  vector<CaffeBlob> params() {
-      vector<CaffeBlob> result;
-      int ix = 0;
-      for (int i = 0; i < net_->layers().size(); ++i) {
-        for (int j = 0; j < net_->layers()[i]->blobs().size(); ++j) {
-          result.push_back(
-              CaffeBlob(net_->params()[ix], net_->layer_names()[i]));
-          ix++;
-        }
-      }
-      return result;
+  vector<CaffeLayer> layers() {
+    vector<CaffeLayer> result;
+    for (int i = 0; i < net_->layers().size(); ++i) {
+      result.push_back(CaffeLayer(net_->layers()[i], net_->layer_names()[i]));
+    }
+    return result;
   }
 
   // The pointer to the internal caffe::Net instant.
@@ -255,7 +266,7 @@ struct CaffeNet {
 
 
 // The boost python module definition.
-BOOST_PYTHON_MODULE(pycaffe) {
+BOOST_PYTHON_MODULE(_caffe) {
   boost::python::class_<CaffeNet>(
       "CaffeNet", boost::python::init<string, string>())
       .def("Forward",          &CaffeNet::Forward)
@@ -266,8 +277,8 @@ BOOST_PYTHON_MODULE(pycaffe) {
       .def("set_phase_train",  &CaffeNet::set_phase_train)
       .def("set_phase_test",   &CaffeNet::set_phase_test)
       .def("set_device",       &CaffeNet::set_device)
-      .def("blobs",            &CaffeNet::blobs)
-      .def("params",           &CaffeNet::params);
+      .add_property("blobs",   &CaffeNet::blobs)
+      .add_property("layers",  &CaffeNet::layers);
 
   boost::python::class_<CaffeBlob, CaffeBlobWrap>(
       "CaffeBlob", boost::python::no_init)
@@ -280,8 +291,16 @@ BOOST_PYTHON_MODULE(pycaffe) {
       .add_property("data",     &CaffeBlobWrap::get_data)
       .add_property("diff",     &CaffeBlobWrap::get_diff);
 
+  boost::python::class_<CaffeLayer>(
+      "CaffeLayer", boost::python::no_init)
+      .add_property("name",  &CaffeLayer::name)
+      .add_property("blobs", &CaffeLayer::blobs);
+
   boost::python::class_<vector<CaffeBlob> >("BlobVec")
       .def(vector_indexing_suite<vector<CaffeBlob>, true>());
 
+  boost::python::class_<vector<CaffeLayer> >("LayerVec")
+      .def(vector_indexing_suite<vector<CaffeLayer>, true>());
+
   import_array();
 }
diff --git a/python/caffe/detection/detector.py b/python/caffe/detection/detector.py
@@ -332,15 +332,15 @@ def config(model_def, pretrained_model, gpu, image_dim, image_mean_file):
   # Initialize network by loading model definition and weights.
   t = time.time()
   print("Loading Caffe model.")
-  NET = caffe.CaffeNet(model_def, pretrained_model)
+  NET = caffe.Net(model_def, pretrained_model)
   NET.set_phase_test()
   if gpu:
     NET.set_mode_gpu()
   print("Caffe model loaded in {:.3f} s".format(time.time() - t))
 
   # Configure for input/output data
   IMAGE_DIM = image_dim
-  CROPPED_DIM = NET.blobs()[0].width
+  CROPPED_DIM = NET.blobs.values()[0].width
   IMAGE_CENTER = int((IMAGE_DIM - CROPPED_DIM) / 2)
 
     # Load the data set mean file
@@ -349,8 +349,8 @@ def config(model_def, pretrained_model, gpu, image_dim, image_mean_file):
   CROPPED_IMAGE_MEAN = IMAGE_MEAN[IMAGE_CENTER:IMAGE_CENTER + CROPPED_DIM,
                                   IMAGE_CENTER:IMAGE_CENTER + CROPPED_DIM,
                                   :]
-  BATCH_SIZE = NET.blobs()[0].num  # network batch size
-  NUM_OUTPUT = NET.blobs()[-1].channels  # number of output classes
+  BATCH_SIZE = NET.blobs.values()[0].num  # network batch size
+  NUM_OUTPUT = NET.blobs.values()[-1].channels  # number of output classes
 
 
 if __name__ == "__main__":

diff --git a/python/caffe/imagenet/wrapper.py b/python/caffe/imagenet/wrapper.py
@@ -76,7 +76,7 @@ def __init__(self, model_def_file, pretrained_model, center_only=False,
       num = 1
     else:
       num = 10
-    self.caffenet = caffe.CaffeNet(model_def_file, pretrained_model)
+    self.caffenet = caffe.Net(model_def_file, pretrained_model)
     self._output_blobs = [np.empty((num, num_output, 1, 1), dtype=np.float32)]
     self._center_only = center_only
 

diff --git a/python/caffe/pycaffe.py b/python/caffe/pycaffe.py
@@ -0,0 +1,28 @@
+"""
+Wrap the internal caffe C++ module (_caffe.so) with a clean, Pythonic
+interface.
+"""
+
+from ._caffe import CaffeNet
+from collections import OrderedDict
+
+class Net(CaffeNet):
+    """
+    The direct Python interface to caffe, exposing Forward and Backward
+    passes, data, gradients, and layer parameters
+    """
+    def __init__(self, param_file, pretrained_param_file):
+        super(Net, self).__init__(param_file, pretrained_param_file)
+        self._blobs = OrderedDict([(bl.name, bl)
+                                   for bl in super(Net, self).blobs])
+        self.params = OrderedDict([(lr.name, lr.blobs)
+                                   for lr in super(Net, self).layers
+                                   if len(lr.blobs) > 0])
+
+    @property
+    def blobs(self):
+        """
+        An OrderedDict (bottom to top, i.e., input to output) of network
+        blobs indexed by name
+        """
+        return self._blobs