[Enhancement] Update data reading process (PaddlePaddle#1719)

Author: wuyefeilin

contrib/PP-HumanSeg/README_cn.md

@@ -127,7 +127,7 @@ python data/download_data.py
 ```bash
 # 通过电脑摄像头进行实时分割处理
 python bg_replace.py \
---config export_model/ppseg_lite_portrait_398x224_with_softmax/deploy.yaml
+--config export_model/pphumanseg_lite_portrait_398x224_with_softmax/deploy.yaml
 
 # 对人像视频进行分割处理
 python bg_replace.py \
@@ -143,7 +143,7 @@ python bg_replace.py \
 ```bash
 # 增加光流后处理
 python bg_replace.py \
---config export_model/ppseg_lite_portrait_398x224_with_softmax/deploy.yaml \
+--config export_model/pphumanseg_lite_portrait_398x224_with_softmax/deploy.yaml \
 --use_optic_flow
 ```
 
@@ -152,7 +152,7 @@ python bg_replace.py \
 ```bash
 # 通过电脑摄像头进行实时背景替换处理。可通过'--background_video_path'传入背景视频
 python bg_replace.py \
---config export_model/ppseg_lite_portrait_398x224_with_softmax/deploy.yaml \
+--config export_model/pphumanseg_lite_portrait_398x224_with_softmax/deploy.yaml \
 --input_shape 224 398 \
 --bg_img_path data/background.jpg
 
@@ -164,7 +164,7 @@ python bg_replace.py \
 
 # 对单张图像进行背景替换
 python bg_replace.py \
---config export_model/ppseg_lite_portrait_398x224_with_softmax/deploy.yaml \
+--config export_model/pphumanseg_lite_portrait_398x224_with_softmax/deploy.yaml \
 --input_shape 224 398 \
 --img_path data/human_image.jpg \
 --bg_img_path data/background.jpg

contrib/PP-HumanSeg/scripts/train.py

@@ -33,17 +33,25 @@ def check_logits_losses(logits_list, losses):
             .format(len_logits, len_losses))
 
 
-def loss_computation(logits_list, labels, losses, edges=None):
+def loss_computation(logits_list, label_dict, losses):
     check_logits_losses(logits_list, losses)
     loss_list = []
     for i in range(len(logits_list)):
         logits = logits_list[i]
         loss_i = losses['types'][i]
-        # Whether to use edges as labels According to loss type.
+        coef_i = losses['coef'][i]
         if loss_i.__class__.__name__ in ('BCELoss', ) and loss_i.edge_label:
-            loss_list.append(losses['coef'][i] * loss_i(logits, edges))
+            # Use edges as labels According to loss type.
+            loss_list.append(coef_i * loss_i(logits, label_dict['edge']))
+        elif loss_i.__class__.__name__ == 'MixedLoss':
+            mixed_loss_list = loss_i(logits, label_dict['label'])
+            for mixed_loss in mixed_loss_list:
+                loss_list.append(coef_i * mixed_loss)
+        elif loss_i.__class__.__name__ in ("KLLoss", ):
+            loss_list.append(coef_i *
+                             loss_i(logits_list[0], logits_list[1].detach()))
         else:
-            loss_list.append(losses['coef'][i] * loss_i(logits, labels))
+            loss_list.append(coef_i * loss_i(logits, label_dict['label']))
     return loss_list
 
 
@@ -132,21 +140,18 @@ def train(model,
             if iter > iters:
                 break
             reader_cost_averager.record(time.time() - batch_start)
-            images = data[0]
-            labels = data[1].astype('int64')
+            images = data['img']
+            labels = data['label'].astype('int64')
             edges = None
-            if len(data) == 3:
-                edges = data[2].astype('int64')
+            if 'edge' in data.keys():
+                edges = data['edge'].astype('int64')
 
             if nranks > 1:
                 logits_list = ddp_model(images)
             else:
                 logits_list = model(images)
             loss_list = loss_computation(
-                logits_list=logits_list,
-                labels=labels,
-                losses=losses,
-                edges=edges)
+                logits_list=logits_list, label_dict=data, losses=losses)
             loss = sum(loss_list)
             loss.backward()
 

deploy/python/infer.py

@@ -387,7 +387,9 @@ def run(self, imgs_path):
         logger.info("Finish")
 
     def _preprocess(self, img):
-        return self.cfg.transforms(img)[0]
+        data = {}
+        data['img'] = img
+        return self.cfg.transforms(data)['img']
 
     def _postprocess(self, results):
         if self.args.with_argmax:

deploy/python/infer_benchmark.py

@@ -168,8 +168,9 @@ def run(self, img_path):
         logger.info("Average time: %.3f ms/img" % avg_time)
 
     def _preprocess(self, img_path):
+        data = {'img': img_path}
         if self.args.resize_width == 0 and self.args.resize_height == 0:
-            return self.cfg.transforms(img_path)[0]
+            return self.cfg.transforms(data)['img']
         else:
             assert args.resize_width > 0 and args.resize_height > 0
             with codecs.open(args.cfg, 'r', 'utf-8') as file:
@@ -180,7 +181,7 @@ def _preprocess(self, img_path):
                 'target_size': [args.resize_width, args.resize_height]
             })
             transforms = DeployConfig.load_transforms(transforms_dic)
-            return transforms(img_path)[0]
+            return transforms(data)['img']
 
     def _save_imgs(self, results):
         for i in range(results.shape[0]):

paddleseg/core/infer.py

@@ -21,96 +21,24 @@
 import paddle.nn.functional as F
 
 
-def get_reverse_list(ori_shape, transforms):
-    """
-    get reverse list of transform.
-
-    Args:
-        ori_shape (list): Origin shape of image.
-        transforms (list): List of transform.
-
-    Returns:
-        list: List of tuple, there are two format:
-            ('resize', (h, w)) The image shape before resize,
-            ('padding', (h, w)) The image shape before padding.
-    """
-    reverse_list = []
-    h, w = ori_shape[0], ori_shape[1]
-    for op in transforms:
-        if op.__class__.__name__ in ['Resize']:
-            reverse_list.append(('resize', (h, w)))
-            h, w = op.target_size[0], op.target_size[1]
-        if op.__class__.__name__ in ['ResizeByLong']:
-            reverse_list.append(('resize', (h, w)))
-            long_edge = max(h, w)
-            short_edge = min(h, w)
-            short_edge = int(round(short_edge * op.long_size / long_edge))
-            long_edge = op.long_size
-            if h > w:
-                h = long_edge
-                w = short_edge
-            else:
-                w = long_edge
-                h = short_edge
-        if op.__class__.__name__ in ['ResizeByShort']:
-            reverse_list.append(('resize', (h, w)))
-            long_edge = max(h, w)
-            short_edge = min(h, w)
-            long_edge = int(round(long_edge * op.short_size / short_edge))
-            short_edge = op.short_size
-            if h > w:
-                h = long_edge
-                w = short_edge
-            else:
-                w = long_edge
-                h = short_edge
-        if op.__class__.__name__ in ['Padding']:
-            reverse_list.append(('padding', (h, w)))
-            w, h = op.target_size[0], op.target_size[1]
-        if op.__class__.__name__ in ['PaddingByAspectRatio']:
-            reverse_list.append(('padding', (h, w)))
-            ratio = w / h
-            if ratio == op.aspect_ratio:
-                pass
-            elif ratio > op.aspect_ratio:
-                h = int(w / op.aspect_ratio)
-            else:
-                w = int(h * op.aspect_ratio)
-        if op.__class__.__name__ in ['LimitLong']:
-            long_edge = max(h, w)
-            short_edge = min(h, w)
-            if ((op.max_long is not None) and (long_edge > op.max_long)):
-                reverse_list.append(('resize', (h, w)))
-                long_edge = op.max_long
-                short_edge = int(round(short_edge * op.max_long / long_edge))
-            elif ((op.min_long is not None) and (long_edge < op.min_long)):
-                reverse_list.append(('resize', (h, w)))
-                long_edge = op.min_long
-                short_edge = int(round(short_edge * op.min_long / long_edge))
-            if h > w:
-                h = long_edge
-                w = short_edge
-            else:
-                w = long_edge
-                h = short_edge
-    return reverse_list
-
-
-def reverse_transform(pred, ori_shape, transforms, mode='nearest'):
+def reverse_transform(pred, trans_info, mode='nearest'):
     """recover pred to origin shape"""
-    reverse_list = get_reverse_list(ori_shape, transforms)
     intTypeList = [paddle.int8, paddle.int16, paddle.int32, paddle.int64]
     dtype = pred.dtype
-    for item in reverse_list[::-1]:
-        if item[0] == 'resize':
+    for item in trans_info[::-1]:
+        if isinstance(item[0], list):
+            trans_mode = item[0][0]
+        else:
+            trans_mode = item[0]
+        if trans_mode == 'resize':
             h, w = item[1][0], item[1][1]
             if paddle.get_device() == 'cpu' and dtype in intTypeList:
                 pred = paddle.cast(pred, 'float32')
                 pred = F.interpolate(pred, (h, w), mode=mode)
                 pred = paddle.cast(pred, dtype)
             else:
                 pred = F.interpolate(pred, (h, w), mode=mode)
-        elif item[0] == 'padding':
+        elif trans_mode == 'padding':
             h, w = item[1][0], item[1][1]
             pred = pred[:, :, 0:h, 0:w]
         else:
@@ -205,8 +133,7 @@ def slide_inference(model, im, crop_size, stride):
 
 def inference(model,
               im,
-              ori_shape=None,
-              transforms=None,
+              trans_info=None,
               is_slide=False,
               stride=None,
               crop_size=None):
@@ -216,8 +143,7 @@ def inference(model,
     Args:
         model (paddle.nn.Layer): model to get logits of image.
         im (Tensor): the input image.
-        ori_shape (list): Origin shape of image.
-        transforms (list): Transforms for image.
+        trans_info (list): Image shape informating changed process. Default: None.
         is_slide (bool): Whether to infer by sliding window. Default: False.
         crop_size (tuple|list). The size of sliding window, (w, h). It should be probided if is_slide is True.
         stride (tuple|list). The size of stride, (w, h). It should be probided if is_slide is True.
@@ -239,8 +165,8 @@ def inference(model,
         logit = slide_inference(model, im, crop_size=crop_size, stride=stride)
     if hasattr(model, 'data_format') and model.data_format == 'NHWC':
         logit = logit.transpose((0, 3, 1, 2))
-    if ori_shape is not None:
-        logit = reverse_transform(logit, ori_shape, transforms, mode='bilinear')
+    if trans_info is not None:
+        logit = reverse_transform(logit, trans_info, mode='bilinear')
         pred = paddle.argmax(logit, axis=1, keepdim=True, dtype='int32')
         return pred, logit
     else:
@@ -249,8 +175,7 @@ def inference(model,
 
 def aug_inference(model,
                   im,
-                  ori_shape,
-                  transforms,
+                  trans_info,
                   scales=1.0,
                   flip_horizontal=False,
                   flip_vertical=False,
@@ -263,8 +188,7 @@ def aug_inference(model,
     Args:
         model (paddle.nn.Layer): model to get logits of image.
         im (Tensor): the input image.
-        ori_shape (list): Origin shape of image.
-        transforms (list): Transforms for image.
+        trans_info (list): Transforms for image.
         scales (float|tuple|list):  Scales for resize. Default: 1.
         flip_horizontal (bool): Whether to flip horizontally. Default: False.
         flip_vertical (bool): Whether to flip vertically. Default: False.
@@ -302,8 +226,7 @@ def aug_inference(model,
             logit = F.softmax(logit, axis=1)
             final_logit = final_logit + logit
 
-    final_logit = reverse_transform(
-        final_logit, ori_shape, transforms, mode='bilinear')
+    final_logit = reverse_transform(final_logit, trans_info, mode='bilinear')
     pred = paddle.argmax(final_logit, axis=1, keepdim=True, dtype='int32')
 
     return pred, final_logit

paddleseg/core/predict.py

@@ -36,6 +36,15 @@ def partition_list(arr, m):
     return [arr[i:i + n] for i in range(0, len(arr), n)]
 
 
+def preprocess(im_path, transforms):
+    data = {}
+    data['img'] = im_path
+    data = transforms(data)
+    data['img'] = data['img'][np.newaxis, ...]
+    data['img'] = paddle.to_tensor(data['img'])
+    return data
+
+
 def predict(model,
             model_path,
             transforms,
@@ -89,18 +98,13 @@ def predict(model,
     color_map = visualize.get_color_map_list(256, custom_color=custom_color)
     with paddle.no_grad():
         for i, im_path in enumerate(img_lists[local_rank]):
-            im = cv2.imread(im_path)
-            ori_shape = im.shape[:2]
-            im, _ = transforms(im)
-            im = im[np.newaxis, ...]
-            im = paddle.to_tensor(im)
+            data = preprocess(im_path, transforms)
 
             if aug_pred:
                 pred, _ = infer.aug_inference(
                     model,
-                    im,
-                    ori_shape=ori_shape,
-                    transforms=transforms.transforms,
+                    data['img'],
+                    trans_info=data['trans_info'],
                     scales=scales,
                     flip_horizontal=flip_horizontal,
                     flip_vertical=flip_vertical,
@@ -110,9 +114,8 @@ def predict(model,
             else:
                 pred, _ = infer.inference(
                     model,
-                    im,
-                    ori_shape=ori_shape,
-                    transforms=transforms.transforms,
+                    data['img'],
+                    trans_info=data['trans_info'],
                     is_slide=is_slide,
                     stride=stride,
                     crop_size=crop_size)
@@ -141,9 +144,4 @@ def predict(model,
             mkdir(pred_saved_path)
             pred_mask.save(pred_saved_path)
 
-            # pred_im = utils.visualize(im_path, pred, weight=0.0)
-            # pred_saved_path = os.path.join(pred_saved_dir, im_file)
-            # mkdir(pred_saved_path)
-            # cv2.imwrite(pred_saved_path, pred_im)
-
             progbar_pred.update(i + 1)