expose input and output blob names to python as lists
authorEvan Shelhamer <shelhamer@imaginarynumber.net>
Wed, 14 May 2014 01:53:36 +0000 (18:53 -0700)
committerEvan Shelhamer <shelhamer@imaginarynumber.net>
Wed, 14 May 2014 20:44:02 +0000 (13:44 -0700)
include/caffe/net.hpp
python/caffe/_caffe.cpp

index 179eadb..46d72a6 100644 (file)
@@ -89,6 +89,8 @@ class Net {
   inline int num_outputs() { return net_output_blobs_.size(); }
   inline vector<Blob<Dtype>*>& input_blobs() { return net_input_blobs_; }
   inline vector<Blob<Dtype>*>& output_blobs() { return net_output_blobs_; }
+  inline vector<int>& input_blob_indices() { return net_input_blob_indices_; }
+  inline vector<int>& output_blob_indices() { return net_output_blob_indices_; }
   // has_blob and blob_by_name are inspired by
   // https://github.com/kencoken/caffe/commit/f36e71569455c9fbb4bf8a63c2d53224e32a4e7b
   // Access intermediary computation layers, testing with centre image only
index bd2b281..e1ee652 100644 (file)
@@ -350,6 +350,24 @@ struct CaffeNet {
     return result;
   }
 
+  list inputs() {
+    list input_blob_names;
+    for (vector<int>::iterator it = net_->input_blob_indices().begin();
+        it != net_->input_blob_indices().end(); ++it) {
+      input_blob_names.append(net_->blob_names()[*it]);
+    }
+    return input_blob_names;
+  }
+
+  list outputs() {
+    list output_blob_names;
+    for (vector<int>::iterator it = net_->output_blob_indices().begin();
+        it != net_->output_blob_indices().end(); ++it) {
+      output_blob_names.append(net_->blob_names()[*it]);
+    }
+    return output_blob_names;
+  }
+
   // The pointer to the internal caffe::Net instant.
   shared_ptr<Net<float> > net_;
   // if taking input from an ndarray, we need to hold references
@@ -399,6 +417,8 @@ BOOST_PYTHON_MODULE(_caffe) {
       .def("set_device",        &CaffeNet::set_device)
       .add_property("_blobs",   &CaffeNet::blobs)
       .add_property("layers",   &CaffeNet::layers)
+      .add_property("inputs",    &CaffeNet::inputs)
+      .add_property("outputs",   &CaffeNet::outputs)
       .def("_set_input_arrays", &CaffeNet::set_input_arrays);
 
   boost::python::class_<CaffeBlob, CaffeBlobWrap>(