update model_builder.py

pdseg/models/model_builder.py

@@ -124,32 +124,12 @@ def sigmoid_to_softmax(logit):
     logit = fluid.layers.transpose(logit, [0, 3, 1, 2])
     return logit
 
-
-def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
-    if not ModelPhase.is_valid_phase(phase):
-        raise ValueError("ModelPhase {} is not valid!".format(phase))
-    if ModelPhase.is_train(phase):
-        width = cfg.TRAIN_CROP_SIZE[0]
-        height = cfg.TRAIN_CROP_SIZE[1]
-    else:
+def export_preprocess(image):
+    """导出模型的预处理流程"""
     width = cfg.EVAL_CROP_SIZE[0]
     height = cfg.EVAL_CROP_SIZE[1]
 
-    image_shape = [cfg.DATASET.DATA_DIM, height, width]
-    grt_shape = [1, height, width]
-    class_num = cfg.DATASET.NUM_CLASSES
-
-    with fluid.program_guard(main_prog, start_prog):
-        with fluid.unique_name.guard():
-            # 在导出模型的时候，增加图像标准化预处理,减小预测部署时图像的处理流程
-            # 预测部署时只须对输入图像增加batch_size维度即可
-            if ModelPhase.is_predict(phase):
-                origin_image = fluid.layers.data(
-                    name='image',
-                    shape=[-1, -1, -1, cfg.DATASET.DATA_DIM],
-                    dtype='float32',
-                    append_batch_size=False)
-                image = fluid.layers.transpose(origin_image, [0, 3, 1, 2])
+    image = fluid.layers.transpose(image, [0, 3, 1, 2])
     origin_shape = fluid.layers.shape(image)[-2:]
 
     # 不同AUG_METHOD方法的resize
@@ -185,7 +165,7 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
     image = fluid.layers.pad2d(
         image, paddings=paddings, pad_value=127.5)
 
-                #normalize
+    # normalize
     mean = np.array(cfg.MEAN).reshape(1, len(cfg.MEAN), 1, 1)
     mean = fluid.layers.assign(mean.astype('float32'))
     std = np.array(cfg.STD).reshape(1, len(cfg.STD), 1, 1)
@@ -194,6 +174,34 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
     # 很有必要，使后面的网络能通过image.shape获取特征图的shape
     image = fluid.layers.reshape(
         image, shape=[-1, cfg.DATASET.DATA_DIM, height, width])
+    return image, valid_shape, origin_shape
+
+
+def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
+    if not ModelPhase.is_valid_phase(phase):
+        raise ValueError("ModelPhase {} is not valid!".format(phase))
+    if ModelPhase.is_train(phase):
+        width = cfg.TRAIN_CROP_SIZE[0]
+        height = cfg.TRAIN_CROP_SIZE[1]
+    else:
+        width = cfg.EVAL_CROP_SIZE[0]
+        height = cfg.EVAL_CROP_SIZE[1]
+
+    image_shape = [cfg.DATASET.DATA_DIM, height, width]
+    grt_shape = [1, height, width]
+    class_num = cfg.DATASET.NUM_CLASSES
+
+    with fluid.program_guard(main_prog, start_prog):
+        with fluid.unique_name.guard():
+            # 在导出模型的时候，增加图像标准化预处理,减小预测部署时图像的处理流程
+            # 预测部署时只须对输入图像增加batch_size维度即可
+            if ModelPhase.is_predict(phase):
+                origin_image = fluid.layers.data(
+                    name='image',
+                    shape=[-1, -1, -1, cfg.DATASET.DATA_DIM],
+                    dtype='float32',
+                    append_batch_size=False)
+                image, valid_shape, origin_shape = export_preprocess(origin_image)
 
             else:
                 image = fluid.layers.data(