Merge pull request #80 from SunAhong1993/develop

fix the --caffe_proto

README.md

@@ -40,7 +40,7 @@ x2paddle --framework=caffe --prototxt=deploy.proto --weight=deploy.caffemodel --
 |--weight | 当framework为caffe时，该参数指定caffe模型的参数文件路径 |
 |--save_dir | 指定转换后的模型保存目录路径 |
 |--model | 当framework为tensorflow时，该参数指定tensorflow的pb模型文件路径 |
-|--caffe_proto | [可选]由caffe.proto编译成caffe_pb2.py文件的存放路径，当没有安装caffe或者使用自定义Layer时使用 |
+|--caffe_proto | [可选]由caffe.proto编译成caffe_pb2.py文件的存放路径，当没有安装caffe或者使用自定义Layer时使用，默认为None |
 
 ## 使用转换后的模型
 转换后的模型包括`model_with_code`和`inference_model`两个目录。  

x2paddle/convert.py

@@ -45,11 +45,12 @@ def arg_parser():
                         type=_text_type,
                         default=None,
                         help="define which deeplearning framework")
-    parser.add_argument("--caffe_proto",
+    parser.add_argument(
+        "--caffe_proto",
         "-c",
         type=_text_type,
         default=None,
-                        help="caffe proto file of caffe model")
+        help="the .py file compiled by caffe proto file of caffe model")
     parser.add_argument("--version",
                         "-v",
                         action="store_true",
@@ -92,8 +93,8 @@ def tf2paddle(model_path, save_dir):
 def caffe2paddle(proto, weight, save_dir, caffe_proto):
     if caffe_proto is not None:
         import os
-        if not os.path.isfile(caffe_proto + 'caffe_pb2.py'):
-            print("The file that resolve caffe is not exist.")
+        if caffe_proto is not None and not os.path.isfile(caffe_proto):
+            print("The .py file compiled by caffe.proto is not exist.")
             return
     else:
         try:

x2paddle/decoder/caffe_decoder.py

@@ -22,18 +22,21 @@ from x2paddle.op_mapper import caffe_shape
 
 
 class CaffeResolver(object):
-    def __init__(self, caffe_proto_folder=None):
-        self.proto_path = caffe_proto_folder
-        if self.proto_path == None:
+    def __init__(self, caffe_proto):
+        self.proto_path = caffe_proto
+        if self.proto_path is None:
             self.use_default = True
         else:
             self.use_default = False
         self.import_caffe()
 
     def import_caffepb(self):
-        sys.path.append(self.proto_path)
-        import caffe_pb2
-        return caffe_pb2
+        (filepath,
+         tempfilename) = os.path.split(os.path.abspath(self.proto_path))
+        (filename, extension) = os.path.splitext(tempfilename)
+        sys.path.append(filepath)
+        out = __import__(filename)
+        return out
 
     def import_caffe(self):
         self.caffe = None
@@ -139,9 +142,17 @@ class CaffeGraph(Graph):
                                 dim=[dims[0], dims[1], dims[2], dims[3]
                                      ]))).to_proto().layer[0])
                     except:
-                        raise ImportError(
-                            'The .proto file does not work for the old style prototxt. You must install the caffe or modify the old style to new style in .protottx file.'
+                        print(
+                            "The .py file compiled by .proto file does not work for the old style prototxt. "
                         )
+                        print("There are 2 solutions for you as below:")
+                        print(
+                            "1. install caffe and don\'t set \'--caffe_proto\'."
+                        )
+                        print(
+                            "2. modify your .prototxt from the old style to the new style."
+                        )
+                        sys.exit(-1)
                     data.name = self.model.input[i]
                     data.top[0] = self.model.input[i]
             else:
@@ -155,9 +166,17 @@ class CaffeGraph(Graph):
                                 dim=[dims[0], dims[1], dims[2], dims[3]
                                      ]))).to_proto().layer[0])
                     except:
-                        raise ImportError(
-                            'The .proto file does not work for the old style prototxt. You must install the caffe or modify the old style to new style in .protottx file.'
+                        print(
+                            "The .py file compiled by .proto file does not work for the old style prototxt. "
+                        )
+                        print("There are 2 solutions for you as below:")
+                        print(
+                            "1. install caffe and don\'t set \'--caffe_proto\'."
+                        )
+                        print(
+                            "2. modify your .prototxt from the old style to the new style."
                         )
+                        sys.exit(-1)
                     data.name = self.model.input[i]
                     data.top[0] = self.model.input[i]
             layers = [data] + layers
@@ -202,11 +221,11 @@ class CaffeGraph(Graph):
 
 
 class CaffeDecoder(object):
-    def __init__(self, proto_path, model_path, caffe_proto_folder=None):
+    def __init__(self, proto_path, model_path, caffe_proto=None):
         self.proto_path = proto_path
         self.model_path = model_path
 
-        self.resolver = CaffeResolver(caffe_proto_folder=caffe_proto_folder)
+        self.resolver = CaffeResolver(caffe_proto=caffe_proto)
         self.net = self.resolver.NetParameter()
         with open(proto_path, 'rb') as proto_file:
             proto_str = proto_file.read()