# 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