Depend on project caffe_pb2 instead of system caffe_pb2, fix ISSUE #12.

solrex · Oct 17, 2017 · 0a8a8bd · 0a8a8bd
1 parent 63641e7
commit 0a8a8bd
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tools/prototxt2bin.py b/tools/prototxt2bin.py
@@ -1,9 +1,13 @@
 #!/usr/bin/env python
 import sys
 from os import path
-from caffe.proto import caffe_pb2
 from google.protobuf import text_format
 
+# Search path for caffe_pb2.py
+sys.path.append(path.abspath("../build/include/caffe/proto/"))
+sys.path.append(path.abspath("../build_armeabi/include/caffe/proto/"))
+import caffe_pb2
+
 if len(sys.argv) < 2:
     print("Usage: ./prototxt2bin.py net.prototxt")