ppmatting/models/human_matting.py

# copyright (c) 2022 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.

from collections import defaultdict
import time

import paddle
import paddle.nn as nn
import paddle.nn.functional as F
import paddleseg
from paddleseg.models import layers
from paddleseg import utils
from paddleseg.cvlibs import manager

from ppmatting.models.losses import MRSD


def conv_up_psp(in_channels, out_channels, up_sample):
 return nn.Sequential(
 layers.ConvBNReLU(
 in_channels, out_channels, 3, padding=1),
 nn.Upsample(
 scale_factor=up_sample, mode='bilinear', align_corners=False))


@manager.MODELS.add_component
class HumanMatting(nn.Layer):
 """A model for """

 def __init__(self,
 backbone,
 pretrained=None,
 backbone_scale=0.25,
 refine_kernel_size=3,
 if_refine=True):
 super().__init__()
 if if_refine:
 if backbone_scale > 0.5:
 raise ValueError(
 'Backbone_scale should not be greater than 1/2, but it is {}'
 .format(backbone_scale))
 else:
 backbone_scale = 1

 self.backbone = backbone
 self.backbone_scale = backbone_scale
 self.pretrained = pretrained
 self.if_refine = if_refine
 if if_refine:
 self.refiner = Refiner(kernel_size=refine_kernel_size)
 self.loss_func_dict = None

 self.backbone_channels = backbone.feat_channels
 ######################
 ### Decoder part - Glance
 ######################
 self.psp_module = layers.PPModule(
 self.backbone_channels[-1],
 512,
 bin_sizes=(1, 3, 5),
 dim_reduction=False,
 align_corners=False)
 self.psp4 = conv_up_psp(512, 256, 2)
 self.psp3 = conv_up_psp(512, 128, 4)
 self.psp2 = conv_up_psp(512, 64, 8)
 self.psp1 = conv_up_psp(512, 64, 16)
 # stage 5g
 self.decoder5_g = nn.Sequential(
 layers.ConvBNReLU(
 512 + self.backbone_channels[-1], 512, 3, padding=1),
 layers.ConvBNReLU(
 512, 512, 3, padding=2, dilation=2),
 layers.ConvBNReLU(
 512, 256, 3, padding=2, dilation=2),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 4g
 self.decoder4_g = nn.Sequential(
 layers.ConvBNReLU(
 512, 256, 3, padding=1),
 layers.ConvBNReLU(
 256, 256, 3, padding=1),
 layers.ConvBNReLU(
 256, 128, 3, padding=1),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 3g
 self.decoder3_g = nn.Sequential(
 layers.ConvBNReLU(
 256, 128, 3, padding=1),
 layers.ConvBNReLU(
 128, 128, 3, padding=1),
 layers.ConvBNReLU(
 128, 64, 3, padding=1),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 2g
 self.decoder2_g = nn.Sequential(
 layers.ConvBNReLU(
 128, 128, 3, padding=1),
 layers.ConvBNReLU(
 128, 128, 3, padding=1),
 layers.ConvBNReLU(
 128, 64, 3, padding=1),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 1g
 self.decoder1_g = nn.Sequential(
 layers.ConvBNReLU(
 128, 64, 3, padding=1),
 layers.ConvBNReLU(
 64, 64, 3, padding=1),
 layers.ConvBNReLU(
 64, 64, 3, padding=1),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 0g
 self.decoder0_g = nn.Sequential(
 layers.ConvBNReLU(
 64, 64, 3, padding=1),
 layers.ConvBNReLU(
 64, 64, 3, padding=1),
 nn.Conv2D(
 64, 3, 3, padding=1))

 ##########################
 ### Decoder part - FOCUS
 ##########################
 self.bridge_block = nn.Sequential(
 layers.ConvBNReLU(
 self.backbone_channels[-1], 512, 3, dilation=2, padding=2),
 layers.ConvBNReLU(
 512, 512, 3, dilation=2, padding=2),
 layers.ConvBNReLU(
 512, 512, 3, dilation=2, padding=2))
 # stage 5f
 self.decoder5_f = nn.Sequential(
 layers.ConvBNReLU(
 512 + self.backbone_channels[-1], 512, 3, padding=1),
 layers.ConvBNReLU(
 512, 512, 3, padding=2, dilation=2),
 layers.ConvBNReLU(
 512, 256, 3, padding=2, dilation=2),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 4f
 self.decoder4_f = nn.Sequential(
 layers.ConvBNReLU(
 256 + self.backbone_channels[-2], 256, 3, padding=1),
 layers.ConvBNReLU(
 256, 256, 3, padding=1),
 layers.ConvBNReLU(
 256, 128, 3, padding=1),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 3f
 self.decoder3_f = nn.Sequential(
 layers.ConvBNReLU(
 128 + self.backbone_channels[-3], 128, 3, padding=1),
 layers.ConvBNReLU(
 128, 128, 3, padding=1),
 layers.ConvBNReLU(
 128, 64, 3, padding=1),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 2f
 self.decoder2_f = nn.Sequential(
 layers.ConvBNReLU(
 64 + self.backbone_channels[-4], 128, 3, padding=1),
 layers.ConvBNReLU(
 128, 128, 3, padding=1),
 layers.ConvBNReLU(
 128, 64, 3, padding=1),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 1f
 self.decoder1_f = nn.Sequential(
 layers.ConvBNReLU(
 64 + self.backbone_channels[-5], 64, 3, padding=1),
 layers.ConvBNReLU(
 64, 64, 3, padding=1),
 layers.ConvBNReLU(
 64, 64, 3, padding=1),
 nn.Upsample(
 scale_factor=2, mode='bilinear', align_corners=False))
 # stage 0f
 self.decoder0_f = nn.Sequential(
 layers.ConvBNReLU(
 64, 64, 3, padding=1),
 layers.ConvBNReLU(
 64, 64, 3, padding=1),
 nn.Conv2D(
 64, 1 + 1 + 32, 3, padding=1))
 self.init_weight()

 def forward(self, data):
 src = data['img']
 src_h, src_w = paddle.shape(src)[2:]
 if self.if_refine:
 # It is not need when exporting.
 if isinstance(src_h, paddle.Tensor):
 if (src_h % 4 != 0) or (src_w % 4) != 0:
 raise ValueError(
 'The input image must have width and height that are divisible by 4'
 )

 # Downsample src for backbone
 src_sm = F.interpolate(
 src,
 scale_factor=self.backbone_scale,
 mode='bilinear',
 align_corners=False)

 # Base
 fea_list = self.backbone(src_sm)
 ##########################
 ### Decoder part - GLANCE
 ##########################
 #psp: N, 512, H/32, W/32
 psp = self.psp_module(fea_list[-1])
 #d6_g: N, 512, H/16, W/16
 d5_g = self.decoder5_g(paddle.concat((psp, fea_list[-1]), 1))
 #d5_g: N, 512, H/8, W/8
 d4_g = self.decoder4_g(paddle.concat((self.psp4(psp), d5_g), 1))
 #d4_g: N, 256, H/4, W/4
 d3_g = self.decoder3_g(paddle.concat((self.psp3(psp), d4_g), 1))
 #d4_g: N, 128, H/2, W/2
 d2_g = self.decoder2_g(paddle.concat((self.psp2(psp), d3_g), 1))
 #d2_g: N, 64, H, W
 d1_g = self.decoder1_g(paddle.concat((self.psp1(psp), d2_g), 1))
 #d0_g: N, 3, H, W
 d0_g = self.decoder0_g(d1_g)
 # The 1st channel is foreground. The 2nd is transition region. The 3rd is background.
 # glance_sigmoid = F.sigmoid(d0_g)
 glance_sigmoid = F.softmax(d0_g, axis=1)

 ##########################
 ### Decoder part - FOCUS
 ##########################
 bb = self.bridge_block(fea_list[-1])
 #bg: N, 512, H/32, W/32
 d5_f = self.decoder5_f(paddle.concat((bb, fea_list[-1]), 1))
 #d5_f: N, 256, H/16, W/16
 d4_f = self.decoder4_f(paddle.concat((d5_f, fea_list[-2]), 1))
 #d4_f: N, 128, H/8, W/8
 d3_f = self.decoder3_f(paddle.concat((d4_f, fea_list[-3]), 1))
 #d3_f: N, 64, H/4, W/4
 d2_f = self.decoder2_f(paddle.concat((d3_f, fea_list[-4]), 1))
 #d2_f: N, 64, H/2, W/2
 d1_f = self.decoder1_f(paddle.concat((d2_f, fea_list[-5]), 1))
 #d1_f: N, 64, H, W
 d0_f = self.decoder0_f(d1_f)
 #d0_f: N, 1, H, W
 focus_sigmoid = F.sigmoid(d0_f[:, 0:1, :, :])
 pha_sm = self.fusion(glance_sigmoid, focus_sigmoid)
 err_sm = d0_f[:, 1:2, :, :]
 err_sm = paddle.clip(err_sm, 0., 1.)
 hid_sm = F.relu(d0_f[:, 2:, :, :])

 # Refiner
 if self.if_refine:
 pha = self.refiner(
 src=src, pha=pha_sm, err=err_sm, hid=hid_sm, tri=glance_sigmoid)
 # Clamp outputs
 pha = paddle.clip(pha, 0., 1.)

 if self.training:
 logit_dict = {
 'glance': glance_sigmoid,
 'focus': focus_sigmoid,
 'fusion': pha_sm,
 'error': err_sm
 }
 if self.if_refine:
 logit_dict['refine'] = pha
 loss_dict = self.loss(logit_dict, data)
 return logit_dict, loss_dict
 else:
 return pha if self.if_refine else pha_sm

 def loss(self, logit_dict, label_dict, loss_func_dict=None):
 if loss_func_dict is None:
 if self.loss_func_dict is None:
 self.loss_func_dict = defaultdict(list)
 self.loss_func_dict['glance'].append(nn.NLLLoss())
 self.loss_func_dict['focus'].append(MRSD())
 self.loss_func_dict['cm'].append(MRSD())
 self.loss_func_dict['err'].append(paddleseg.models.MSELoss())
 self.loss_func_dict['refine'].append(paddleseg.models.L1Loss())
 else:
 self.loss_func_dict = loss_func_dict

 loss = {}

 # glance loss computation
 # get glance label
 glance_label = F.interpolate(
 label_dict['trimap'],
 logit_dict['glance'].shape[2:],
 mode='nearest',
 align_corners=False)
 glance_label_trans = (glance_label == 128).astype('int64')
 glance_label_bg = (glance_label == 0).astype('int64')
 glance_label = glance_label_trans + glance_label_bg * 2
 loss_glance = self.loss_func_dict['glance'][0](
 paddle.log(logit_dict['glance'] + 1e-6), glance_label.squeeze(1))
 loss['glance'] = loss_glance

 # focus loss computation
 focus_label = F.interpolate(
 label_dict['alpha'],
 logit_dict['focus'].shape[2:],
 mode='bilinear',
 align_corners=False)
 loss_focus = self.loss_func_dict['focus'][0](
 logit_dict['focus'], focus_label, glance_label_trans)
 loss['focus'] = loss_focus

 # collaborative matting loss
 loss_cm_func = self.loss_func_dict['cm']
 # fusion_sigmoid loss
 loss_cm = loss_cm_func[0](logit_dict['fusion'], focus_label)
 loss['cm'] = loss_cm

 # error loss
 err = F.interpolate(
 logit_dict['error'],
 label_dict['alpha'].shape[2:],
 mode='bilinear',
 align_corners=False)
 err_label = (F.interpolate(
 logit_dict['fusion'],
 label_dict['alpha'].shape[2:],
 mode='bilinear',
 align_corners=False) - label_dict['alpha']).abs()
 loss_err = self.loss_func_dict['err'][0](err, err_label)
 loss['err'] = loss_err

 loss_all = 0.25 * loss_glance + 0.25 * loss_focus + 0.25 * loss_cm + loss_err

 # refine loss
 if self.if_refine:
 loss_refine = self.loss_func_dict['refine'][0](logit_dict['refine'],
 label_dict['alpha'])
 loss['refine'] = loss_refine
 loss_all = loss_all + loss_refine

 loss['all'] = loss_all
 return loss

 def fusion(self, glance_sigmoid, focus_sigmoid):
 # glance_sigmoid [N, 3, H, W].
 # In index, 0 is foreground, 1 is transition, 2 is backbone.
 # After fusion, the foreground is 1, the background is 0, and the transion is between (0, 1).
 index = paddle.argmax(glance_sigmoid, axis=1, keepdim=True)
 transition_mask = (index == 1).astype('float32')
 fg = (index == 0).astype('float32')
 fusion_sigmoid = focus_sigmoid * transition_mask + fg
 return fusion_sigmoid

 def init_weight(self):
 if self.pretrained is not None:
 utils.load_entire_model(self, self.pretrained)


class Refiner(nn.Layer):
 '''
 Refiner refines the coarse output to full resolution.

 Args:
 kernel_size: The convolution kernel_size. Options: [1, 3]. Default: 3.
 '''

 def __init__(self, kernel_size=3):
 super().__init__()
 if kernel_size not in [1, 3]:
 raise ValueError("kernel_size must be in [1, 3]")

 self.kernel_size = kernel_size

 channels = [32, 24, 16, 12, 1]
 self.conv1 = layers.ConvBNReLU(
 channels[0] + 4 + 3,
 channels[1],
 kernel_size,
 padding=0,
 bias_attr=False)
 self.conv2 = layers.ConvBNReLU(
 channels[1], channels[2], kernel_size, padding=0, bias_attr=False)
 self.conv3 = layers.ConvBNReLU(
 channels[2] + 3,
 channels[3],
 kernel_size,
 padding=0,
 bias_attr=False)
 self.conv4 = nn.Conv2D(
 channels[3], channels[4], kernel_size, padding=0, bias_attr=True)

 def forward(self, src, pha, err, hid, tri):
 '''
 Args：
 src: (B, 3, H, W) full resolution source image.
 pha: (B, 1, Hc, Wc) coarse alpha prediction.
 err: (B, 1, Hc, Hc) coarse error prediction.
 hid: (B, 32, Hc, Hc) coarse hidden encoding.
 tri: (B, 1, Hc, Hc) trimap prediction.
 '''
 h_full, w_full = paddle.shape(src)[2:]
 h_half, w_half = h_full // 2, w_full // 2
 h_quat, w_quat = h_full // 4, w_full // 4

 x = paddle.concat([hid, pha, tri], axis=1)
 x = F.interpolate(
 x,
 paddle.concat((h_half, w_half)),
 mode='bilinear',
 align_corners=False)
 y = F.interpolate(
 src,
 paddle.concat((h_half, w_half)),
 mode='bilinear',
 align_corners=False)

 if self.kernel_size == 3:
 x = F.pad(x, [3, 3, 3, 3])
 y = F.pad(y, [3, 3, 3, 3])

 x = self.conv1(paddle.concat([x, y], axis=1))
 x = self.conv2(x)

 if self.kernel_size == 3:
 x = F.interpolate(x, paddle.concat((h_full + 4, w_full + 4)))
 y = F.pad(src, [2, 2, 2, 2])
 else:
 x = F.interpolate(
 x, paddle.concat((h_full, w_full)), mode='nearest')
 y = src

 x = self.conv3(paddle.concat([x, y], axis=1))
 x = self.conv4(x)

 pha = x
 return pha