add files

app.py

@@ -0,0 +1,53 @@
+import requests
+import gradio as gr
+
+import paddle
+from paddleseg.cvlibs import Config
+
+from matting.core import predict
+from matting.model import *
+from matting.dataset import MattingDataset
+
+
+def download_file(http_address, file_name):
+    r = requests.get(http_address, allow_redirects=True)
+    open(file_name, 'wb').write(r.content)
+
+
+cfgs = ['configs/modnet/modnet_mobilenetv2.yml', 'configs/modnet/modnet_resnet50_vd.yml', 'configs/modnet/modnet_hrnet_w18.yml']
+
+download_file('https://paddleseg.bj.bcebos.com/matting/models/modnet-mobilenetv2.pdparams', 'modnet-mobilenetv2.pdparams')
+download_file('https://paddleseg.bj.bcebos.com/matting/models/modnet-resnet50_vd.pdparams', 'modnet-resnet50_vd.pdparams')
+download_file('https://paddleseg.bj.bcebos.com/matting/models/modnet-hrnet_w18.pdparams', 'modnet-hrnet_w18.pdparams')
+models_paths = ['modnet-mobilenetv2.pdparams', 'modnet-resnet50_vd.pdparams', 'modnet-hrnet_w18.pdparams']
+
+
+def inference(image, chosen_model):
+    paddle.set_device('cpu')
+    cfg = Config(cfgs[chosen_model])
+
+    val_dataset = cfg.val_dataset
+    model = cfg.model
+    img_transforms = val_dataset.transforms
+
+    alpha_pred = predict(model,
+                         model_path=models_paths[chosen_model],
+                         transforms=img_transforms,
+                         image_list=[image])
+
+    return alpha_pred
+
+
+inputs = [gr.inputs.Image(label='Input Image'),
+          gr.inputs.Radio(['MobileNetV2', 'ResNet50_vd', 'HRNet_W18'], label='Model', type='index')]
+
+
+gr.Interface(
+    inference, 
+    inputs,
+    gr.outputs.Image(label='Output'),
+    title='PaddleSeg - Matting',
+    examples=[['images/armchair.jpg', 'MobileNetV2'],
+              ['images/cat.jpg', 'ResNet50_vd'],
+              ['images/plant.jpg', 'HRNet_W18']]
+    ).launch()

configs/modnet/modnet_hrnet_w18.yml

@@ -0,0 +1,5 @@
+_base_: modnet_mobilenetv2.yml
+model:
+  backbone:
+    type: HRNet_W18
+    pretrained: https://bj.bcebos.com/paddleseg/dygraph/hrnet_w18_ssld.tar.gz

configs/modnet/modnet_mobilenetv2.yml

@@ -0,0 +1,49 @@
+batch_size: 16
+iters: 100000
+
+train_dataset:
+  type: MattingDataset
+  dataset_root: .
+  train_file: train.txt
+  transforms:
+    # - type: LoadImages
+    - type: ResizeByShort
+      short_size: 512
+    - type: RandomCrop
+      crop_size: [512, 512]
+    - type: RandomDistort
+    - type: RandomBlur
+    - type: RandomHorizontalFlip
+    - type: Normalize
+  mode: train
+
+val_dataset:
+  type: MattingDataset
+  dataset_root: .
+  val_file: val.txt
+  transforms:
+    # - type: LoadImages
+    - type: ResizeByShort
+      short_size: 512
+    - type: ResizeToIntMult
+      mult_int: 32
+    - type: Normalize
+  mode: val
+  get_trimap: False
+
+model:
+  type: MODNet
+  backbone:
+    type: MobileNetV2
+    pretrained: https://paddleseg.bj.bcebos.com/matting/models/MobileNetV2_pretrained/model.pdparams
+  pretrained: Null
+
+optimizer:
+  type: sgd
+  momentum: 0.9
+  weight_decay: 4.0e-5
+
+lr_scheduler:
+  type: PiecewiseDecay
+  boundaries: [40000, 80000]
+  values: [0.02, 0.002, 0.0002]

configs/modnet/modnet_resnet50_vd.yml

@@ -0,0 +1,5 @@
+_base_: modnet_mobilenetv2.yml
+model:
+  backbone:
+    type: ResNet50_vd
+    pretrained: https://bj.bcebos.com/paddleseg/dygraph/resnet50_vd_ssld_v2.tar.gz

matting/core/__init__.py

@@ -0,0 +1 @@
+from .predict import predict

matting/core/predict.py

@@ -0,0 +1,163 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import math
+import time
+
+import cv2
+import numpy as np
+import paddle
+import paddle.nn.functional as F
+from paddleseg import utils
+from paddleseg.core import infer
+from paddleseg.utils import logger, progbar, TimeAverager
+
+from matting.utils import mkdir
+
+
+def partition_list(arr, m):
+    """split the list 'arr' into m pieces"""
+    n = int(math.ceil(len(arr) / float(m)))
+    return [arr[i:i + n] for i in range(0, len(arr), n)]
+
+
+def save_alpha_pred(alpha, path, trimap=None):
+    """
+    The value of alpha is range [0, 1], shape should be [h,w]
+    """
+    dirname = os.path.dirname(path)
+    if not os.path.exists(dirname):
+        os.makedirs(dirname)
+
+    trimap = cv2.imread(trimap, 0)
+    alpha[trimap == 0] = 0
+    alpha[trimap == 255] = 255
+    alpha = (alpha).astype('uint8')
+    cv2.imwrite(path, alpha)
+
+
+def reverse_transform(alpha, trans_info):
+    """recover pred to origin shape"""
+    for item in trans_info[::-1]:
+        if item[0] == 'resize':
+            h, w = item[1][0], item[1][1]
+            alpha = F.interpolate(alpha, [h, w], mode='bilinear')
+        elif item[0] == 'padding':
+            h, w = item[1][0], item[1][1]
+            alpha = alpha[:, :, 0:h, 0:w]
+        else:
+            raise Exception("Unexpected info '{}' in im_info".format(item[0]))
+    return alpha
+
+
+def preprocess(img, transforms, trimap=None):
+    data = {}
+    data['img'] = img
+    if trimap is not None:
+        data['trimap'] = trimap
+        data['gt_fields'] = ['trimap']
+    data['trans_info'] = []
+    data = transforms(data)
+    data['img'] = paddle.to_tensor(data['img'])
+    data['img'] = data['img'].unsqueeze(0)
+    if trimap is not None:
+        data['trimap'] = paddle.to_tensor(data['trimap'])
+        data['trimap'] = data['trimap'].unsqueeze((0, 1))
+
+    return data
+
+
+def predict(model,
+            model_path,
+            transforms,
+            image_list,
+            image_dir=None,
+            trimap_list=None,
+            save_dir='output'):
+    """
+    predict and visualize the image_list.
+
+    Args:
+        model (nn.Layer): Used to predict for input image.
+        model_path (str): The path of pretrained model.
+        transforms (transforms.Compose): Preprocess for input image.
+        image_list (list): A list of image path to be predicted.
+        image_dir (str, optional): The root directory of the images predicted. Default: None.
+        trimap_list (list, optional): A list of trimap of image_list. Default: None.
+        save_dir (str, optional): The directory to save the visualized results. Default: 'output'.
+    """
+    utils.utils.load_entire_model(model, model_path)
+    model.eval()
+    nranks = paddle.distributed.get_world_size()
+    local_rank = paddle.distributed.get_rank()
+    if nranks > 1:
+        img_lists = partition_list(image_list, nranks)
+        trimap_lists = partition_list(
+            trimap_list, nranks) if trimap_list is not None else None
+    else:
+        img_lists = [image_list]
+        trimap_lists = [trimap_list] if trimap_list is not None else None
+
+    logger.info("Start to predict...")
+    progbar_pred = progbar.Progbar(target=len(img_lists[0]), verbose=1)
+    preprocess_cost_averager = TimeAverager()
+    infer_cost_averager = TimeAverager()
+    postprocess_cost_averager = TimeAverager()
+    batch_start = time.time()
+    with paddle.no_grad():
+        for i, im_path in enumerate(img_lists[local_rank]):
+            preprocess_start = time.time()
+            trimap = trimap_lists[local_rank][
+                i] if trimap_list is not None else None
+            data = preprocess(img=im_path, transforms=transforms, trimap=trimap)
+            preprocess_cost_averager.record(time.time() - preprocess_start)
+
+            infer_start = time.time()
+            alpha_pred = model(data)
+            infer_cost_averager.record(time.time() - infer_start)
+
+            postprocess_start = time.time()
+            alpha_pred = reverse_transform(alpha_pred, data['trans_info'])
+            alpha_pred = (alpha_pred.numpy()).squeeze()
+            alpha_pred = (alpha_pred * 255).astype('uint8')
+
+            # get the saved name
+            # if image_dir is not None:
+            #     im_file = im_path.replace(image_dir, '')
+            # else:
+            #     im_file = os.path.basename(im_path)
+            # if im_file[0] == '/' or im_file[0] == '\\':
+            #     im_file = im_file[1:]
+
+            # save_path = os.path.join(save_dir, im_file)
+            # mkdir(save_path)
+            # save_alpha_pred(alpha_pred, save_path, trimap=trimap)
+
+            postprocess_cost_averager.record(time.time() - postprocess_start)
+
+            preprocess_cost = preprocess_cost_averager.get_average()
+            infer_cost = infer_cost_averager.get_average()
+            postprocess_cost = postprocess_cost_averager.get_average()
+            if local_rank == 0:
+                progbar_pred.update(i + 1,
+                                    [('preprocess_cost', preprocess_cost),
+                                     ('infer_cost cost', infer_cost),
+                                     ('postprocess_cost', postprocess_cost)])
+
+            preprocess_cost_averager.reset()
+            infer_cost_averager.reset()
+            postprocess_cost_averager.reset()
+            
+    return alpha_pred

matting/dataset/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .matting_dataset import MattingDataset

matting/dataset/matting_dataset.py

@@ -0,0 +1,229 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import math
+
+import cv2
+import numpy as np
+import random
+import paddle
+from paddleseg.cvlibs import manager
+
+import matting.transforms as T
+
+
+@manager.DATASETS.add_component
+class MattingDataset(paddle.io.Dataset):
+    """
+    Pass in a dataset that conforms to the format.
+        matting_dataset/
+        |--bg/
+        |
+        |--train/
+        |  |--fg/
+        |  |--alpha/
+        |
+        |--val/
+        |  |--fg/
+        |  |--alpha/
+        |  |--trimap/ (if existing)
+        |
+        |--train.txt
+        |
+        |--val.txt
+    See README.md for more information of dataset.
+
+    Args:
+        dataset_root(str): The root path of dataset.
+        transforms(list):  Transforms for image.
+        mode (str, optional): which part of dataset to use. it is one of ('train', 'val', 'trainval'). Default: 'train'.
+        train_file (str|list, optional): File list is used to train. It should be `foreground_image.png background_image.png`
+            or `foreground_image.png`. It shold be provided if mode equal to 'train'. Default: None.
+        val_file (str|list, optional): File list is used to evaluation. It should be `foreground_image.png background_image.png`
+            or `foreground_image.png` or ``foreground_image.png background_image.png trimap_image.png`.
+            It shold be provided if mode equal to 'val'. Default: None.
+        get_trimap (bool, optional): Whether to get triamp. Default: True.
+        separator (str, optional): The separator of train_file or val_file. If file name contains ' ', '|' may be perfect. Default: ' '.
+    """
+
+    def __init__(self,
+                 dataset_root,
+                 transforms,
+                 mode='train',
+                 train_file=None,
+                 val_file=None,
+                 get_trimap=True,
+                 separator=' '):
+        super().__init__()
+        self.dataset_root = dataset_root
+        self.transforms = T.Compose(transforms)
+        self.mode = mode
+        self.get_trimap = get_trimap
+        self.separator = separator
+
+        # check file
+        if mode == 'train' or mode == 'trainval':
+            if train_file is None:
+                raise ValueError(
+                    "When `mode` is 'train' or 'trainval', `train_file must be provided!"
+                )
+            if isinstance(train_file, str):
+                train_file = [train_file]
+            file_list = train_file
+
+        if mode == 'val' or mode == 'trainval':
+            if val_file is None:
+                raise ValueError(
+                    "When `mode` is 'val' or 'trainval', `val_file must be provided!"
+                )
+            if isinstance(val_file, str):
+                val_file = [val_file]
+            file_list = val_file
+
+        if mode == 'trainval':
+            file_list = train_file + val_file
+
+        # read file
+        self.fg_bg_list = []
+        for file in file_list:
+            file = os.path.join(dataset_root, file)
+            with open(file, 'r') as f:
+                lines = f.readlines()
+                for line in lines:
+                    line = line.strip()
+                    self.fg_bg_list.append(line)
+
+    def __getitem__(self, idx):
+        data = {}
+        fg_bg_file = self.fg_bg_list[idx]
+        fg_bg_file = fg_bg_file.split(self.separator)
+        data['img_name'] = fg_bg_file[0]  # using in save prediction results
+        fg_file = os.path.join(self.dataset_root, fg_bg_file[0])
+        alpha_file = fg_file.replace('/fg', '/alpha')
+        fg = cv2.imread(fg_file)
+        alpha = cv2.imread(alpha_file, 0)
+        data['alpha'] = alpha
+        data['gt_fields'] = []
+
+        # line is: fg [bg] [trimap]
+        if len(fg_bg_file) >= 2:
+            bg_file = os.path.join(self.dataset_root, fg_bg_file[1])
+            bg = cv2.imread(bg_file)
+            data['img'], data['bg'] = self.composite(fg, alpha, bg)
+            data['fg'] = fg
+            if self.mode in ['train', 'trainval']:
+                data['gt_fields'].append('fg')
+                data['gt_fields'].append('bg')
+                data['gt_fields'].append('alpha')
+            if len(fg_bg_file) == 3 and self.get_trimap:
+                if self.mode == 'val':
+                    trimap_path = os.path.join(self.dataset_root, fg_bg_file[2])
+                    if os.path.exists(trimap_path):
+                        data['trimap'] = trimap_path
+                        data['gt_fields'].append('trimap')
+                        data['ori_trimap'] = cv2.imread(trimap_path, 0)
+                    else:
+                        raise FileNotFoundError(
+                            'trimap is not Found: {}'.format(fg_bg_file[2]))
+        else:
+            data['img'] = fg
+            if self.mode in ['train', 'trainval']:
+                data['fg'] = fg.copy()
+                data['bg'] = fg.copy()
+                data['gt_fields'].append('fg')
+                data['gt_fields'].append('bg')
+                data['gt_fields'].append('alpha')
+
+        data['trans_info'] = []  # Record shape change information
+
+        # Generate trimap from alpha if no trimap file provided
+        if self.get_trimap:
+            if 'trimap' not in data:
+                data['trimap'] = self.gen_trimap(
+                    data['alpha'], mode=self.mode).astype('float32')
+                data['gt_fields'].append('trimap')
+                if self.mode == 'val':
+                    data['ori_trimap'] = data['trimap'].copy()
+
+        data = self.transforms(data)
+
+        # When evaluation, gt should not be transforms.
+        if self.mode == 'val':
+            data['gt_fields'].append('alpha')
+
+        data['img'] = data['img'].astype('float32')
+        for key in data.get('gt_fields', []):
+            data[key] = data[key].astype('float32')
+
+        if 'trimap' in data:
+            data['trimap'] = data['trimap'][np.newaxis, :, :]
+        if 'ori_trimap' in data:
+            data['ori_trimap'] = data['ori_trimap'][np.newaxis, :, :]
+
+        data['alpha'] = data['alpha'][np.newaxis, :, :] / 255.
+
+        return data
+
+    def __len__(self):
+        return len(self.fg_bg_list)
+
+    def composite(self, fg, alpha, ori_bg):
+        fg_h, fg_w = fg.shape[:2]
+        ori_bg_h, ori_bg_w = ori_bg.shape[:2]
+
+        wratio = fg_w / ori_bg_w
+        hratio = fg_h / ori_bg_h
+        ratio = wratio if wratio > hratio else hratio
+
+        # Resize ori_bg if it is smaller than fg.
+        if ratio > 1:
+            resize_h = math.ceil(ori_bg_h * ratio)
+            resize_w = math.ceil(ori_bg_w * ratio)
+            bg = cv2.resize(
+                ori_bg, (resize_w, resize_h), interpolation=cv2.INTER_LINEAR)
+        else:
+            bg = ori_bg
+
+        bg = bg[0:fg_h, 0:fg_w, :]
+        alpha = alpha / 255
+        alpha = np.expand_dims(alpha, axis=2)
+        image = alpha * fg + (1 - alpha) * bg
+        image = image.astype(np.uint8)
+        return image, bg
+
+    @staticmethod
+    def gen_trimap(alpha, mode='train', eval_kernel=7):
+        if mode == 'train':
+            k_size = random.choice(range(2, 5))
+            iterations = np.random.randint(5, 15)
+            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
+                                               (k_size, k_size))
+            dilated = cv2.dilate(alpha, kernel, iterations=iterations)
+            eroded = cv2.erode(alpha, kernel, iterations=iterations)
+            trimap = np.zeros(alpha.shape)
+            trimap.fill(128)
+            trimap[eroded > 254.5] = 255
+            trimap[dilated < 0.5] = 0
+        else:
+            k_size = eval_kernel
+            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
+                                               (k_size, k_size))
+            dilated = cv2.dilate(alpha, kernel)
+            trimap = np.zeros(alpha.shape)
+            trimap.fill(128)
+            trimap[alpha >= 250] = 255
+            trimap[dilated <= 5] = 0
+
+        return trimap

matting/model/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .vgg import *
+from .resnet_vd import *
+from .mobilenet_v2 import *
+from .hrnet import *
+from .dim import DIM
+from .loss import MRSD
+from .modnet import MODNet

matting/model/dim.py

@@ -0,0 +1,203 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import defaultdict
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddleseg.models import layers
+from paddleseg import utils
+from paddleseg.cvlibs import manager
+
+from .loss import MRSD
+
+
+@manager.MODELS.add_component
+class DIM(nn.Layer):
+    """
+    The DIM implementation based on PaddlePaddle.
+
+    The original article refers to
+    Ning Xu, et, al. "Deep Image Matting"
+    (https://arxiv.org/pdf/1908.07919.pdf).
+
+    Args:
+        backbone: backbone model.
+        stage (int, optional): The stage of model. Defautl: 3.
+        decoder_input_channels(int, optional): The channel of decoder input. Default: 512.
+        pretrained(str, optional): The path of pretrianed model. Defautl: None.
+
+    """
+
+    def __init__(self,
+                 backbone,
+                 stage=3,
+                 decoder_input_channels=512,
+                 pretrained=None):
+        super().__init__()
+        self.backbone = backbone
+        self.pretrained = pretrained
+        self.stage = stage
+
+        decoder_output_channels = [64, 128, 256, 512]
+        self.decoder = Decoder(
+            input_channels=decoder_input_channels,
+            output_channels=decoder_output_channels)
+        if self.stage == 2:
+            for param in self.backbone.parameters():
+                param.stop_gradient = True
+            for param in self.decoder.parameters():
+                param.stop_gradient = True
+        if self.stage >= 2:
+            self.refine = Refine()
+        self.init_weight()
+
+    def forward(self, inputs):
+        input_shape = paddle.shape(inputs['img'])[-2:]
+        x = paddle.concat([inputs['img'], inputs['trimap'] / 255], axis=1)
+        fea_list = self.backbone(x)
+
+        # decoder stage
+        up_shape = []
+        for i in range(5):
+            up_shape.append(paddle.shape(fea_list[i])[-2:])
+        alpha_raw = self.decoder(fea_list, up_shape)
+        alpha_raw = F.interpolate(
+            alpha_raw, input_shape, mode='bilinear', align_corners=False)
+        logit_dict = {'alpha_raw': alpha_raw}
+        if self.stage < 2:
+            return logit_dict
+
+        if self.stage >= 2:
+            # refine stage
+            refine_input = paddle.concat([inputs['img'], alpha_raw], axis=1)
+            alpha_refine = self.refine(refine_input)
+
+            # finally alpha
+            alpha_pred = alpha_refine + alpha_raw
+            alpha_pred = F.interpolate(
+                alpha_pred, input_shape, mode='bilinear', align_corners=False)
+            if not self.training:
+                alpha_pred = paddle.clip(alpha_pred, min=0, max=1)
+            logit_dict['alpha_pred'] = alpha_pred
+        if self.training:
+            return logit_dict
+        else:
+            return alpha_pred
+
+    def loss(self, logit_dict, label_dict, loss_func_dict=None):
+        if loss_func_dict is None:
+            loss_func_dict = defaultdict(list)
+            loss_func_dict['alpha_raw'].append(MRSD())
+            loss_func_dict['comp'].append(MRSD())
+            loss_func_dict['alpha_pred'].append(MRSD())
+
+        loss = {}
+        mask = label_dict['trimap'] == 128
+        loss['all'] = 0
+
+        if self.stage != 2:
+            loss['alpha_raw'] = loss_func_dict['alpha_raw'][0](
+                logit_dict['alpha_raw'], label_dict['alpha'], mask)
+            loss['alpha_raw'] = 0.5 * loss['alpha_raw']
+            loss['all'] = loss['all'] + loss['alpha_raw']
+
+        if self.stage == 1 or self.stage == 3:
+            comp_pred = logit_dict['alpha_raw'] * label_dict['fg'] + \
+                (1 - logit_dict['alpha_raw']) * label_dict['bg']
+            loss['comp'] = loss_func_dict['comp'][0](comp_pred,
+                                                     label_dict['img'], mask)
+            loss['comp'] = 0.5 * loss['comp']
+            loss['all'] = loss['all'] + loss['comp']
+
+        if self.stage == 2 or self.stage == 3:
+            loss['alpha_pred'] = loss_func_dict['alpha_pred'][0](
+                logit_dict['alpha_pred'], label_dict['alpha'], mask)
+            loss['all'] = loss['all'] + loss['alpha_pred']
+
+        return loss
+
+    def init_weight(self):
+        if self.pretrained is not None:
+            utils.load_entire_model(self, self.pretrained)
+
+
+# bilinear interpolate skip connect
+class Up(nn.Layer):
+    def __init__(self, input_channels, output_channels):
+        super().__init__()
+        self.conv = layers.ConvBNReLU(
+            input_channels,
+            output_channels,
+            kernel_size=5,
+            padding=2,
+            bias_attr=False)
+
+    def forward(self, x, skip, output_shape):
+        x = F.interpolate(
+            x, size=output_shape, mode='bilinear', align_corners=False)
+        x = x + skip
+        x = self.conv(x)
+        x = F.relu(x)
+
+        return x
+
+
+class Decoder(nn.Layer):
+    def __init__(self, input_channels, output_channels=(64, 128, 256, 512)):
+        super().__init__()
+        self.deconv6 = nn.Conv2D(
+            input_channels, input_channels, kernel_size=1, bias_attr=False)
+        self.deconv5 = Up(input_channels, output_channels[-1])
+        self.deconv4 = Up(output_channels[-1], output_channels[-2])
+        self.deconv3 = Up(output_channels[-2], output_channels[-3])
+        self.deconv2 = Up(output_channels[-3], output_channels[-4])
+        self.deconv1 = Up(output_channels[-4], 64)
+
+        self.alpha_conv = nn.Conv2D(
+            64, 1, kernel_size=5, padding=2, bias_attr=False)
+
+    def forward(self, fea_list, shape_list):
+        x = fea_list[-1]
+        x = self.deconv6(x)
+        x = self.deconv5(x, fea_list[4], shape_list[4])
+        x = self.deconv4(x, fea_list[3], shape_list[3])
+        x = self.deconv3(x, fea_list[2], shape_list[2])
+        x = self.deconv2(x, fea_list[1], shape_list[1])
+        x = self.deconv1(x, fea_list[0], shape_list[0])
+        alpha = self.alpha_conv(x)
+        alpha = F.sigmoid(alpha)
+
+        return alpha
+
+
+class Refine(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = layers.ConvBNReLU(
+            4, 64, kernel_size=3, padding=1, bias_attr=False)
+        self.conv2 = layers.ConvBNReLU(
+            64, 64, kernel_size=3, padding=1, bias_attr=False)
+        self.conv3 = layers.ConvBNReLU(
+            64, 64, kernel_size=3, padding=1, bias_attr=False)
+        self.alpha_pred = layers.ConvBNReLU(
+            64, 1, kernel_size=3, padding=1, bias_attr=False)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        alpha = self.alpha_pred(x)
+
+        return alpha

matting/model/hrnet.py

@@ -0,0 +1,835 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from paddleseg.cvlibs import manager, param_init
+from paddleseg.models import layers
+from paddleseg.utils import utils
+
+__all__ = [
+    "HRNet_W18_Small_V1", "HRNet_W18_Small_V2", "HRNet_W18", "HRNet_W30",
+    "HRNet_W32", "HRNet_W40", "HRNet_W44", "HRNet_W48", "HRNet_W60", "HRNet_W64"
+]
+
+
+class HRNet(nn.Layer):
+    """
+    The HRNet implementation based on PaddlePaddle.
+
+    The original article refers to
+    Jingdong Wang, et, al. "HRNet：Deep High-Resolution Representation Learning for Visual Recognition"
+    (https://arxiv.org/pdf/1908.07919.pdf).
+
+    Args:
+        pretrained (str, optional): The path of pretrained model.
+        stage1_num_modules (int, optional): Number of modules for stage1. Default 1.
+        stage1_num_blocks (list, optional): Number of blocks per module for stage1. Default (4).
+        stage1_num_channels (list, optional): Number of channels per branch for stage1. Default (64).
+        stage2_num_modules (int, optional): Number of modules for stage2. Default 1.
+        stage2_num_blocks (list, optional): Number of blocks per module for stage2. Default (4, 4).
+        stage2_num_channels (list, optional): Number of channels per branch for stage2. Default (18, 36).
+        stage3_num_modules (int, optional): Number of modules for stage3. Default 4.
+        stage3_num_blocks (list, optional): Number of blocks per module for stage3. Default (4, 4, 4).
+        stage3_num_channels (list, optional): Number of channels per branch for stage3. Default [18, 36, 72).
+        stage4_num_modules (int, optional): Number of modules for stage4. Default 3.
+        stage4_num_blocks (list, optional): Number of blocks per module for stage4. Default (4, 4, 4, 4).
+        stage4_num_channels (list, optional): Number of channels per branch for stage4. Default (18, 36, 72. 144).
+        has_se (bool, optional): Whether to use Squeeze-and-Excitation module. Default False.
+        align_corners (bool, optional): An argument of F.interpolate. It should be set to False when the feature size is even,
+            e.g. 1024x512, otherwise it is True, e.g. 769x769. Default: False.
+    """
+
+    def __init__(self,
+                 input_channels=3,
+                 pretrained=None,
+                 stage1_num_modules=1,
+                 stage1_num_blocks=(4, ),
+                 stage1_num_channels=(64, ),
+                 stage2_num_modules=1,
+                 stage2_num_blocks=(4, 4),
+                 stage2_num_channels=(18, 36),
+                 stage3_num_modules=4,
+                 stage3_num_blocks=(4, 4, 4),
+                 stage3_num_channels=(18, 36, 72),
+                 stage4_num_modules=3,
+                 stage4_num_blocks=(4, 4, 4, 4),
+                 stage4_num_channels=(18, 36, 72, 144),
+                 has_se=False,
+                 align_corners=False,
+                 padding_same=True):
+        super(HRNet, self).__init__()
+        self.pretrained = pretrained
+        self.stage1_num_modules = stage1_num_modules
+        self.stage1_num_blocks = stage1_num_blocks
+        self.stage1_num_channels = stage1_num_channels
+        self.stage2_num_modules = stage2_num_modules
+        self.stage2_num_blocks = stage2_num_blocks
+        self.stage2_num_channels = stage2_num_channels
+        self.stage3_num_modules = stage3_num_modules
+        self.stage3_num_blocks = stage3_num_blocks
+        self.stage3_num_channels = stage3_num_channels
+        self.stage4_num_modules = stage4_num_modules
+        self.stage4_num_blocks = stage4_num_blocks
+        self.stage4_num_channels = stage4_num_channels
+        self.has_se = has_se
+        self.align_corners = align_corners
+
+        self.feat_channels = [i for i in stage4_num_channels]
+        self.feat_channels = [64] + self.feat_channels
+
+        self.conv_layer1_1 = layers.ConvBNReLU(
+            in_channels=input_channels,
+            out_channels=64,
+            kernel_size=3,
+            stride=2,
+            padding=1 if not padding_same else 'same',
+            bias_attr=False)
+
+        self.conv_layer1_2 = layers.ConvBNReLU(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=2,
+            padding=1 if not padding_same else 'same',
+            bias_attr=False)
+
+        self.la1 = Layer1(
+            num_channels=64,
+            num_blocks=self.stage1_num_blocks[0],
+            num_filters=self.stage1_num_channels[0],
+            has_se=has_se,
+            name="layer2",
+            padding_same=padding_same)
+
+        self.tr1 = TransitionLayer(
+            in_channels=[self.stage1_num_channels[0] * 4],
+            out_channels=self.stage2_num_channels,
+            name="tr1",
+            padding_same=padding_same)
+
+        self.st2 = Stage(
+            num_channels=self.stage2_num_channels,
+            num_modules=self.stage2_num_modules,
+            num_blocks=self.stage2_num_blocks,
+            num_filters=self.stage2_num_channels,
+            has_se=self.has_se,
+            name="st2",
+            align_corners=align_corners,
+            padding_same=padding_same)
+
+        self.tr2 = TransitionLayer(
+            in_channels=self.stage2_num_channels,
+            out_channels=self.stage3_num_channels,
+            name="tr2",
+            padding_same=padding_same)
+        self.st3 = Stage(
+            num_channels=self.stage3_num_channels,
+            num_modules=self.stage3_num_modules,
+            num_blocks=self.stage3_num_blocks,
+            num_filters=self.stage3_num_channels,
+            has_se=self.has_se,
+            name="st3",
+            align_corners=align_corners,
+            padding_same=padding_same)
+
+        self.tr3 = TransitionLayer(
+            in_channels=self.stage3_num_channels,
+            out_channels=self.stage4_num_channels,
+            name="tr3",
+            padding_same=padding_same)
+        self.st4 = Stage(
+            num_channels=self.stage4_num_channels,
+            num_modules=self.stage4_num_modules,
+            num_blocks=self.stage4_num_blocks,
+            num_filters=self.stage4_num_channels,
+            has_se=self.has_se,
+            name="st4",
+            align_corners=align_corners,
+            padding_same=padding_same)
+
+        self.init_weight()
+
+    def forward(self, x):
+        feat_list = []
+        conv1 = self.conv_layer1_1(x)
+        feat_list.append(conv1)
+        conv2 = self.conv_layer1_2(conv1)
+
+        la1 = self.la1(conv2)
+
+        tr1 = self.tr1([la1])
+        st2 = self.st2(tr1)
+
+        tr2 = self.tr2(st2)
+        st3 = self.st3(tr2)
+
+        tr3 = self.tr3(st3)
+        st4 = self.st4(tr3)
+
+        feat_list = feat_list + st4
+
+        return feat_list
+
+    def init_weight(self):
+        for layer in self.sublayers():
+            if isinstance(layer, nn.Conv2D):
+                param_init.normal_init(layer.weight, std=0.001)
+            elif isinstance(layer, (nn.BatchNorm, nn.SyncBatchNorm)):
+                param_init.constant_init(layer.weight, value=1.0)
+                param_init.constant_init(layer.bias, value=0.0)
+        if self.pretrained is not None:
+            utils.load_pretrained_model(self, self.pretrained)
+
+
+class Layer1(nn.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_filters,
+                 num_blocks,
+                 has_se=False,
+                 name=None,
+                 padding_same=True):
+        super(Layer1, self).__init__()
+
+        self.bottleneck_block_list = []
+
+        for i in range(num_blocks):
+            bottleneck_block = self.add_sublayer(
+                "bb_{}_{}".format(name, i + 1),
+                BottleneckBlock(
+                    num_channels=num_channels if i == 0 else num_filters * 4,
+                    num_filters=num_filters,
+                    has_se=has_se,
+                    stride=1,
+                    downsample=True if i == 0 else False,
+                    name=name + '_' + str(i + 1),
+                    padding_same=padding_same))
+            self.bottleneck_block_list.append(bottleneck_block)
+
+    def forward(self, x):
+        conv = x
+        for block_func in self.bottleneck_block_list:
+            conv = block_func(conv)
+        return conv
+
+
+class TransitionLayer(nn.Layer):
+    def __init__(self, in_channels, out_channels, name=None, padding_same=True):
+        super(TransitionLayer, self).__init__()
+
+        num_in = len(in_channels)
+        num_out = len(out_channels)
+        self.conv_bn_func_list = []
+        for i in range(num_out):
+            residual = None
+            if i < num_in:
+                if in_channels[i] != out_channels[i]:
+                    residual = self.add_sublayer(
+                        "transition_{}_layer_{}".format(name, i + 1),
+                        layers.ConvBNReLU(
+                            in_channels=in_channels[i],
+                            out_channels=out_channels[i],
+                            kernel_size=3,
+                            padding=1 if not padding_same else 'same',
+                            bias_attr=False))
+            else:
+                residual = self.add_sublayer(
+                    "transition_{}_layer_{}".format(name, i + 1),
+                    layers.ConvBNReLU(
+                        in_channels=in_channels[-1],
+                        out_channels=out_channels[i],
+                        kernel_size=3,
+                        stride=2,
+                        padding=1 if not padding_same else 'same',
+                        bias_attr=False))
+            self.conv_bn_func_list.append(residual)
+
+    def forward(self, x):
+        outs = []
+        for idx, conv_bn_func in enumerate(self.conv_bn_func_list):
+            if conv_bn_func is None:
+                outs.append(x[idx])
+            else:
+                if idx < len(x):
+                    outs.append(conv_bn_func(x[idx]))
+                else:
+                    outs.append(conv_bn_func(x[-1]))
+        return outs
+
+
+class Branches(nn.Layer):
+    def __init__(self,
+                 num_blocks,
+                 in_channels,
+                 out_channels,
+                 has_se=False,
+                 name=None,
+                 padding_same=True):
+        super(Branches, self).__init__()
+
+        self.basic_block_list = []
+
+        for i in range(len(out_channels)):
+            self.basic_block_list.append([])
+            for j in range(num_blocks[i]):
+                in_ch = in_channels[i] if j == 0 else out_channels[i]
+                basic_block_func = self.add_sublayer(
+                    "bb_{}_branch_layer_{}_{}".format(name, i + 1, j + 1),
+                    BasicBlock(
+                        num_channels=in_ch,
+                        num_filters=out_channels[i],
+                        has_se=has_se,
+                        name=name + '_branch_layer_' + str(i + 1) + '_' +
+                        str(j + 1),
+                        padding_same=padding_same))
+                self.basic_block_list[i].append(basic_block_func)
+
+    def forward(self, x):
+        outs = []
+        for idx, input in enumerate(x):
+            conv = input
+            for basic_block_func in self.basic_block_list[idx]:
+                conv = basic_block_func(conv)
+            outs.append(conv)
+        return outs
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_filters,
+                 has_se,
+                 stride=1,
+                 downsample=False,
+                 name=None,
+                 padding_same=True):
+        super(BottleneckBlock, self).__init__()
+
+        self.has_se = has_se
+        self.downsample = downsample
+
+        self.conv1 = layers.ConvBNReLU(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=1,
+            bias_attr=False)
+
+        self.conv2 = layers.ConvBNReLU(
+            in_channels=num_filters,
+            out_channels=num_filters,
+            kernel_size=3,
+            stride=stride,
+            padding=1 if not padding_same else 'same',
+            bias_attr=False)
+
+        self.conv3 = layers.ConvBN(
+            in_channels=num_filters,
+            out_channels=num_filters * 4,
+            kernel_size=1,
+            bias_attr=False)
+
+        if self.downsample:
+            self.conv_down = layers.ConvBN(
+                in_channels=num_channels,
+                out_channels=num_filters * 4,
+                kernel_size=1,
+                bias_attr=False)
+
+        if self.has_se:
+            self.se = SELayer(
+                num_channels=num_filters * 4,
+                num_filters=num_filters * 4,
+                reduction_ratio=16,
+                name=name + '_fc')
+
+        self.add = layers.Add()
+        self.relu = layers.Activation("relu")
+
+    def forward(self, x):
+        residual = x
+        conv1 = self.conv1(x)
+        conv2 = self.conv2(conv1)
+        conv3 = self.conv3(conv2)
+
+        if self.downsample:
+            residual = self.conv_down(x)
+
+        if self.has_se:
+            conv3 = self.se(conv3)
+
+        y = self.add(conv3, residual)
+        y = self.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_filters,
+                 stride=1,
+                 has_se=False,
+                 downsample=False,
+                 name=None,
+                 padding_same=True):
+        super(BasicBlock, self).__init__()
+
+        self.has_se = has_se
+        self.downsample = downsample
+
+        self.conv1 = layers.ConvBNReLU(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=3,
+            stride=stride,
+            padding=1 if not padding_same else 'same',
+            bias_attr=False)
+        self.conv2 = layers.ConvBN(
+            in_channels=num_filters,
+            out_channels=num_filters,
+            kernel_size=3,
+            padding=1 if not padding_same else 'same',
+            bias_attr=False)
+
+        if self.downsample:
+            self.conv_down = layers.ConvBNReLU(
+                in_channels=num_channels,
+                out_channels=num_filters,
+                kernel_size=1,
+                bias_attr=False)
+
+        if self.has_se:
+            self.se = SELayer(
+                num_channels=num_filters,
+                num_filters=num_filters,
+                reduction_ratio=16,
+                name=name + '_fc')
+
+        self.add = layers.Add()
+        self.relu = layers.Activation("relu")
+
+    def forward(self, x):
+        residual = x
+        conv1 = self.conv1(x)
+        conv2 = self.conv2(conv1)
+
+        if self.downsample:
+            residual = self.conv_down(x)
+
+        if self.has_se:
+            conv2 = self.se(conv2)
+
+        y = self.add(conv2, residual)
+        y = self.relu(y)
+        return y
+
+
+class SELayer(nn.Layer):
+    def __init__(self, num_channels, num_filters, reduction_ratio, name=None):
+        super(SELayer, self).__init__()
+
+        self.pool2d_gap = nn.AdaptiveAvgPool2D(1)
+
+        self._num_channels = num_channels
+
+        med_ch = int(num_channels / reduction_ratio)
+        stdv = 1.0 / math.sqrt(num_channels * 1.0)
+        self.squeeze = nn.Linear(
+            num_channels,
+            med_ch,
+            weight_attr=paddle.ParamAttr(
+                initializer=nn.initializer.Uniform(-stdv, stdv)))
+
+        stdv = 1.0 / math.sqrt(med_ch * 1.0)
+        self.excitation = nn.Linear(
+            med_ch,
+            num_filters,
+            weight_attr=paddle.ParamAttr(
+                initializer=nn.initializer.Uniform(-stdv, stdv)))
+
+    def forward(self, x):
+        pool = self.pool2d_gap(x)
+        pool = paddle.reshape(pool, shape=[-1, self._num_channels])
+        squeeze = self.squeeze(pool)
+        squeeze = F.relu(squeeze)
+        excitation = self.excitation(squeeze)
+        excitation = F.sigmoid(excitation)
+        excitation = paddle.reshape(
+            excitation, shape=[-1, self._num_channels, 1, 1])
+        out = x * excitation
+        return out
+
+
+class Stage(nn.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_modules,
+                 num_blocks,
+                 num_filters,
+                 has_se=False,
+                 multi_scale_output=True,
+                 name=None,
+                 align_corners=False,
+                 padding_same=True):
+        super(Stage, self).__init__()
+
+        self._num_modules = num_modules
+
+        self.stage_func_list = []
+        for i in range(num_modules):
+            if i == num_modules - 1 and not multi_scale_output:
+                stage_func = self.add_sublayer(
+                    "stage_{}_{}".format(name, i + 1),
+                    HighResolutionModule(
+                        num_channels=num_channels,
+                        num_blocks=num_blocks,
+                        num_filters=num_filters,
+                        has_se=has_se,
+                        multi_scale_output=False,
+                        name=name + '_' + str(i + 1),
+                        align_corners=align_corners,
+                        padding_same=padding_same))
+            else:
+                stage_func = self.add_sublayer(
+                    "stage_{}_{}".format(name, i + 1),
+                    HighResolutionModule(
+                        num_channels=num_channels,
+                        num_blocks=num_blocks,
+                        num_filters=num_filters,
+                        has_se=has_se,
+                        name=name + '_' + str(i + 1),
+                        align_corners=align_corners,
+                        padding_same=padding_same))
+
+            self.stage_func_list.append(stage_func)
+
+    def forward(self, x):
+        out = x
+        for idx in range(self._num_modules):
+            out = self.stage_func_list[idx](out)
+        return out
+
+
+class HighResolutionModule(nn.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_blocks,
+                 num_filters,
+                 has_se=False,
+                 multi_scale_output=True,
+                 name=None,
+                 align_corners=False,
+                 padding_same=True):
+        super(HighResolutionModule, self).__init__()
+
+        self.branches_func = Branches(
+            num_blocks=num_blocks,
+            in_channels=num_channels,
+            out_channels=num_filters,
+            has_se=has_se,
+            name=name,
+            padding_same=padding_same)
+
+        self.fuse_func = FuseLayers(
+            in_channels=num_filters,
+            out_channels=num_filters,
+            multi_scale_output=multi_scale_output,
+            name=name,
+            align_corners=align_corners,
+            padding_same=padding_same)
+
+    def forward(self, x):
+        out = self.branches_func(x)
+        out = self.fuse_func(out)
+        return out
+
+
+class FuseLayers(nn.Layer):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 multi_scale_output=True,
+                 name=None,
+                 align_corners=False,
+                 padding_same=True):
+        super(FuseLayers, self).__init__()
+
+        self._actual_ch = len(in_channels) if multi_scale_output else 1
+        self._in_channels = in_channels
+        self.align_corners = align_corners
+
+        self.residual_func_list = []
+        for i in range(self._actual_ch):
+            for j in range(len(in_channels)):
+                if j > i:
+                    residual_func = self.add_sublayer(
+                        "residual_{}_layer_{}_{}".format(name, i + 1, j + 1),
+                        layers.ConvBN(
+                            in_channels=in_channels[j],
+                            out_channels=out_channels[i],
+                            kernel_size=1,
+                            bias_attr=False))
+                    self.residual_func_list.append(residual_func)
+                elif j < i:
+                    pre_num_filters = in_channels[j]
+                    for k in range(i - j):
+                        if k == i - j - 1:
+                            residual_func = self.add_sublayer(
+                                "residual_{}_layer_{}_{}_{}".format(
+                                    name, i + 1, j + 1, k + 1),
+                                layers.ConvBN(
+                                    in_channels=pre_num_filters,
+                                    out_channels=out_channels[i],
+                                    kernel_size=3,
+                                    stride=2,
+                                    padding=1 if not padding_same else 'same',
+                                    bias_attr=False))
+                            pre_num_filters = out_channels[i]
+                        else:
+                            residual_func = self.add_sublayer(
+                                "residual_{}_layer_{}_{}_{}".format(
+                                    name, i + 1, j + 1, k + 1),
+                                layers.ConvBNReLU(
+                                    in_channels=pre_num_filters,
+                                    out_channels=out_channels[j],
+                                    kernel_size=3,
+                                    stride=2,
+                                    padding=1 if not padding_same else 'same',
+                                    bias_attr=False))
+                            pre_num_filters = out_channels[j]
+                        self.residual_func_list.append(residual_func)
+
+    def forward(self, x):
+        outs = []
+        residual_func_idx = 0
+        for i in range(self._actual_ch):
+            residual = x[i]
+            residual_shape = paddle.shape(residual)[-2:]
+            for j in range(len(self._in_channels)):
+                if j > i:
+                    y = self.residual_func_list[residual_func_idx](x[j])
+                    residual_func_idx += 1
+
+                    y = F.interpolate(
+                        y,
+                        residual_shape,
+                        mode='bilinear',
+                        align_corners=self.align_corners)
+                    residual = residual + y
+                elif j < i:
+                    y = x[j]
+                    for k in range(i - j):
+                        y = self.residual_func_list[residual_func_idx](y)
+                        residual_func_idx += 1
+
+                    residual = residual + y
+
+            residual = F.relu(residual)
+            outs.append(residual)
+
+        return outs
+
+
+@manager.BACKBONES.add_component
+def HRNet_W18_Small_V1(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[1],
+        stage1_num_channels=[32],
+        stage2_num_modules=1,
+        stage2_num_blocks=[2, 2],
+        stage2_num_channels=[16, 32],
+        stage3_num_modules=1,
+        stage3_num_blocks=[2, 2, 2],
+        stage3_num_channels=[16, 32, 64],
+        stage4_num_modules=1,
+        stage4_num_blocks=[2, 2, 2, 2],
+        stage4_num_channels=[16, 32, 64, 128],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W18_Small_V2(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[2],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[2, 2],
+        stage2_num_channels=[18, 36],
+        stage3_num_modules=3,
+        stage3_num_blocks=[2, 2, 2],
+        stage3_num_channels=[18, 36, 72],
+        stage4_num_modules=2,
+        stage4_num_blocks=[2, 2, 2, 2],
+        stage4_num_channels=[18, 36, 72, 144],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W18(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[4],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[4, 4],
+        stage2_num_channels=[18, 36],
+        stage3_num_modules=4,
+        stage3_num_blocks=[4, 4, 4],
+        stage3_num_channels=[18, 36, 72],
+        stage4_num_modules=3,
+        stage4_num_blocks=[4, 4, 4, 4],
+        stage4_num_channels=[18, 36, 72, 144],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W30(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[4],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[4, 4],
+        stage2_num_channels=[30, 60],
+        stage3_num_modules=4,
+        stage3_num_blocks=[4, 4, 4],
+        stage3_num_channels=[30, 60, 120],
+        stage4_num_modules=3,
+        stage4_num_blocks=[4, 4, 4, 4],
+        stage4_num_channels=[30, 60, 120, 240],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W32(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[4],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[4, 4],
+        stage2_num_channels=[32, 64],
+        stage3_num_modules=4,
+        stage3_num_blocks=[4, 4, 4],
+        stage3_num_channels=[32, 64, 128],
+        stage4_num_modules=3,
+        stage4_num_blocks=[4, 4, 4, 4],
+        stage4_num_channels=[32, 64, 128, 256],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W40(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[4],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[4, 4],
+        stage2_num_channels=[40, 80],
+        stage3_num_modules=4,
+        stage3_num_blocks=[4, 4, 4],
+        stage3_num_channels=[40, 80, 160],
+        stage4_num_modules=3,
+        stage4_num_blocks=[4, 4, 4, 4],
+        stage4_num_channels=[40, 80, 160, 320],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W44(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[4],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[4, 4],
+        stage2_num_channels=[44, 88],
+        stage3_num_modules=4,
+        stage3_num_blocks=[4, 4, 4],
+        stage3_num_channels=[44, 88, 176],
+        stage4_num_modules=3,
+        stage4_num_blocks=[4, 4, 4, 4],
+        stage4_num_channels=[44, 88, 176, 352],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W48(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[4],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[4, 4],
+        stage2_num_channels=[48, 96],
+        stage3_num_modules=4,
+        stage3_num_blocks=[4, 4, 4],
+        stage3_num_channels=[48, 96, 192],
+        stage4_num_modules=3,
+        stage4_num_blocks=[4, 4, 4, 4],
+        stage4_num_channels=[48, 96, 192, 384],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W60(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[4],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[4, 4],
+        stage2_num_channels=[60, 120],
+        stage3_num_modules=4,
+        stage3_num_blocks=[4, 4, 4],
+        stage3_num_channels=[60, 120, 240],
+        stage4_num_modules=3,
+        stage4_num_blocks=[4, 4, 4, 4],
+        stage4_num_channels=[60, 120, 240, 480],
+        **kwargs)
+    return model
+
+
+@manager.BACKBONES.add_component
+def HRNet_W64(**kwargs):
+    model = HRNet(
+        stage1_num_modules=1,
+        stage1_num_blocks=[4],
+        stage1_num_channels=[64],
+        stage2_num_modules=1,
+        stage2_num_blocks=[4, 4],
+        stage2_num_channels=[64, 128],
+        stage3_num_modules=4,
+        stage3_num_blocks=[4, 4, 4],
+        stage3_num_channels=[64, 128, 256],
+        stage4_num_modules=3,
+        stage4_num_blocks=[4, 4, 4, 4],
+        stage4_num_channels=[64, 128, 256, 512],
+        **kwargs)
+    return model

matting/model/loss.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from paddleseg.cvlibs import manager
+
+
+@manager.LOSSES.add_component
+class MRSD(nn.Layer):
+    def __init__(self, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, logit, label, mask=None):
+        """
+        Forward computation.
+
+        Args:
+            logit (Tensor): Logit tensor, the data type is float32, float64.
+            label (Tensor): Label tensor, the data type is float32, float64. The shape should equal to logit.
+            mask (Tensor, optional): The mask where the loss valid. Default： None.
+        """
+        if len(label.shape) == 3:
+            label = label.unsqueeze(1)
+        sd = paddle.square(logit - label)
+        loss = paddle.sqrt(sd + self.eps)
+        if mask is not None:
+            mask = mask.astype('float32')
+            if len(mask.shape) == 3:
+                mask = mask.unsqueeze(1)
+            loss = loss * mask
+            loss = loss.sum() / (mask.sum() + self.eps)
+            mask.stop_gradient = True
+        else:
+            loss = loss.mean()
+
+        return loss

matting/model/mobilenet_v2.py

@@ -0,0 +1,241 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import numpy as np
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.nn import Conv2D, BatchNorm, Linear, Dropout
+from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D
+
+from paddleseg import utils
+from paddleseg.cvlibs import manager
+
+MODEL_URLS = {
+    "MobileNetV2_x0_25":
+    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV2_x0_25_pretrained.pdparams",
+    "MobileNetV2_x0_5":
+    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV2_x0_5_pretrained.pdparams",
+    "MobileNetV2_x0_75":
+    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV2_x0_75_pretrained.pdparams",
+    "MobileNetV2":
+    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV2_pretrained.pdparams",
+    "MobileNetV2_x1_5":
+    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV2_x1_5_pretrained.pdparams",
+    "MobileNetV2_x2_0":
+    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV2_x2_0_pretrained.pdparams"
+}
+
+__all__ = ["MobileNetV2"]
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(self,
+                 num_channels,
+                 filter_size,
+                 num_filters,
+                 stride,
+                 padding,
+                 channels=None,
+                 num_groups=1,
+                 name=None,
+                 use_cudnn=True):
+        super(ConvBNLayer, self).__init__()
+
+        self._conv = Conv2D(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=padding,
+            groups=num_groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False)
+
+        self._batch_norm = BatchNorm(
+            num_filters,
+            param_attr=ParamAttr(name=name + "_bn_scale"),
+            bias_attr=ParamAttr(name=name + "_bn_offset"),
+            moving_mean_name=name + "_bn_mean",
+            moving_variance_name=name + "_bn_variance")
+
+    def forward(self, inputs, if_act=True):
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        if if_act:
+            y = F.relu6(y)
+        return y
+
+
+class InvertedResidualUnit(nn.Layer):
+    def __init__(self, num_channels, num_in_filter, num_filters, stride,
+                 filter_size, padding, expansion_factor, name):
+        super(InvertedResidualUnit, self).__init__()
+        num_expfilter = int(round(num_in_filter * expansion_factor))
+        self._expand_conv = ConvBNLayer(
+            num_channels=num_channels,
+            num_filters=num_expfilter,
+            filter_size=1,
+            stride=1,
+            padding=0,
+            num_groups=1,
+            name=name + "_expand")
+
+        self._bottleneck_conv = ConvBNLayer(
+            num_channels=num_expfilter,
+            num_filters=num_expfilter,
+            filter_size=filter_size,
+            stride=stride,
+            padding=padding,
+            num_groups=num_expfilter,
+            use_cudnn=False,
+            name=name + "_dwise")
+
+        self._linear_conv = ConvBNLayer(
+            num_channels=num_expfilter,
+            num_filters=num_filters,
+            filter_size=1,
+            stride=1,
+            padding=0,
+            num_groups=1,
+            name=name + "_linear")
+
+    def forward(self, inputs, ifshortcut):
+        y = self._expand_conv(inputs, if_act=True)
+        y = self._bottleneck_conv(y, if_act=True)
+        y = self._linear_conv(y, if_act=False)
+        if ifshortcut:
+            y = paddle.add(inputs, y)
+        return y
+
+
+class InvresiBlocks(nn.Layer):
+    def __init__(self, in_c, t, c, n, s, name):
+        super(InvresiBlocks, self).__init__()
+
+        self._first_block = InvertedResidualUnit(
+            num_channels=in_c,
+            num_in_filter=in_c,
+            num_filters=c,
+            stride=s,
+            filter_size=3,
+            padding=1,
+            expansion_factor=t,
+            name=name + "_1")
+
+        self._block_list = []
+        for i in range(1, n):
+            block = self.add_sublayer(
+                name + "_" + str(i + 1),
+                sublayer=InvertedResidualUnit(
+                    num_channels=c,
+                    num_in_filter=c,
+                    num_filters=c,
+                    stride=1,
+                    filter_size=3,
+                    padding=1,
+                    expansion_factor=t,
+                    name=name + "_" + str(i + 1)))
+            self._block_list.append(block)
+
+    def forward(self, inputs):
+        y = self._first_block(inputs, ifshortcut=False)
+        for block in self._block_list:
+            y = block(y, ifshortcut=True)
+        return y
+
+
+class MobileNet(nn.Layer):
+    def __init__(self,
+                 input_channels=3,
+                 scale=1.0,
+                 pretrained=None,
+                 prefix_name=""):
+        super(MobileNet, self).__init__()
+        self.scale = scale
+
+        bottleneck_params_list = [
+            (1, 16, 1, 1),
+            (6, 24, 2, 2),
+            (6, 32, 3, 2),
+            (6, 64, 4, 2),
+            (6, 96, 3, 1),
+            (6, 160, 3, 2),
+            (6, 320, 1, 1),
+        ]
+
+        self.conv1 = ConvBNLayer(
+            num_channels=input_channels,
+            num_filters=int(32 * scale),
+            filter_size=3,
+            stride=2,
+            padding=1,
+            name=prefix_name + "conv1_1")
+
+        self.block_list = []
+        i = 1
+        in_c = int(32 * scale)
+        for layer_setting in bottleneck_params_list:
+            t, c, n, s = layer_setting
+            i += 1
+            block = self.add_sublayer(
+                prefix_name + "conv" + str(i),
+                sublayer=InvresiBlocks(
+                    in_c=in_c,
+                    t=t,
+                    c=int(c * scale),
+                    n=n,
+                    s=s,
+                    name=prefix_name + "conv" + str(i)))
+            self.block_list.append(block)
+            in_c = int(c * scale)
+
+        self.out_c = int(1280 * scale) if scale > 1.0 else 1280
+        self.conv9 = ConvBNLayer(
+            num_channels=in_c,
+            num_filters=self.out_c,
+            filter_size=1,
+            stride=1,
+            padding=0,
+            name=prefix_name + "conv9")
+
+        self.feat_channels = [int(i * scale) for i in [16, 24, 32, 96, 1280]]
+        self.pretrained = pretrained
+        self.init_weight()
+
+    def forward(self, inputs):
+        feat_list = []
+        y = self.conv1(inputs, if_act=True)
+
+        block_index = 0
+        for block in self.block_list:
+            y = block(y)
+            if block_index in [0, 1, 2, 4]:
+                feat_list.append(y)
+            block_index += 1
+        y = self.conv9(y, if_act=True)
+        feat_list.append(y)
+        return feat_list
+
+    def init_weight(self):
+        utils.load_pretrained_model(self, self.pretrained)
+
+
+@manager.BACKBONES.add_component
+def MobileNetV2(**kwargs):
+    model = MobileNet(scale=1.0, **kwargs)
+    return model

matting/model/modnet.py

@@ -0,0 +1,481 @@
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import defaultdict
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+import paddleseg
+from paddleseg.models import layers, losses
+from paddleseg import utils
+from paddleseg.cvlibs import manager, param_init
+import numpy as np
+import scipy
+
+
+@manager.MODELS.add_component
+class MODNet(nn.Layer):
+    """
+    The MODNet implementation based on PaddlePaddle.
+
+    The original article refers to
+    Zhanghan Ke, et, al. "Is a Green Screen Really Necessary for Real-Time Portrait Matting?"
+    (https://arxiv.org/pdf/2011.11961.pdf).
+
+    Args:
+        backbone: backbone model.
+        hr(int, optional): The channels of high resolutions branch. Defautl: None.
+        pretrained(str, optional): The path of pretrianed model. Defautl: None.
+
+    """
+
+    def __init__(self, backbone, hr_channels=32, pretrained=None):
+        super().__init__()
+        self.backbone = backbone
+        self.pretrained = pretrained
+
+        self.head = MODNetHead(
+            hr_channels=hr_channels, backbone_channels=backbone.feat_channels)
+        self.init_weight()
+        self.blurer = GaussianBlurLayer(1, 3)
+
+    def forward(self, inputs):
+        """
+        If training, return a dict.
+        If evaluation, return the final alpha prediction.
+        """
+        x = inputs['img']
+        feat_list = self.backbone(x)
+        y = self.head(inputs=inputs, feat_list=feat_list)
+
+        return y
+
+    def loss(self, logit_dict, label_dict, loss_func_dict=None):
+        if loss_func_dict is None:
+            loss_func_dict = defaultdict(list)
+            loss_func_dict['semantic'].append(paddleseg.models.MSELoss())
+            loss_func_dict['detail'].append(paddleseg.models.L1Loss())
+            loss_func_dict['fusion'].append(paddleseg.models.L1Loss())
+            loss_func_dict['fusion'].append(paddleseg.models.L1Loss())
+
+        loss = {}
+        # semantic loss
+        semantic_gt = F.interpolate(
+            label_dict['alpha'],
+            scale_factor=1 / 16,
+            mode='bilinear',
+            align_corners=False)
+        semantic_gt = self.blurer(semantic_gt)
+        #         semantic_gt.stop_gradient=True
+        loss['semantic'] = loss_func_dict['semantic'][0](logit_dict['semantic'],
+                                                         semantic_gt)
+
+        # detail loss
+        trimap = label_dict['trimap']
+        mask = (trimap == 128).astype('float32')
+        logit_detail = logit_dict['detail'] * mask
+        label_detail = label_dict['alpha'] * mask
+        loss_detail = loss_func_dict['detail'][0](logit_detail, label_detail)
+        loss_detail = loss_detail / (mask.mean() + 1e-6)
+        loss['detail'] = 10 * loss_detail
+
+        # fusion loss
+        matte = logit_dict['matte']
+        alpha = label_dict['alpha']
+        transition_mask = label_dict['trimap'] == 128
+        matte_boundary = paddle.where(transition_mask, matte, alpha)
+        # l1 loss
+        loss_fusion_l1 = loss_func_dict['fusion'][0](
+            matte,
+            alpha) + 4 * loss_func_dict['fusion'][0](matte_boundary, alpha)
+        # composition loss
+        loss_fusion_comp = loss_func_dict['fusion'][1](
+            matte * label_dict['img'],
+            alpha * label_dict['img']) + 4 * loss_func_dict['fusion'][1](
+                matte_boundary * label_dict['img'], alpha * label_dict['img'])
+        # consisten loss with semantic
+        transition_mask = F.interpolate(
+            label_dict['trimap'],
+            scale_factor=1 / 16,
+            mode='nearest',
+            align_corners=False)
+        transition_mask = transition_mask == 128
+        matte_con_sem = F.interpolate(
+            matte, scale_factor=1 / 16, mode='bilinear', align_corners=False)
+        matte_con_sem = self.blurer(matte_con_sem)
+        logit_semantic = logit_dict['semantic'].clone()
+        logit_semantic.stop_gradient = True
+        matte_con_sem = paddle.where(transition_mask, logit_semantic,
+                                     matte_con_sem)
+        if False:
+            import cv2
+            matte_con_sem_num = matte_con_sem.numpy()
+            matte_con_sem_num = matte_con_sem_num[0].squeeze()
+            matte_con_sem_num = (matte_con_sem_num * 255).astype('uint8')
+            semantic = logit_dict['semantic'].numpy()
+            semantic = semantic[0].squeeze()
+            semantic = (semantic * 255).astype('uint8')
+            transition_mask = transition_mask.astype('uint8')
+            transition_mask = transition_mask.numpy()
+            transition_mask = (transition_mask[0].squeeze()) * 255
+            cv2.imwrite('matte_con.png', matte_con_sem_num)
+            cv2.imwrite('semantic.png', semantic)
+            cv2.imwrite('transition.png', transition_mask)
+        mse_loss = paddleseg.models.MSELoss()
+        loss_fusion_con_sem = mse_loss(matte_con_sem, logit_dict['semantic'])
+        loss_fusion = loss_fusion_l1 + loss_fusion_comp + loss_fusion_con_sem
+        loss['fusion'] = loss_fusion
+        loss['fusion_l1'] = loss_fusion_l1
+        loss['fusion_comp'] = loss_fusion_comp
+        loss['fusion_con_sem'] = loss_fusion_con_sem
+
+        loss['all'] = loss['semantic'] + loss['detail'] + loss['fusion']
+
+        return loss
+
+    def init_weight(self):
+        if self.pretrained is not None:
+            utils.load_entire_model(self, self.pretrained)
+
+
+class MODNetHead(nn.Layer):
+    def __init__(self, hr_channels, backbone_channels):
+        super().__init__()
+
+        self.lr_branch = LRBranch(backbone_channels)
+        self.hr_branch = HRBranch(hr_channels, backbone_channels)
+        self.f_branch = FusionBranch(hr_channels, backbone_channels)
+        self.init_weight()
+
+    def forward(self, inputs, feat_list):
+        pred_semantic, lr8x, [enc2x, enc4x] = self.lr_branch(feat_list)
+        pred_detail, hr2x = self.hr_branch(inputs['img'], enc2x, enc4x, lr8x)
+        pred_matte = self.f_branch(inputs['img'], lr8x, hr2x)
+
+        if self.training:
+            logit_dict = {
+                'semantic': pred_semantic,
+                'detail': pred_detail,
+                'matte': pred_matte
+            }
+            return logit_dict
+        else:
+            return pred_matte
+
+    def init_weight(self):
+        for layer in self.sublayers():
+            if isinstance(layer, nn.Conv2D):
+                param_init.kaiming_uniform(layer.weight)
+
+
+class FusionBranch(nn.Layer):
+    def __init__(self, hr_channels, enc_channels):
+        super().__init__()
+        self.conv_lr4x = Conv2dIBNormRelu(
+            enc_channels[2], hr_channels, 5, stride=1, padding=2)
+
+        self.conv_f2x = Conv2dIBNormRelu(
+            2 * hr_channels, hr_channels, 3, stride=1, padding=1)
+        self.conv_f = nn.Sequential(
+            Conv2dIBNormRelu(
+                hr_channels + 3, int(hr_channels / 2), 3, stride=1, padding=1),
+            Conv2dIBNormRelu(
+                int(hr_channels / 2),
+                1,
+                1,
+                stride=1,
+                padding=0,
+                with_ibn=False,
+                with_relu=False))
+
+    def forward(self, img, lr8x, hr2x):
+        lr4x = F.interpolate(
+            lr8x, scale_factor=2, mode='bilinear', align_corners=False)
+        lr4x = self.conv_lr4x(lr4x)
+        lr2x = F.interpolate(
+            lr4x, scale_factor=2, mode='bilinear', align_corners=False)
+
+        f2x = self.conv_f2x(paddle.concat((lr2x, hr2x), axis=1))
+        f = F.interpolate(
+            f2x, scale_factor=2, mode='bilinear', align_corners=False)
+        f = self.conv_f(paddle.concat((f, img), axis=1))
+        pred_matte = F.sigmoid(f)
+
+        return pred_matte
+
+
+class HRBranch(nn.Layer):
+    """
+    High Resolution Branch of MODNet
+    """
+
+    def __init__(self, hr_channels, enc_channels):
+        super().__init__()
+
+        self.tohr_enc2x = Conv2dIBNormRelu(
+            enc_channels[0], hr_channels, 1, stride=1, padding=0)
+        self.conv_enc2x = Conv2dIBNormRelu(
+            hr_channels + 3, hr_channels, 3, stride=2, padding=1)
+
+        self.tohr_enc4x = Conv2dIBNormRelu(
+            enc_channels[1], hr_channels, 1, stride=1, padding=0)
+        self.conv_enc4x = Conv2dIBNormRelu(
+            2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1)
+
+        self.conv_hr4x = nn.Sequential(
+            Conv2dIBNormRelu(
+                2 * hr_channels + enc_channels[2] + 3,
+                2 * hr_channels,
+                3,
+                stride=1,
+                padding=1),
+            Conv2dIBNormRelu(
+                2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1),
+            Conv2dIBNormRelu(
+                2 * hr_channels, hr_channels, 3, stride=1, padding=1))
+
+        self.conv_hr2x = nn.Sequential(
+            Conv2dIBNormRelu(
+                2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1),
+            Conv2dIBNormRelu(
+                2 * hr_channels, hr_channels, 3, stride=1, padding=1),
+            Conv2dIBNormRelu(hr_channels, hr_channels, 3, stride=1, padding=1),
+            Conv2dIBNormRelu(hr_channels, hr_channels, 3, stride=1, padding=1))
+
+        self.conv_hr = nn.Sequential(
+            Conv2dIBNormRelu(
+                hr_channels + 3, hr_channels, 3, stride=1, padding=1),
+            Conv2dIBNormRelu(
+                hr_channels,
+                1,
+                1,
+                stride=1,
+                padding=0,
+                with_ibn=False,
+                with_relu=False))
+
+    def forward(self, img, enc2x, enc4x, lr8x):
+        img2x = F.interpolate(
+            img, scale_factor=1 / 2, mode='bilinear', align_corners=False)
+        img4x = F.interpolate(
+            img, scale_factor=1 / 4, mode='bilinear', align_corners=False)
+
+        enc2x = self.tohr_enc2x(enc2x)
+        hr4x = self.conv_enc2x(paddle.concat((img2x, enc2x), axis=1))
+
+        enc4x = self.tohr_enc4x(enc4x)
+        hr4x = self.conv_enc4x(paddle.concat((hr4x, enc4x), axis=1))
+
+        lr4x = F.interpolate(
+            lr8x, scale_factor=2, mode='bilinear', align_corners=False)
+        hr4x = self.conv_hr4x(paddle.concat((hr4x, lr4x, img4x), axis=1))
+
+        hr2x = F.interpolate(
+            hr4x, scale_factor=2, mode='bilinear', align_corners=False)
+        hr2x = self.conv_hr2x(paddle.concat((hr2x, enc2x), axis=1))
+
+        pred_detail = None
+        if self.training:
+            hr = F.interpolate(
+                hr2x, scale_factor=2, mode='bilinear', align_corners=False)
+            hr = self.conv_hr(paddle.concat((hr, img), axis=1))
+            pred_detail = F.sigmoid(hr)
+
+        return pred_detail, hr2x
+
+
+class LRBranch(nn.Layer):
+    def __init__(self, backbone_channels):
+        super().__init__()
+        self.se_block = SEBlock(backbone_channels[4], reduction=4)
+        self.conv_lr16x = Conv2dIBNormRelu(
+            backbone_channels[4], backbone_channels[3], 5, stride=1, padding=2)
+        self.conv_lr8x = Conv2dIBNormRelu(
+            backbone_channels[3], backbone_channels[2], 5, stride=1, padding=2)
+        self.conv_lr = Conv2dIBNormRelu(
+            backbone_channels[2],
+            1,
+            3,
+            stride=2,
+            padding=1,
+            with_ibn=False,
+            with_relu=False)
+
+    def forward(self, feat_list):
+        enc2x, enc4x, enc32x = feat_list[0], feat_list[1], feat_list[4]
+
+        enc32x = self.se_block(enc32x)
+        lr16x = F.interpolate(
+            enc32x, scale_factor=2, mode='bilinear', align_corners=False)
+        lr16x = self.conv_lr16x(lr16x)
+        lr8x = F.interpolate(
+            lr16x, scale_factor=2, mode='bilinear', align_corners=False)
+        lr8x = self.conv_lr8x(lr8x)
+
+        pred_semantic = None
+        if self.training:
+            lr = self.conv_lr(lr8x)
+            pred_semantic = F.sigmoid(lr)
+
+        return pred_semantic, lr8x, [enc2x, enc4x]
+
+
+class IBNorm(nn.Layer):
+    """
+    Combine Instance Norm and Batch Norm into One Layer
+    """
+
+    def __init__(self, in_channels):
+        super().__init__()
+        self.bnorm_channels = in_channels // 2
+        self.inorm_channels = in_channels - self.bnorm_channels
+
+        self.bnorm = nn.BatchNorm2D(self.bnorm_channels)
+        self.inorm = nn.InstanceNorm2D(self.inorm_channels)
+
+    def forward(self, x):
+        bn_x = self.bnorm(x[:, :self.bnorm_channels, :, :])
+        in_x = self.inorm(x[:, self.bnorm_channels:, :, :])
+
+        return paddle.concat((bn_x, in_x), 1)
+
+
+class Conv2dIBNormRelu(nn.Layer):
+    """
+    Convolution + IBNorm + Relu
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias_attr=None,
+                 with_ibn=True,
+                 with_relu=True):
+
+        super().__init__()
+
+        layers = [
+            nn.Conv2D(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+                groups=groups,
+                bias_attr=bias_attr)
+        ]
+
+        if with_ibn:
+            layers.append(IBNorm(out_channels))
+
+        if with_relu:
+            layers.append(nn.ReLU())
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class SEBlock(nn.Layer):
+    """
+    SE Block Proposed in https://arxiv.org/pdf/1709.01507.pdf
+    """
+
+    def __init__(self, num_channels, reduction=1):
+        super().__init__()
+        self.pool = nn.AdaptiveAvgPool2D(1)
+        self.conv = nn.Sequential(
+            nn.Conv2D(
+                num_channels,
+                int(num_channels // reduction),
+                1,
+                bias_attr=False), nn.ReLU(),
+            nn.Conv2D(
+                int(num_channels // reduction),
+                num_channels,
+                1,
+                bias_attr=False), nn.Sigmoid())
+
+    def forward(self, x):
+        w = self.pool(x)
+        w = self.conv(w)
+        return w * x
+
+
+class GaussianBlurLayer(nn.Layer):
+    """ Add Gaussian Blur to a 4D tensors
+    This layer takes a 4D tensor of {N, C, H, W} as input.
+    The Gaussian blur will be performed in given channel number (C) splitly.
+    """
+
+    def __init__(self, channels, kernel_size):
+        """
+        Args:
+            channels (int): Channel for input tensor
+            kernel_size (int): Size of the kernel used in blurring
+        """
+
+        super(GaussianBlurLayer, self).__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        assert self.kernel_size % 2 != 0
+
+        self.op = nn.Sequential(
+            nn.Pad2D(int(self.kernel_size / 2), mode='reflect'),
+            nn.Conv2D(
+                channels,
+                channels,
+                self.kernel_size,
+                stride=1,
+                padding=0,
+                bias_attr=False,
+                groups=channels))
+
+        self._init_kernel()
+        self.op[1].weight.stop_gradient = True
+
+    def forward(self, x):
+        """
+        Args:
+            x (paddle.Tensor): input 4D tensor
+        Returns:
+            paddle.Tensor: Blurred version of the input
+        """
+
+        if not len(list(x.shape)) == 4:
+            print('\'GaussianBlurLayer\' requires a 4D tensor as input\n')
+            exit()
+        elif not x.shape[1] == self.channels:
+            print('In \'GaussianBlurLayer\', the required channel ({0}) is'
+                  'not the same as input ({1})\n'.format(
+                      self.channels, x.shape[1]))
+            exit()
+
+        return self.op(x)
+
+    def _init_kernel(self):
+        sigma = 0.3 * ((self.kernel_size - 1) * 0.5 - 1) + 0.8
+
+        n = np.zeros((self.kernel_size, self.kernel_size))
+        i = int(self.kernel_size / 2)
+        n[i, i] = 1
+        kernel = scipy.ndimage.gaussian_filter(n, sigma)
+        kernel = kernel.astype('float32')
+        kernel = kernel[np.newaxis, np.newaxis, :, :]
+        paddle.assign(kernel, self.op[1].weight)

matting/model/resnet_vd.py

@@ -0,0 +1,368 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from paddleseg.cvlibs import manager
+from paddleseg.models import layers
+from paddleseg.utils import utils
+
+__all__ = [
+    "ResNet18_vd", "ResNet34_vd", "ResNet50_vd", "ResNet101_vd", "ResNet152_vd"
+]
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=1,
+            dilation=1,
+            groups=1,
+            is_vd_mode=False,
+            act=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self.is_vd_mode = is_vd_mode
+        self._pool2d_avg = nn.AvgPool2D(
+            kernel_size=2, stride=2, padding=0, ceil_mode=True)
+        self._conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2 if dilation == 1 else 0,
+            dilation=dilation,
+            groups=groups,
+            bias_attr=False)
+
+        self._batch_norm = layers.SyncBatchNorm(out_channels)
+        self._act_op = layers.Activation(act=act)
+
+    def forward(self, inputs):
+        if self.is_vd_mode:
+            inputs = self._pool2d_avg(inputs)
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        y = self._act_op(y)
+
+        return y
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 stride,
+                 shortcut=True,
+                 if_first=False,
+                 dilation=1):
+        super(BottleneckBlock, self).__init__()
+
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            act='relu')
+
+        self.dilation = dilation
+
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act='relu',
+            dilation=dilation)
+        self.conv2 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels * 4,
+            kernel_size=1,
+            act=None)
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels * 4,
+                kernel_size=1,
+                stride=1,
+                is_vd_mode=False if if_first or stride == 1 else True)
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+
+        ####################################################################
+        # If given dilation rate > 1, using corresponding padding.
+        # The performance drops down without the follow padding.
+        if self.dilation > 1:
+            padding = self.dilation
+            y = F.pad(y, [padding, padding, padding, padding])
+        #####################################################################
+
+        conv1 = self.conv1(y)
+        conv2 = self.conv2(conv1)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+
+        y = paddle.add(x=short, y=conv2)
+        y = F.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 stride,
+                 shortcut=True,
+                 if_first=False):
+        super(BasicBlock, self).__init__()
+        self.stride = stride
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act='relu')
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            act=None)
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                is_vd_mode=False if if_first else True)
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv1)
+        y = F.relu(y)
+
+        return y
+
+
+class ResNet_vd(nn.Layer):
+    """
+    The ResNet_vd implementation based on PaddlePaddle.
+
+    The original article refers to Jingdong
+    Tong He, et, al. "Bag of Tricks for Image Classification with Convolutional Neural Networks"
+    (https://arxiv.org/pdf/1812.01187.pdf).
+
+    Args:
+        layers (int, optional): The layers of ResNet_vd. The supported layers are (18, 34, 50, 101, 152, 200). Default: 50.
+        output_stride (int, optional): The stride of output features compared to input images. It is 8 or 16. Default: 8.
+        multi_grid (tuple|list, optional): The grid of stage4. Defult: (1, 1, 1).
+        pretrained (str, optional): The path of pretrained model.
+
+    """
+
+    def __init__(self,
+                 input_channels=3,
+                 layers=50,
+                 output_stride=32,
+                 multi_grid=(1, 1, 1),
+                 pretrained=None):
+        super(ResNet_vd, self).__init__()
+
+        self.conv1_logit = None  # for gscnn shape stream
+        self.layers = layers
+        supported_layers = [18, 34, 50, 101, 152, 200]
+        assert layers in supported_layers, \
+            "supported layers are {} but input layer is {}".format(
+                supported_layers, layers)
+
+        if layers == 18:
+            depth = [2, 2, 2, 2]
+        elif layers == 34 or layers == 50:
+            depth = [3, 4, 6, 3]
+        elif layers == 101:
+            depth = [3, 4, 23, 3]
+        elif layers == 152:
+            depth = [3, 8, 36, 3]
+        elif layers == 200:
+            depth = [3, 12, 48, 3]
+        num_channels = [64, 256, 512, 1024
+                        ] if layers >= 50 else [64, 64, 128, 256]
+        num_filters = [64, 128, 256, 512]
+
+        # for channels of four returned stages
+        self.feat_channels = [c * 4 for c in num_filters
+                              ] if layers >= 50 else num_filters
+        self.feat_channels = [64] + self.feat_channels
+
+        dilation_dict = None
+        if output_stride == 8:
+            dilation_dict = {2: 2, 3: 4}
+        elif output_stride == 16:
+            dilation_dict = {3: 2}
+
+        self.conv1_1 = ConvBNLayer(
+            in_channels=input_channels,
+            out_channels=32,
+            kernel_size=3,
+            stride=2,
+            act='relu')
+        self.conv1_2 = ConvBNLayer(
+            in_channels=32,
+            out_channels=32,
+            kernel_size=3,
+            stride=1,
+            act='relu')
+        self.conv1_3 = ConvBNLayer(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            act='relu')
+        self.pool2d_max = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
+
+        # self.block_list = []
+        self.stage_list = []
+        if layers >= 50:
+            for block in range(len(depth)):
+                shortcut = False
+                block_list = []
+                for i in range(depth[block]):
+                    if layers in [101, 152] and block == 2:
+                        if i == 0:
+                            conv_name = "res" + str(block + 2) + "a"
+                        else:
+                            conv_name = "res" + str(block + 2) + "b" + str(i)
+                    else:
+                        conv_name = "res" + str(block + 2) + chr(97 + i)
+
+                    ###############################################################################
+                    # Add dilation rate for some segmentation tasks, if dilation_dict is not None.
+                    dilation_rate = dilation_dict[
+                        block] if dilation_dict and block in dilation_dict else 1
+
+                    # Actually block here is 'stage', and i is 'block' in 'stage'
+                    # At the stage 4, expand the the dilation_rate if given multi_grid
+                    if block == 3:
+                        dilation_rate = dilation_rate * multi_grid[i]
+                    ###############################################################################
+
+                    bottleneck_block = self.add_sublayer(
+                        'bb_%d_%d' % (block, i),
+                        BottleneckBlock(
+                            in_channels=num_channels[block]
+                            if i == 0 else num_filters[block] * 4,
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0
+                            and dilation_rate == 1 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            dilation=dilation_rate))
+
+                    block_list.append(bottleneck_block)
+                    shortcut = True
+                self.stage_list.append(block_list)
+        else:
+            for block in range(len(depth)):
+                shortcut = False
+                block_list = []
+                for i in range(depth[block]):
+                    conv_name = "res" + str(block + 2) + chr(97 + i)
+                    basic_block = self.add_sublayer(
+                        'bb_%d_%d' % (block, i),
+                        BasicBlock(
+                            in_channels=num_channels[block]
+                            if i == 0 else num_filters[block],
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0))
+                    block_list.append(basic_block)
+                    shortcut = True
+                self.stage_list.append(block_list)
+
+        self.pretrained = pretrained
+        self.init_weight()
+
+    def forward(self, inputs):
+        feat_list = []
+        y = self.conv1_1(inputs)
+        y = self.conv1_2(y)
+        y = self.conv1_3(y)
+        feat_list.append(y)
+
+        y = self.pool2d_max(y)
+
+        # A feature list saves the output feature map of each stage.
+        for stage in self.stage_list:
+            for block in stage:
+                y = block(y)
+            feat_list.append(y)
+
+        return feat_list
+
+    def init_weight(self):
+        utils.load_pretrained_model(self, self.pretrained)
+
+
+@manager.BACKBONES.add_component
+def ResNet18_vd(**args):
+    model = ResNet_vd(layers=18, **args)
+    return model
+
+
+def ResNet34_vd(**args):
+    model = ResNet_vd(layers=34, **args)
+    return model
+
+
+@manager.BACKBONES.add_component
+def ResNet50_vd(**args):
+    model = ResNet_vd(layers=50, **args)
+    return model
+
+
+@manager.BACKBONES.add_component
+def ResNet101_vd(**args):
+    model = ResNet_vd(layers=101, **args)
+    return model
+
+
+def ResNet152_vd(**args):
+    model = ResNet_vd(layers=152, **args)
+    return model
+
+
+def ResNet200_vd(**args):
+    model = ResNet_vd(layers=200, **args)
+    return model

matting/model/vgg.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.nn import Conv2D, BatchNorm, Linear, Dropout
+from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D
+
+from paddleseg.cvlibs import manager
+from paddleseg.utils import utils
+
+
+class ConvBlock(nn.Layer):
+    def __init__(self, input_channels, output_channels, groups, name=None):
+        super(ConvBlock, self).__init__()
+
+        self.groups = groups
+        self._conv_1 = Conv2D(
+            in_channels=input_channels,
+            out_channels=output_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=ParamAttr(name=name + "1_weights"),
+            bias_attr=False)
+        if groups == 2 or groups == 3 or groups == 4:
+            self._conv_2 = Conv2D(
+                in_channels=output_channels,
+                out_channels=output_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                weight_attr=ParamAttr(name=name + "2_weights"),
+                bias_attr=False)
+        if groups == 3 or groups == 4:
+            self._conv_3 = Conv2D(
+                in_channels=output_channels,
+                out_channels=output_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                weight_attr=ParamAttr(name=name + "3_weights"),
+                bias_attr=False)
+        if groups == 4:
+            self._conv_4 = Conv2D(
+                in_channels=output_channels,
+                out_channels=output_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                weight_attr=ParamAttr(name=name + "4_weights"),
+                bias_attr=False)
+
+        self._pool = MaxPool2D(
+            kernel_size=2, stride=2, padding=0, return_mask=True)
+
+    def forward(self, inputs):
+        x = self._conv_1(inputs)
+        x = F.relu(x)
+        if self.groups == 2 or self.groups == 3 or self.groups == 4:
+            x = self._conv_2(x)
+            x = F.relu(x)
+        if self.groups == 3 or self.groups == 4:
+            x = self._conv_3(x)
+            x = F.relu(x)
+        if self.groups == 4:
+            x = self._conv_4(x)
+            x = F.relu(x)
+        skip = x
+        x, max_indices = self._pool(x)
+        return x, max_indices, skip
+
+
+class VGGNet(nn.Layer):
+    def __init__(self, input_channels=3, layers=11, pretrained=None):
+        super(VGGNet, self).__init__()
+        self.pretrained = pretrained
+
+        self.layers = layers
+        self.vgg_configure = {
+            11: [1, 1, 2, 2, 2],
+            13: [2, 2, 2, 2, 2],
+            16: [2, 2, 3, 3, 3],
+            19: [2, 2, 4, 4, 4]
+        }
+        assert self.layers in self.vgg_configure.keys(), \
+            "supported layers are {} but input layer is {}".format(
+                self.vgg_configure.keys(), layers)
+        self.groups = self.vgg_configure[self.layers]
+
+        # matting的第一层卷积输入为4通道，初始化是直接初始化为0
+        self._conv_block_1 = ConvBlock(
+            input_channels, 64, self.groups[0], name="conv1_")
+        self._conv_block_2 = ConvBlock(64, 128, self.groups[1], name="conv2_")
+        self._conv_block_3 = ConvBlock(128, 256, self.groups[2], name="conv3_")
+        self._conv_block_4 = ConvBlock(256, 512, self.groups[3], name="conv4_")
+        self._conv_block_5 = ConvBlock(512, 512, self.groups[4], name="conv5_")
+
+        # 这一层的初始化需要利用vgg fc6的参数转换后进行初始化，可以暂时不考虑初始化
+        self._conv_6 = Conv2D(
+            512, 512, kernel_size=3, padding=1, bias_attr=False)
+
+        self.init_weight()
+
+    def forward(self, inputs):
+        fea_list = []
+        ids_list = []
+        x, ids, skip = self._conv_block_1(inputs)
+        fea_list.append(skip)
+        ids_list.append(ids)
+        x, ids, skip = self._conv_block_2(x)
+        fea_list.append(skip)
+        ids_list.append(ids)
+        x, ids, skip = self._conv_block_3(x)
+        fea_list.append(skip)
+        ids_list.append(ids)
+        x, ids, skip = self._conv_block_4(x)
+        fea_list.append(skip)
+        ids_list.append(ids)
+        x, ids, skip = self._conv_block_5(x)
+        fea_list.append(skip)
+        ids_list.append(ids)
+        x = F.relu(self._conv_6(x))
+        fea_list.append(x)
+        return fea_list
+
+    def init_weight(self):
+        if self.pretrained is not None:
+            utils.load_pretrained_model(self, self.pretrained)
+
+
+@manager.BACKBONES.add_component
+def VGG11(**args):
+    model = VGGNet(layers=11, **args)
+    return model
+
+
+@manager.BACKBONES.add_component
+def VGG13(**args):
+    model = VGGNet(layers=13, **args)
+    return model
+
+
+@manager.BACKBONES.add_component
+def VGG16(**args):
+    model = VGGNet(layers=16, **args)
+    return model
+
+
+@manager.BACKBONES.add_component
+def VGG19(**args):
+    model = VGGNet(layers=19, **args)
+    return model

matting/transforms.py

@@ -0,0 +1,530 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+
+import cv2
+import numpy as np
+from paddleseg.transforms import functional
+from paddleseg.cvlibs import manager
+from PIL import Image
+
+
+@manager.TRANSFORMS.add_component
+class Compose:
+    """
+    Do transformation on input data with corresponding pre-processing and augmentation operations.
+    The shape of input data to all operations is [height, width, channels].
+    """
+
+    def __init__(self, transforms, to_rgb=True):
+        if not isinstance(transforms, list):
+            raise TypeError('The transforms must be a list!')
+        self.transforms = transforms
+        self.to_rgb = to_rgb
+
+    def __call__(self, data):
+        """
+        Args:
+            data (dict): The data to transform.
+
+        Returns:
+            dict: Data after transformation
+        """
+        if 'trans_info' not in data:
+            data['trans_info'] = []
+        for op in self.transforms:
+            data = op(data)
+            if data is None:
+                return None
+
+        data['img'] = np.transpose(data['img'], (2, 0, 1))
+        for key in data.get('gt_fields', []):
+            if len(data[key].shape) == 2:
+                continue
+            data[key] = np.transpose(data[key], (2, 0, 1))
+
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class LoadImages:
+    def __init__(self, to_rgb=True):
+        self.to_rgb = to_rgb
+
+    def __call__(self, data):
+        if isinstance(data['img'], str):
+            data['img'] = cv2.imread(data['img'])
+        for key in data.get('gt_fields', []):
+            if isinstance(data[key], str):
+                data[key] = cv2.imread(data[key], cv2.IMREAD_UNCHANGED)
+            # if alpha and trimap has 3 channels, extract one.
+            if key in ['alpha', 'trimap']:
+                if len(data[key].shape) > 2:
+                    data[key] = data[key][:, :, 0]
+
+        if self.to_rgb:
+            data['img'] = cv2.cvtColor(data['img'], cv2.COLOR_BGR2RGB)
+            for key in data.get('gt_fields', []):
+                if len(data[key].shape) == 2:
+                    continue
+                data[key] = cv2.cvtColor(data[key], cv2.COLOR_BGR2RGB)
+
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class Resize:
+    def __init__(self, target_size=(512, 512)):
+        if isinstance(target_size, list) or isinstance(target_size, tuple):
+            if len(target_size) != 2:
+                raise ValueError(
+                    '`target_size` should include 2 elements, but it is {}'.
+                    format(target_size))
+        else:
+            raise TypeError(
+                "Type of `target_size` is invalid. It should be list or tuple, but it is {}"
+                .format(type(target_size)))
+
+        self.target_size = target_size
+
+    def __call__(self, data):
+        data['trans_info'].append(('resize', data['img'].shape[0:2]))
+        data['img'] = functional.resize(data['img'], self.target_size)
+        for key in data.get('gt_fields', []):
+            data[key] = functional.resize(data[key], self.target_size)
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class ResizeByLong:
+    """
+    Resize the long side of an image to given size, and then scale the other side proportionally.
+
+    Args:
+        long_size (int): The target size of long side.
+    """
+
+    def __init__(self, long_size):
+        self.long_size = long_size
+
+    def __call__(self, data):
+        data['trans_info'].append(('resize', data['img'].shape[0:2]))
+        data['img'] = functional.resize_long(data['img'], self.long_size)
+        for key in data.get('gt_fields', []):
+            data[key] = functional.resize_long(data[key], self.long_size)
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class ResizeByShort:
+    """
+    Resize the short side of an image to given size, and then scale the other side proportionally.
+
+    Args:
+        short_size (int): The target size of short side.
+    """
+
+    def __init__(self, short_size):
+        self.short_size = short_size
+
+    def __call__(self, data):
+        data['trans_info'].append(('resize', data['img'].shape[0:2]))
+        data['img'] = functional.resize_short(data['img'], self.short_size)
+        for key in data.get('gt_fields', []):
+            data[key] = functional.resize_short(data[key], self.short_size)
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class ResizeToIntMult:
+    """
+    Resize to some int muitple, d.g. 32.
+    """
+
+    def __init__(self, mult_int=32):
+        self.mult_int = mult_int
+
+    def __call__(self, data):
+        data['trans_info'].append(('resize', data['img'].shape[0:2]))
+
+        h, w = data['img'].shape[0:2]
+        rw = w - w % 32
+        rh = h - h % 32
+        data['img'] = functional.resize(data['img'], (rw, rh))
+        for key in data.get('gt_fields', []):
+            data[key] = functional.resize(data[key], (rw, rh))
+
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class Normalize:
+    """
+    Normalize an image.
+
+    Args:
+        mean (list, optional): The mean value of a data set. Default: [0.5, 0.5, 0.5].
+        std (list, optional): The standard deviation of a data set. Default: [0.5, 0.5, 0.5].
+
+    Raises:
+        ValueError: When mean/std is not list or any value in std is 0.
+    """
+
+    def __init__(self, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)):
+        self.mean = mean
+        self.std = std
+        if not (isinstance(self.mean, (list, tuple))
+                and isinstance(self.std, (list, tuple))):
+            raise ValueError(
+                "{}: input type is invalid. It should be list or tuple".format(
+                    self))
+        from functools import reduce
+        if reduce(lambda x, y: x * y, self.std) == 0:
+            raise ValueError('{}: std is invalid!'.format(self))
+
+    def __call__(self, data):
+        mean = np.array(self.mean)[np.newaxis, np.newaxis, :]
+        std = np.array(self.std)[np.newaxis, np.newaxis, :]
+        data['img'] = functional.normalize(data['img'], mean, std)
+        if 'fg' in data.get('gt_fields', []):
+            data['fg'] = functional.normalize(data['fg'], mean, std)
+        if 'bg' in data.get('gt_fields', []):
+            data['bg'] = functional.normalize(data['bg'], mean, std)
+
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class RandomCropByAlpha:
+    """
+    Randomly crop while centered on uncertain area by a certain probability.
+
+    Args:
+        crop_size (tuple|list): The size you want to crop from image.
+        p (float): The probability centered on uncertain area.
+
+    """
+
+    def __init__(self, crop_size=((320, 320), (480, 480), (640, 640)),
+                 prob=0.5):
+        self.crop_size = crop_size
+        self.prob = prob
+
+    def __call__(self, data):
+        idex = np.random.randint(low=0, high=len(self.crop_size))
+        crop_w, crop_h = self.crop_size[idex]
+
+        img_h = data['img'].shape[0]
+        img_w = data['img'].shape[1]
+        if np.random.rand() < self.prob:
+            crop_center = np.where((data['alpha'] > 0) & (data['alpha'] < 255))
+            center_h_array, center_w_array = crop_center
+            if len(center_h_array) == 0:
+                return data
+            rand_ind = np.random.randint(len(center_h_array))
+            center_h = center_h_array[rand_ind]
+            center_w = center_w_array[rand_ind]
+            delta_h = crop_h // 2
+            delta_w = crop_w // 2
+            start_h = max(0, center_h - delta_h)
+            start_w = max(0, center_w - delta_w)
+        else:
+            start_h = 0
+            start_w = 0
+            if img_h > crop_h:
+                start_h = np.random.randint(img_h - crop_h + 1)
+            if img_w > crop_w:
+                start_w = np.random.randint(img_w - crop_w + 1)
+
+        end_h = min(img_h, start_h + crop_h)
+        end_w = min(img_w, start_w + crop_w)
+
+        data['img'] = data['img'][start_h:end_h, start_w:end_w]
+        for key in data.get('gt_fields', []):
+            data[key] = data[key][start_h:end_h, start_w:end_w]
+
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class RandomCrop:
+    """
+    Randomly crop
+
+    Args:
+    crop_size (tuple|list): The size you want to crop from image.
+    """
+
+    def __init__(self, crop_size=((320, 320), (480, 480), (640, 640))):
+        if not isinstance(crop_size[0], (list, tuple)):
+            crop_size = [crop_size]
+        self.crop_size = crop_size
+
+    def __call__(self, data):
+        idex = np.random.randint(low=0, high=len(self.crop_size))
+        crop_w, crop_h = self.crop_size[idex]
+        img_h, img_w = data['img'].shape[0:2]
+
+        start_h = 0
+        start_w = 0
+        if img_h > crop_h:
+            start_h = np.random.randint(img_h - crop_h + 1)
+        if img_w > crop_w:
+            start_w = np.random.randint(img_w - crop_w + 1)
+
+        end_h = min(img_h, start_h + crop_h)
+        end_w = min(img_w, start_w + crop_w)
+
+        data['img'] = data['img'][start_h:end_h, start_w:end_w]
+        for key in data.get('gt_fields', []):
+            data[key] = data[key][start_h:end_h, start_w:end_w]
+
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class LimitLong:
+    """
+    Limit the long edge of image.
+
+    If the long edge is larger than max_long, resize the long edge
+    to max_long, while scale the short edge proportionally.
+
+    If the long edge is smaller than min_long, resize the long edge
+    to min_long, while scale the short edge proportionally.
+
+    Args:
+        max_long (int, optional): If the long edge of image is larger than max_long,
+            it will be resize to max_long. Default: None.
+        min_long (int, optional): If the long edge of image is smaller than min_long,
+            it will be resize to min_long. Default: None.
+    """
+
+    def __init__(self, max_long=None, min_long=None):
+        if max_long is not None:
+            if not isinstance(max_long, int):
+                raise TypeError(
+                    "Type of `max_long` is invalid. It should be int, but it is {}"
+                    .format(type(max_long)))
+        if min_long is not None:
+            if not isinstance(min_long, int):
+                raise TypeError(
+                    "Type of `min_long` is invalid. It should be int, but it is {}"
+                    .format(type(min_long)))
+        if (max_long is not None) and (min_long is not None):
+            if min_long > max_long:
+                raise ValueError(
+                    '`max_long should not smaller than min_long, but they are {} and {}'
+                    .format(max_long, min_long))
+        self.max_long = max_long
+        self.min_long = min_long
+
+    def __call__(self, data):
+        h, w = data['img'].shape[:2]
+        long_edge = max(h, w)
+        target = long_edge
+        if (self.max_long is not None) and (long_edge > self.max_long):
+            target = self.max_long
+        elif (self.min_long is not None) and (long_edge < self.min_long):
+            target = self.min_long
+
+        if target != long_edge:
+            data['trans_info'].append(('resize', data['img'].shape[0:2]))
+            data['img'] = functional.resize_long(data['img'], target)
+            for key in data.get('gt_fields', []):
+                data[key] = functional.resize_long(data[key], target)
+
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class RandomHorizontalFlip:
+    """
+    Flip an image horizontally with a certain probability.
+
+    Args:
+        prob (float, optional): A probability of horizontally flipping. Default: 0.5.
+    """
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, data):
+        if random.random() < self.prob:
+            data['img'] = functional.horizontal_flip(data['img'])
+            for key in data.get('gt_fields', []):
+                data[key] = functional.horizontal_flip(data[key])
+
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class RandomBlur:
+    """
+    Blurring an image by a Gaussian function with a certain probability.
+
+    Args:
+        prob (float, optional): A probability of blurring an image. Default: 0.1.
+    """
+
+    def __init__(self, prob=0.1):
+        self.prob = prob
+
+    def __call__(self, data):
+        if self.prob <= 0:
+            n = 0
+        elif self.prob >= 1:
+            n = 1
+        else:
+            n = int(1.0 / self.prob)
+        if n > 0:
+            if np.random.randint(0, n) == 0:
+                radius = np.random.randint(3, 10)
+                if radius % 2 != 1:
+                    radius = radius + 1
+                if radius > 9:
+                    radius = 9
+                data['img'] = cv2.GaussianBlur(data['img'], (radius, radius), 0,
+                                               0)
+                for key in data.get('gt_fields', []):
+                    data[key] = cv2.GaussianBlur(data[key], (radius, radius), 0,
+                                                 0)
+        return data
+
+
+@manager.TRANSFORMS.add_component
+class RandomDistort:
+    """
+    Distort an image with random configurations.
+
+    Args:
+        brightness_range (float, optional): A range of brightness. Default: 0.5.
+        brightness_prob (float, optional): A probability of adjusting brightness. Default: 0.5.
+        contrast_range (float, optional): A range of contrast. Default: 0.5.
+        contrast_prob (float, optional): A probability of adjusting contrast. Default: 0.5.
+        saturation_range (float, optional): A range of saturation. Default: 0.5.
+        saturation_prob (float, optional): A probability of adjusting saturation. Default: 0.5.
+        hue_range (int, optional): A range of hue. Default: 18.
+        hue_prob (float, optional): A probability of adjusting hue. Default: 0.5.
+    """
+
+    def __init__(self,
+                 brightness_range=0.5,
+                 brightness_prob=0.5,
+                 contrast_range=0.5,
+                 contrast_prob=0.5,
+                 saturation_range=0.5,
+                 saturation_prob=0.5,
+                 hue_range=18,
+                 hue_prob=0.5):
+        self.brightness_range = brightness_range
+        self.brightness_prob = brightness_prob
+        self.contrast_range = contrast_range
+        self.contrast_prob = contrast_prob
+        self.saturation_range = saturation_range
+        self.saturation_prob = saturation_prob
+        self.hue_range = hue_range
+        self.hue_prob = hue_prob
+
+    def __call__(self, data):
+        brightness_lower = 1 - self.brightness_range
+        brightness_upper = 1 + self.brightness_range
+        contrast_lower = 1 - self.contrast_range
+        contrast_upper = 1 + self.contrast_range
+        saturation_lower = 1 - self.saturation_range
+        saturation_upper = 1 + self.saturation_range
+        hue_lower = -self.hue_range
+        hue_upper = self.hue_range
+        ops = [
+            functional.brightness, functional.contrast, functional.saturation,
+            functional.hue
+        ]
+        random.shuffle(ops)
+        params_dict = {
+            'brightness': {
+                'brightness_lower': brightness_lower,
+                'brightness_upper': brightness_upper
+            },
+            'contrast': {
+                'contrast_lower': contrast_lower,
+                'contrast_upper': contrast_upper
+            },
+            'saturation': {
+                'saturation_lower': saturation_lower,
+                'saturation_upper': saturation_upper
+            },
+            'hue': {
+                'hue_lower': hue_lower,
+                'hue_upper': hue_upper
+            }
+        }
+        prob_dict = {
+            'brightness': self.brightness_prob,
+            'contrast': self.contrast_prob,
+            'saturation': self.saturation_prob,
+            'hue': self.hue_prob
+        }
+
+        im = data['img'].astype('uint8')
+        im = Image.fromarray(im)
+        for id in range(len(ops)):
+            params = params_dict[ops[id].__name__]
+            params['im'] = im
+            prob = prob_dict[ops[id].__name__]
+            if np.random.uniform(0, 1) < prob:
+                im = ops[id](**params)
+        data['img'] = np.asarray(im)
+
+        for key in data.get('gt_fields', []):
+            if key in ['alpha', 'trimap']:
+                continue
+            else:
+                im = data[key].astype('uint8')
+                im = Image.fromarray(im)
+                for id in range(len(ops)):
+                    params = params_dict[ops[id].__name__]
+                    params['im'] = im
+                    prob = prob_dict[ops[id].__name__]
+                    if np.random.uniform(0, 1) < prob:
+                        im = ops[id](**params)
+                data[key] = np.asarray(im)
+        return data
+
+
+if __name__ == "__main__":
+    transforms = [RandomDistort()]
+    transforms = Compose(transforms)
+    fg_path = '/ssd1/home/chenguowei01/github/PaddleSeg/contrib/matting/data/matting/human_matting/Distinctions-646/train/fg/13(2).png'
+    alpha_path = fg_path.replace('fg', 'alpha')
+    bg_path = '/ssd1/home/chenguowei01/github/PaddleSeg/contrib/matting/data/matting/human_matting/bg/unsplash_bg/attic/photo-1443884590026-2e4d21aee71c?crop=entropy&cs=tinysrgb&fit=max&fm=jpg&ixid=MnwxMjA3fDB8MXxzZWFyY2h8Nzh8fGF0dGljfGVufDB8fHx8MTYyOTY4MDcxNQ&ixlib=rb-1.2.1&q=80&w=400.jpg'
+    data = {}
+    data['fg'] = cv2.imread(fg_path)
+    data['bg'] = cv2.imread(bg_path)
+    h, w, c = data['fg'].shape
+    data['bg'] = cv2.resize(data['bg'], (w, h))
+    alpha = cv2.imread(alpha_path)
+    data['alpha'] = alpha[:, :, 0]
+    alpha = alpha / 255.
+    data['img'] = alpha * data['fg'] + (1 - alpha) * data['bg']
+
+    data['gt_fields'] = ['fg', 'bg']
+    print(data['img'].shape)
+    for key in data['gt_fields']:
+        print(data[key].shape)
+#     import pdb
+#     pdb.set_trace()
+    data = transforms(data)
+    print(data['img'].dtype, data['img'].shape)
+    cv2.imwrite('distort_img.jpg', data['img'].transpose([1, 2, 0]))

matting/utils.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+
+def get_files(root_path):
+    res = []
+    for root, dirs, files in os.walk(root_path, followlinks=True):
+        for f in files:
+            if f.endswith(('.jpg', '.png', '.jpeg', 'JPG')):
+                res.append(os.path.join(root, f))
+    return res
+
+
+def get_image_list(image_path):
+    """Get image list"""
+    valid_suffix = [
+        '.JPEG', '.jpeg', '.JPG', '.jpg', '.BMP', '.bmp', '.PNG', '.png'
+    ]
+    image_list = []
+    image_dir = None
+    if os.path.isfile(image_path):
+        if os.path.splitext(image_path)[-1] in valid_suffix:
+            image_list.append(image_path)
+        else:
+            image_dir = os.path.dirname(image_path)
+            with open(image_path, 'r') as f:
+                for line in f:
+                    line = line.strip()
+                    if len(line.split()) > 1:
+                        raise RuntimeError(
+                            'There should be only one image path per line in `image_path` file. Wrong line: {}'
+                            .format(line))
+                    image_list.append(os.path.join(image_dir, line))
+    elif os.path.isdir(image_path):
+        image_dir = image_path
+        for root, dirs, files in os.walk(image_path):
+            for f in files:
+                if '.ipynb_checkpoints' in root:
+                    continue
+                if os.path.splitext(f)[-1] in valid_suffix:
+                    image_list.append(os.path.join(root, f))
+        image_list.sort()
+    else:
+        raise FileNotFoundError(
+            '`image_path` is not found. it should be an image file or a directory including images'
+        )
+
+    if len(image_list) == 0:
+        raise RuntimeError('There are not image file in `image_path`')
+
+    return image_list, image_dir
+
+
+def mkdir(path):
+    sub_dir = os.path.dirname(path)
+    if not os.path.exists(sub_dir):
+        os.makedirs(sub_dir)

requirements.txt

@@ -0,0 +1,2 @@
+paddlepaddle
+paddleseg