class PolyBlock4a(nn.Layer): """ PolyNet type Mixed-4a block. Parameters: ---------- data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, data_format="channels_last", **kwargs): super(PolyBlock4a, self).__init__(**kwargs) self.branches = Concurrent(data_format=data_format, name="branches") self.branches.add( ConvSeqBranch(in_channels=160, out_channels_list=(64, 96), kernel_size_list=(1, 3), strides_list=(1, 1), padding_list=(0, 0), data_format=data_format, name="branch1")) self.branches.add( ConvSeqBranch(in_channels=160, out_channels_list=(64, 64, 64, 96), kernel_size_list=(1, (7, 1), (1, 7), 3), strides_list=(1, 1, 1, 1), padding_list=(0, (3, 0), (0, 3), 0), data_format=data_format, name="branch2")) def call(self, x, training=None): x = self.branches(x, training=training) return x
class LffdDetectionBlock(nn.Layer): """ LFFD specific detection block. Parameters: ---------- in_channels : int Number of input channels. mid_channels : int Number of middle channels. use_bias : bool Whether the layer uses a bias vector. use_bn : bool Whether to use BatchNorm layer. data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, in_channels, mid_channels, use_bias, use_bn, data_format="channels_last", **kwargs): super(LffdDetectionBlock, self).__init__(**kwargs) self.conv = conv1x1_block( in_channels=in_channels, out_channels=mid_channels, use_bias=use_bias, use_bn=use_bn, data_format=data_format, name="conv") self.branches = Concurrent( data_format=data_format, name="branches") self.branches.add(LffdDetectionBranch( in_channels=mid_channels, out_channels=4, use_bias=use_bias, use_bn=use_bn, data_format=data_format, name="bbox_branch")) self.branches.add(LffdDetectionBranch( in_channels=mid_channels, out_channels=2, use_bias=use_bias, use_bn=use_bn, data_format=data_format, name="score_branch")) def call(self, x, training=None): x = self.conv(x, training=training) x = self.branches(x, training=training) return x
class StemBlock(nn.Layer): """ PeleeNet stem block. Parameters: ---------- in_channels : int Number of input channels. out_channels : int Number of output channels. data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, in_channels, out_channels, data_format="channels_last", **kwargs): super(StemBlock, self).__init__(**kwargs) mid1_channels = out_channels // 2 mid2_channels = out_channels * 2 self.first_conv = conv3x3_block(in_channels=in_channels, out_channels=out_channels, strides=2, data_format=data_format, name="first_conv") self.branches = Concurrent(data_format=data_format, name="branches") self.branches.add( PeleeBranch1(in_channels=out_channels, out_channels=out_channels, mid_channels=mid1_channels, strides=2, data_format=data_format, name="branch1")) self.branches.add( MaxPool2d(pool_size=2, strides=2, padding=0, data_format=data_format, name="branch2")) self.last_conv = conv1x1_block(in_channels=mid2_channels, out_channels=out_channels, data_format=data_format, name="last_conv") def call(self, x, training=None): x = self.first_conv(x, training=training) x = self.branches(x, training=training) x = self.last_conv(x, training=training) return x
class AtrousSpatialPyramidPooling(nn.Layer): """ Atrous Spatial Pyramid Pooling (ASPP) module. Parameters: ---------- in_channels : int Number of input channels. upscale_out_size : tuple of 2 int Spatial size of the input tensor for the bilinear upsampling operation. data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, in_channels, upscale_out_size, data_format="channels_last", **kwargs): super(AtrousSpatialPyramidPooling, self).__init__(**kwargs) atrous_rates = [12, 24, 36] assert (in_channels % 8 == 0) mid_channels = in_channels // 8 project_in_channels = 5 * mid_channels self.branches = Concurrent( data_format=data_format, name="branches") self.branches.add(conv1x1_block( in_channels=in_channels, out_channels=mid_channels, data_format=data_format, name="branch1")) for i, atrous_rate in enumerate(atrous_rates): self.branches.add(conv3x3_block( in_channels=in_channels, out_channels=mid_channels, padding=atrous_rate, dilation=atrous_rate, data_format=data_format, name="branch{}".format(i + 2))) self.branches.add(ASPPAvgBranch( in_channels=in_channels, out_channels=mid_channels, upscale_out_size=upscale_out_size, data_format=data_format, name="branch5")) self.conv = conv1x1_block( in_channels=project_in_channels, out_channels=mid_channels, data_format=data_format, name="conv") self.dropout = nn.Dropout( rate=0.5, name="dropout") def call(self, x, training=None): x = self.branches(x, training=training) x = self.conv(x, training=training) x = self.dropout(x, training=training) return x
class PyramidPooling(nn.Layer): """ Pyramid Pooling module. Parameters: ---------- in_channels : int Number of input channels. upscale_out_size : tuple of 2 int Spatial size of the input tensor for the bilinear upsampling operation. data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, in_channels, upscale_out_size, data_format="channels_last", **kwargs): super(PyramidPooling, self).__init__(**kwargs) pool_out_sizes = [1, 2, 3, 6] assert (len(pool_out_sizes) == 4) assert (in_channels % 4 == 0) mid_channels = in_channels // 4 self.branches = Concurrent( data_format=data_format, name="branches") self.branches.add(Identity(name="branch1")) for i, pool_out_size in enumerate(pool_out_sizes): self.branches.add(PyramidPoolingBranch( in_channels=in_channels, out_channels=mid_channels, pool_out_size=pool_out_size, upscale_out_size=upscale_out_size, data_format=data_format, name="branch{}".format(i + 2))) def call(self, x, training=None): x = self.branches(x, training=training) return x
class PolyBlock5a(nn.Layer): """ PolyNet type Mixed-5a block. Parameters: ---------- data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, data_format="channels_last", **kwargs): super(PolyBlock5a, self).__init__(**kwargs) self.branches = Concurrent(data_format=data_format, name="branches") self.branches.add( MaxPoolBranch(data_format=data_format, name="branch1")) self.branches.add( Conv3x3Branch(in_channels=192, out_channels=192, data_format=data_format, name="branch2")) def call(self, x, training=None): x = self.branches(x, training=training) return x
class TwoWayCBlock(nn.Layer): """ PolyNet type Inception-C block. Parameters: ---------- data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, data_format="channels_last", **kwargs): super(TwoWayCBlock, self).__init__(**kwargs) in_channels = 2048 self.branches = Concurrent(data_format=data_format, name="branches") self.branches.add( ConvSeqBranch(in_channels=in_channels, out_channels_list=(192, 224, 256), kernel_size_list=(1, (1, 3), (3, 1)), strides_list=(1, 1, 1), padding_list=(0, (0, 1), (1, 0)), data_format=data_format, name="branch1")) self.branches.add( Conv1x1Branch(in_channels=in_channels, out_channels=192, data_format=data_format, name="branch2")) self.conv = conv1x1_block(in_channels=448, out_channels=in_channels, activation=None, data_format=data_format, name="conv") def call(self, x, training=None): x = self.branches(x, training=training) x = self.conv(x, training=training) return x
def __init__(self, backbone, backbone_out_channels, channels, return_heatmap=False, topk=40, in_channels=3, in_size=(512, 512), classes=80, data_format="channels_last", **kwargs): super(CenterNet, self).__init__(**kwargs) self.in_size = in_size self.in_channels = in_channels self.return_heatmap = return_heatmap self.data_format = data_format self.backbone = backbone self.backbone._name = "backbone" self.decoder = tf.keras.Sequential(name="decoder") in_channels = backbone_out_channels for i, out_channels in enumerate(channels): self.decoder.add( CenterNetDecoderUnit(in_channels=in_channels, out_channels=out_channels, data_format=data_format, name="unit{}".format(i + 1))) in_channels = out_channels heads = Concurrent(data_format=data_format, name="heads") heads.add( CenterNetHeatmapBlock(in_channels=in_channels, out_channels=classes, do_nms=(not self.return_heatmap), data_format=data_format, name="heapmap_block")) heads.add( CenterNetHeadBlock(in_channels=in_channels, out_channels=2, data_format=data_format, name="wh_block")) heads.add( CenterNetHeadBlock(in_channels=in_channels, out_channels=2, data_format=data_format, name="reg_block")) self.decoder.add(heads) if not self.return_heatmap: self.heatmap_max_det = CenterNetHeatmapMaxDet( topk=topk, scale=4, data_format=data_format, name="heatmap_max_det")
class ReductionBUnit(nn.Layer): """ PolyNet type Reduction-B unit. Parameters: ---------- data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, data_format="channels_last", **kwargs): super(ReductionBUnit, self).__init__(**kwargs) in_channels = 1152 self.branches = Concurrent(data_format=data_format, name="branches") self.branches.add( ConvSeqBranch(in_channels=in_channels, out_channels_list=(256, 256, 256), kernel_size_list=(1, 3, 3), strides_list=(1, 1, 2), padding_list=(0, 1, 0), data_format=data_format, name="branch1")) self.branches.add( ConvSeqBranch(in_channels=in_channels, out_channels_list=(256, 256), kernel_size_list=(1, 3), strides_list=(1, 2), padding_list=(0, 0), data_format=data_format, name="branch2")) self.branches.add( ConvSeqBranch(in_channels=in_channels, out_channels_list=(256, 384), kernel_size_list=(1, 3), strides_list=(1, 2), padding_list=(0, 0), data_format=data_format, name="branch3")) self.branches.add( MaxPoolBranch(data_format=data_format, name="branch4")) def call(self, x, training=None): x = self.branches(x, training=training) return x
class ESPBlock(nn.Layer): """ ESP block, which is based on the following principle: Reduce ---> Split ---> Transform --> Merge. Parameters: ---------- in_channels : int Number of input channels. out_channels : int Number of output channels. kernel_sizes : list of int Convolution window size for branches. scale_factors : list of int Scale factor for branches. use_residual : bool Whether to use residual connection. in_size : tuple of 2 int Spatial size of the output tensor for the bilinear upsampling operation. bn_eps : float Small float added to variance in Batch norm. data_format : str, default 'channels_last' The ordering of the dimensions in tensors. """ def __init__(self, in_channels, out_channels, kernel_sizes, scale_factors, use_residual, in_size, bn_eps, data_format="channels_last", **kwargs): super(ESPBlock, self).__init__(**kwargs) self.use_residual = use_residual groups = len(kernel_sizes) mid_channels = int(out_channels / groups) res_channels = out_channels - groups * mid_channels self.conv = conv1x1(in_channels=in_channels, out_channels=mid_channels, groups=groups, data_format=data_format, name="conv") self.c_shuffle = ChannelShuffle(channels=mid_channels, groups=groups, data_format=data_format, name="c_shuffle") self.branches = Concurrent(data_format=data_format, name="branches") for i in range(groups): out_channels_i = (mid_channels + res_channels) if i == 0 else mid_channels self.branches.add( SBBlock(in_channels=mid_channels, out_channels=out_channels_i, kernel_size=kernel_sizes[i], scale_factor=scale_factors[i], size=in_size, bn_eps=bn_eps, data_format=data_format, name="branch{}".format(i + 1))) self.preactiv = PreActivation(in_channels=out_channels, bn_eps=bn_eps, data_format=data_format, name="preactiv") def call(self, x, training=None): if self.use_residual: identity = x x = self.conv(x) x = self.c_shuffle(x) x = self.branches(x, training=None) if self.use_residual: x = identity + x x = self.preactiv(x, training=None) return x