From 1137e89fef767c68e9368779a57dfc61c6d8d834 Mon Sep 17 00:00:00 2001
From: philkr <philkr@users.noreply.github.com>
Date: Wed, 5 Aug 2015 11:54:08 -0700
Subject: [PATCH] Exposing layer top and bottom names to python

---
 include/caffe/net.hpp   | 12 ++++++++++++
 python/caffe/_caffe.cpp |  4 ++++
 python/caffe/pycaffe.py | 18 ++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/include/caffe/net.hpp b/include/caffe/net.hpp
index 1bf07d2..3b56f30 100644
--- a/include/caffe/net.hpp
+++ b/include/caffe/net.hpp
@@ -149,6 +149,18 @@ class Net {
   inline const vector<vector<Blob<Dtype>*> >& top_vecs() const {
     return top_vecs_;
   }
+  /// @brief returns the ids of the top blobs of layer i
+  inline const vector<int> & top_ids(int i) const {
+    CHECK_GE(i, 0) << "Invalid layer id";
+    CHECK_LT(i, top_id_vecs_.size()) << "Invalid layer id";
+    return top_id_vecs_[i];
+  }
+  /// @brief returns the ids of the bottom blobs of layer i
+  inline const vector<int> & bottom_ids(int i) const {
+    CHECK_GE(i, 0) << "Invalid layer id";
+    CHECK_LT(i, bottom_id_vecs_.size()) << "Invalid layer id";
+    return bottom_id_vecs_[i];
+  }
   inline const vector<vector<bool> >& bottom_need_backward() const {
     return bottom_need_backward_;
   }
diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
index 69d5533..4ea2ec6 100644
--- a/python/caffe/_caffe.cpp
+++ b/python/caffe/_caffe.cpp
@@ -232,6 +232,10 @@ BOOST_PYTHON_MODULE(_caffe) {
     .def("share_with", &Net<Dtype>::ShareTrainedLayersWith)
     .add_property("_blob_loss_weights", bp::make_function(
         &Net<Dtype>::blob_loss_weights, bp::return_internal_reference<>()))
+    .def("_bottom_ids", bp::make_function(&Net<Dtype>::bottom_ids,
+        bp::return_value_policy<bp::copy_const_reference>()))
+    .def("_top_ids", bp::make_function(&Net<Dtype>::top_ids,
+        bp::return_value_policy<bp::copy_const_reference>()))
     .add_property("_blobs", bp::make_function(&Net<Dtype>::blobs,
         bp::return_internal_reference<>()))
     .add_property("layers", bp::make_function(&Net<Dtype>::layers,
diff --git a/python/caffe/pycaffe.py b/python/caffe/pycaffe.py
index 31dc702..3054110 100644
--- a/python/caffe/pycaffe.py
+++ b/python/caffe/pycaffe.py
@@ -276,6 +276,22 @@ def _Net_batch(self, blobs):
                                                  padding])
         yield padded_batch
 
+
+class _Net_IdNameWrapper:
+    """
+    A simple wrapper that allows the ids propery to be accessed as a dict
+    indexed by names. Used for top and bottom names
+    """
+    def __init__(self, net, func):
+        self.net, self.func = net, func
+
+    def __getitem__(self, name):
+        # Map the layer name to id
+        ids = self.func(self.net, list(self.net._layer_names).index(name))
+        # Map the blob id to name
+        id_to_name = list(self.net.blobs)
+        return [id_to_name[i] for i in ids]
+
 # Attach methods to Net.
 Net.blobs = _Net_blobs
 Net.blob_loss_weights = _Net_blob_loss_weights
@@ -288,3 +304,5 @@ Net.set_input_arrays = _Net_set_input_arrays
 Net._batch = _Net_batch
 Net.inputs = _Net_inputs
 Net.outputs = _Net_outputs
+Net.top_names = property(lambda n: _Net_IdNameWrapper(n, Net._top_ids))
+Net.bottom_names = property(lambda n: _Net_IdNameWrapper(n, Net._bottom_ids))
-- 
2.7.4