update model_builder.py for fitting different aug_method

pdseg/models/model_builder.py

@@ -151,16 +151,47 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
                     append_batch_size=False)
                 image = fluid.layers.transpose(origin_image, [0, 3, 1, 2])
                 origin_shape = fluid.layers.shape(image)[-2:]
+
+                # 不同AUG_METHOD方法的resize
+                if cfg.AUG.AUG_METHOD == 'unpadding':
+                    h = cfg.AUG.FIX_RESIZE_SIZE[1]
+                    w = cfg.AUG.FIX_RESIZE_SIZE[0]
+                    image = fluid.layers.resize_bilinear(
+                        image,
+                        out_shape=[h, w],
+                        align_corners=False,
+                        align_mode=0)
+                if cfg.AUG.AUG_METHOD == 'stepscaling':
+                    pass
+                if cfg.AUG.AUG_METHOD == 'rangescaling':
+                    size = cfg.AUG.INF_RESIZE_VALUE
+                    value = fluid.layers.reduce_max(origin_shape)
+                    scale = size / value
+                    image = fluid.layers.resize_bilinear(
+                        image, scale=scale, align_corners=False, align_mode=0)
+
+                # 存储resize后图像shape
+                valid_shape = fluid.layers.shape(image)[-2:]
+
+                # padding 到eval_crop_size大小
+                pad_target = fluid.layers.assign(
+                    np.array([height, width]).astype('float32'))
+                up = fluid.layers.assign(np.array([0]).astype('float32'))
+                down = pad_target[0] - valid_shape[0]
+                left = up
+                right = pad_target[1] - valid_shape[1]
+                paddings = fluid.layers.concat([up, down, left, right])
+                paddings = fluid.layers.cast(paddings, 'int32')
+                image = fluid.layers.pad2d(
+                    image, paddings=paddings, pad_value=127.5)
+
+                #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)
                 std = fluid.layers.assign(std.astype('float32'))
-                image = fluid.layers.resize_bilinear(
-                    image,
-                    out_shape=[height, width],
-                    align_corners=False,
-                    align_mode=0)
                 image = (image / 255 - mean) / std
+
             else:
                 image = fluid.layers.data(
                     name='image', shape=image_shape, dtype='float32')
@@ -252,13 +283,17 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
                     logit = sigmoid_to_softmax(logit)
                 else:
                     logit = softmax(logit)
+
+                # 获取有效部分
+                logit = fluid.layers.slice(
+                    logit, axes=[2, 3], starts=[0, 0], ends=valid_shape)
+
                 logit = fluid.layers.resize_bilinear(
                     logit,
                     out_shape=origin_shape,
                     align_corners=False,
                     align_mode=0)
-                logit = fluid.layers.transpose(logit, [0, 2, 3, 1])
-                logit = fluid.layers.argmax(logit, axis=3)
+                logit = fluid.layers.argmax(logit, axis=1)
                 return origin_image, logit
 
             if class_num == 1: