Spaces:
Configuration error
Configuration error
| # 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.nn as nn | |
| from paddleseg import utils | |
| from paddleseg.cvlibs import manager | |
| class PortraitNet(nn.Layer): | |
| """ | |
| The PortraitNet implementation based on PaddlePaddle. | |
| The original article refers to | |
| Song-Hai Zhanga, Xin Donga, Jia Lib, Ruilong Lia, Yong-Liang Yangc | |
| "PortraitNet: Real-time Portrait Segmentation Network for Mobile Device" | |
| (https://www.yongliangyang.net/docs/mobilePotrait_c&g19.pdf). | |
| Args: | |
| num_classes (int, optional): The unique number of target classes. Default: 2. | |
| backbone (Paddle.nn.Layer): Backbone network, currently support MobileNetV2. | |
| add_edge (bool, optional): Whether output to edge. Default: False | |
| pretrained (str, optional): The path or url of pretrained model. Default: None | |
| """ | |
| def __init__(self, | |
| num_classes, | |
| backbone, | |
| min_channel=16, | |
| channel_ratio=1.0, | |
| add_edge=False, | |
| pretrained=None): | |
| super(PortraitNet, self).__init__() | |
| self.backbone = backbone | |
| self.head = PortraitNetHead(num_classes, min_channel, channel_ratio, | |
| add_edge) | |
| self.pretrained = pretrained | |
| self.init_weight() | |
| def forward(self, x): | |
| feat_list = self.backbone(x) | |
| logits_list = self.head(feat_list) | |
| return [logits_list] | |
| def init_weight(self): | |
| if self.pretrained is not None: | |
| utils.load_entire_model(self, self.pretrained) | |
| class PortraitNetHead(nn.Layer): | |
| def __init__(self, | |
| num_classes, | |
| min_channel=16, | |
| channel_ratio=1.0, | |
| add_edge=False): | |
| super().__init__() | |
| self.min_channel = min_channel | |
| self.channel_ratio = channel_ratio | |
| self.add_edge = add_edge | |
| self.deconv1 = nn.Conv2DTranspose( | |
| self.depth(96), | |
| self.depth(96), | |
| groups=1, | |
| kernel_size=4, | |
| stride=2, | |
| padding=1, | |
| bias_attr=False) | |
| self.deconv2 = nn.Conv2DTranspose( | |
| self.depth(32), | |
| self.depth(32), | |
| groups=1, | |
| kernel_size=4, | |
| stride=2, | |
| padding=1, | |
| bias_attr=False) | |
| self.deconv3 = nn.Conv2DTranspose( | |
| self.depth(24), | |
| self.depth(24), | |
| groups=1, | |
| kernel_size=4, | |
| stride=2, | |
| padding=1, | |
| bias_attr=False) | |
| self.deconv4 = nn.Conv2DTranspose( | |
| self.depth(16), | |
| self.depth(16), | |
| groups=1, | |
| kernel_size=4, | |
| stride=2, | |
| padding=1, | |
| bias_attr=False) | |
| self.deconv5 = nn.Conv2DTranspose( | |
| self.depth(8), | |
| self.depth(8), | |
| groups=1, | |
| kernel_size=4, | |
| stride=2, | |
| padding=1, | |
| bias_attr=False) | |
| self.transit1 = ResidualBlock(self.depth(320), self.depth(96)) | |
| self.transit2 = ResidualBlock(self.depth(96), self.depth(32)) | |
| self.transit3 = ResidualBlock(self.depth(32), self.depth(24)) | |
| self.transit4 = ResidualBlock(self.depth(24), self.depth(16)) | |
| self.transit5 = ResidualBlock(self.depth(16), self.depth(8)) | |
| self.pred = nn.Conv2D( | |
| self.depth(8), num_classes, 3, 1, 1, bias_attr=False) | |
| if self.add_edge: | |
| self.edge = nn.Conv2D( | |
| self.depth(8), num_classes, 3, 1, 1, bias_attr=False) | |
| def depth(self, channels): | |
| min_channel = min(channels, self.min_channel) | |
| return max(min_channel, int(channels * self.channel_ratio)) | |
| def forward(self, feat_list): | |
| feature_1_4, feature_1_8, feature_1_16, feature_1_32 = feat_list | |
| up_1_16 = self.deconv1(self.transit1(feature_1_32)) | |
| up_1_8 = self.deconv2(self.transit2(feature_1_16 + up_1_16)) | |
| up_1_4 = self.deconv3(self.transit3(feature_1_8 + up_1_8)) | |
| up_1_2 = self.deconv4(self.transit4(feature_1_4 + up_1_4)) | |
| up_1_1 = self.deconv5(self.transit5(up_1_2)) | |
| pred = self.pred(up_1_1) | |
| if self.add_edge: | |
| edge = self.edge(up_1_1) | |
| return pred, edge | |
| else: | |
| return pred | |
| class ConvDw(nn.Layer): | |
| def __init__(self, inp, oup, kernel, stride): | |
| super(ConvDw, self).__init__() | |
| self.conv = nn.Sequential( | |
| nn.Conv2D( | |
| inp, | |
| inp, | |
| kernel, | |
| stride, (kernel - 1) // 2, | |
| groups=inp, | |
| bias_attr=False), | |
| nn.BatchNorm2D( | |
| num_features=inp, epsilon=1e-05, momentum=0.1), | |
| nn.ReLU(), | |
| nn.Conv2D( | |
| inp, oup, 1, 1, 0, bias_attr=False), | |
| nn.BatchNorm2D( | |
| num_features=oup, epsilon=1e-05, momentum=0.1), | |
| nn.ReLU(), ) | |
| def forward(self, x): | |
| return self.conv(x) | |
| class ResidualBlock(nn.Layer): | |
| def __init__(self, inp, oup, stride=1): | |
| super(ResidualBlock, self).__init__() | |
| self.block = nn.Sequential( | |
| ConvDw( | |
| inp, oup, 3, stride=stride), | |
| nn.Conv2D( | |
| in_channels=oup, | |
| out_channels=oup, | |
| kernel_size=3, | |
| stride=1, | |
| padding=1, | |
| groups=oup, | |
| bias_attr=False), | |
| nn.BatchNorm2D( | |
| num_features=oup, epsilon=1e-05, momentum=0.1), | |
| nn.ReLU(), | |
| nn.Conv2D( | |
| in_channels=oup, | |
| out_channels=oup, | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| bias_attr=False), | |
| nn.BatchNorm2D( | |
| num_features=oup, epsilon=1e-05, momentum=0.1), ) | |
| if inp == oup: | |
| self.residual = None | |
| else: | |
| self.residual = nn.Sequential( | |
| nn.Conv2D( | |
| in_channels=inp, | |
| out_channels=oup, | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| bias_attr=False), | |
| nn.BatchNorm2D( | |
| num_features=oup, epsilon=1e-05, momentum=0.1), ) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| residual = x | |
| out = self.block(x) | |
| if self.residual is not None: | |
| residual = self.residual(x) | |
| out += residual | |
| out = self.relu(out) | |
| return out | |