diff --git a/pdseg/models/modeling/ocnet.py b/pdseg/models/modeling/ocnet.py new file mode 100644 index 0000000000000000000000000000000000000000..01b2e301ee72de13c2b9681354cc333921919bf9 --- /dev/null +++ b/pdseg/models/modeling/ocnet.py @@ -0,0 +1,309 @@ +# coding: utf8 +# copyright (c) 2019 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 __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import paddle +import paddle.fluid as fluid +from paddle.fluid.initializer import MSRA +from paddle.fluid.param_attr import ParamAttr + +from utils.config import cfg + + +def conv_bn_layer(input, filter_size, num_filters, stride=1, padding=1, num_groups=1, if_act=True, name=None): + conv = fluid.layers.conv2d( + input=input, + num_filters=num_filters, + filter_size=filter_size, + stride=stride, + padding=(filter_size - 1) // 2, + groups=num_groups, + act=None, +# param_attr=ParamAttr(initializer=MSRA(), learning_rate=1.0, name=name + '_weights'), + param_attr=ParamAttr(initializer=fluid.initializer.Normal(scale=0.001), learning_rate=1.0, name=name + '_weights'), + bias_attr=False) + bn_name = name + '_bn' + bn = fluid.layers.batch_norm(input=conv, + param_attr=ParamAttr(name=bn_name + "_scale", + initializer=fluid.initializer.Constant(1.0)), + bias_attr=ParamAttr(name=bn_name + "_offset", + initializer=fluid.initializer.Constant(0.0)), + moving_mean_name=bn_name + '_mean', + moving_variance_name=bn_name + '_variance') + if if_act: + bn = fluid.layers.relu(bn) + return bn + +def basic_block(input, num_filters, stride=1, downsample=False, name=None): + residual = input + conv = conv_bn_layer(input=input, filter_size=3, num_filters=num_filters, stride=stride, name=name + '_conv1') + conv = conv_bn_layer(input=conv, filter_size=3, num_filters=num_filters, if_act=False, name=name + '_conv2') + if downsample: + residual = conv_bn_layer(input=input, filter_size=1, num_filters=num_filters, if_act=False, + name=name + '_downsample') + return fluid.layers.elementwise_add(x=residual, y=conv, act='relu') + +def bottleneck_block(input, num_filters, stride=1, downsample=False, name=None): + residual = input + conv = conv_bn_layer(input=input, filter_size=1, num_filters=num_filters, name=name + '_conv1') + conv = conv_bn_layer(input=conv, filter_size=3, num_filters=num_filters, stride=stride, name=name + '_conv2') + conv = conv_bn_layer(input=conv, filter_size=1, num_filters=num_filters * 4, if_act=False, + name=name + '_conv3') + if downsample: + residual = conv_bn_layer(input=input, filter_size=1, num_filters=num_filters * 4, if_act=False, + name=name + '_downsample') + return fluid.layers.elementwise_add(x=residual, y=conv, act='relu') + +def fuse_layers(x, channels, multi_scale_output=True, name=None): + out = [] + for i in range(len(channels) if multi_scale_output else 1): + residual = x[i] + shape = residual.shape + width = shape[-1] + height = shape[-2] + for j in range(len(channels)): + if j > i: + y = conv_bn_layer(x[j], filter_size=1, num_filters=channels[i], if_act=False, + name=name + '_layer_' + str(i + 1) + '_' + str(j + 1)) + y = fluid.layers.resize_bilinear(input=y, out_shape=[height, width]) + residual = fluid.layers.elementwise_add(x=residual, y=y, act=None) + elif j < i: + y = x[j] + for k in range(i - j): + if k == i - j - 1: + y = conv_bn_layer(y, filter_size=3, num_filters=channels[i], stride=2, if_act=False, + name=name + '_layer_' + str(i + 1) + '_' + str(j + 1) + '_' + str(k + 1)) + else: + y = conv_bn_layer(y, filter_size=3, num_filters=channels[j], stride=2, + name=name + '_layer_' + str(i + 1) + '_' + str(j + 1) + '_' + str(k + 1)) + residual = fluid.layers.elementwise_add(x=residual, y=y, act=None) + + residual = fluid.layers.relu(residual) + out.append(residual) + return out + +def branches(x, block_num, channels, name=None): + out = [] + for i in range(len(channels)): + residual = x[i] + for j in range(block_num): + residual = basic_block(residual, channels[i], + name=name + '_branch_layer_' + str(i + 1) + '_' + str(j + 1)) + out.append(residual) + return out + +def high_resolution_module(x, channels, multi_scale_output=True, name=None): + residual = branches(x, 4, channels, name=name) + out = fuse_layers(residual, channels, multi_scale_output=multi_scale_output, name=name) + return out + +def transition_layer(x, in_channels, out_channels, name=None): + num_in = len(in_channels) + num_out = len(out_channels) + out = [] + for i in range(num_out): + if i < num_in: + if in_channels[i] != out_channels[i]: + residual = conv_bn_layer(x[i], filter_size=3, num_filters=out_channels[i], + name=name + '_layer_' + str(i + 1)) + out.append(residual) + else: + out.append(x[i]) + else: + residual = conv_bn_layer(x[-1], filter_size=3, num_filters=out_channels[i], stride=2, + name=name + '_layer_' + str(i + 1)) + out.append(residual) + return out + +def stage(x, num_modules, channels, multi_scale_output=True, name=None): + out = x + for i in range(num_modules): + if i == num_modules - 1 and multi_scale_output == False: + out = high_resolution_module(out, channels, multi_scale_output=False, name=name + '_' + str(i + 1)) + else: + out = high_resolution_module(out, channels, name=name + '_' + str(i + 1)) + + return out + +def layer1(input, name=None): + conv = input + for i in range(4): + conv = bottleneck_block(conv, num_filters=64, downsample=True if i == 0 else False, + name=name + '_' + str(i + 1)) + return conv + +def aux_head(input, last_inp_channels, num_classes): + x = conv_bn_layer(input=input, filter_size=1, num_filters=last_inp_channels, stride=1, padding=0, name='aux_head_conv1') + x = fluid.layers.conv2d( + input=x, + num_filters=num_classes, + filter_size=1, + stride=1, + padding=0, + act=None, +# param_attr=ParamAttr(initializer=MSRA(), learning_rate=1.0, name='aux_head_conv2_weights'), + param_attr=ParamAttr(initializer=fluid.initializer.Normal(scale=0.001), learning_rate=1.0, name='aux_head_conv2_weights'), + bias_attr=ParamAttr(initializer=fluid.initializer.Constant(0.0), name="aux_head_conv2_bias") + ) + return x + +def conv3x3_ocr(input, ocr_mid_channels): + x = conv_bn_layer(input=input, filter_size=3, num_filters=ocr_mid_channels, stride=1, padding=1, name='conv3x3_ocr') + return x + +def f_pixel(input, key_channels): + x = conv_bn_layer(input=input, filter_size=1, num_filters=key_channels, stride=1, padding=0, name='f_pixel_conv1') + x = conv_bn_layer(input=x, filter_size=1, num_filters=key_channels, stride=1, padding=0, name='f_pixel_conv2') + return x + +def f_object(input, key_channels): + x = conv_bn_layer(input=input, filter_size=1, num_filters=key_channels, stride=1, padding=0, name='f_object_conv1') + x = conv_bn_layer(input=x, filter_size=1, num_filters=key_channels, stride=1, padding=0, name='f_object_conv2') + return x + +def f_down(input, key_channels): + x = conv_bn_layer(input=input, filter_size=1, num_filters=key_channels, stride=1, padding=0, name='f_down_conv') + return x + +def f_up(input, in_channels): + x = conv_bn_layer(input=input, filter_size=1, num_filters=in_channels, stride=1, padding=0, name='f_up_conv') + return x + +def object_context_block(x, proxy, in_channels, key_channels, scale): + batch_size, _, h, w = x.shape + if scale > 1: + x = fluid.layers.pool2d(x, pool_size=[scale, scale], pool_type='max') + + query = f_pixel(x, key_channels) + query = fluid.layers.reshape(query, shape=[batch_size, key_channels, query.shape[2]*query.shape[3]]) + query = fluid.layers.transpose(query, perm=[0, 2, 1]) + + key = f_object(proxy, key_channels) + key = fluid.layers.reshape(key, shape=[batch_size, key_channels, key.shape[2]*key.shape[3]]) + + value = f_down(proxy, key_channels) + value = fluid.layers.reshape(value, shape=[batch_size, key_channels, value.shape[2]*value.shape[3]]) + value = fluid.layers.transpose(value, perm=[0, 2, 1]) + + sim_map = fluid.layers.matmul(query, key) + sim_map = (key_channels**-.5) * sim_map + sim_map = fluid.layers.softmax(sim_map, axis=-1) + + context = fluid.layers.matmul(sim_map, value) + context = fluid.layers.transpose(context, perm=[0, 2, 1]) + context = fluid.layers.reshape(context, shape=[batch_size, key_channels, x.shape[2], x.shape[3]]) + context = f_up(context, in_channels) + + if scale > 1: + context = fluid.layers.resize_bilinear(context, out_shape=[h, w]) + + return context + +def ocr_gather_head(feats, probs, scale=1): + feats = fluid.layers.reshape(feats, shape=[feats.shape[0], feats.shape[1], feats.shape[2]*feats.shape[3]]) + feats = fluid.layers.transpose(feats, perm=[0, 2, 1]) + probs = fluid.layers.reshape(probs, shape=[probs.shape[0], probs.shape[1], probs.shape[2]*probs.shape[3]]) + probs = fluid.layers.softmax(scale * probs, axis=2) + ocr_context = fluid.layers.matmul(probs, feats) + ocr_context = fluid.layers.transpose(ocr_context, perm=[0, 2, 1]) + ocr_context = fluid.layers.unsqueeze(ocr_context, axes=[3]) + return ocr_context + +def ocr_distri_head(feats, proxy_feats, ocr_mid_channels, ocr_key_channels, scale=1, dropout=0.05): + context = object_context_block(feats, proxy_feats, ocr_mid_channels, ocr_key_channels, scale) + x = fluid.layers.concat([context, feats], axis=1) + x = conv_bn_layer(input=x, filter_size=1, num_filters=ocr_mid_channels, stride=1, padding=0, name='spatial_ocr_conv') + x = fluid.layers.dropout(x, dropout_prob=dropout) + return x + +def cls_head(input, num_classes): + x = fluid.layers.conv2d( + input=input, + num_filters=num_classes, + filter_size=1, + stride=1, + padding=0, + act=None, +# param_attr=ParamAttr(initializer=MSRA(), learning_rate=1.0, name='cls_head_conv_weights'), + param_attr=ParamAttr(initializer=fluid.initializer.Normal(scale=0.001), learning_rate=1.0, name='cls_head_conv_weights'), + bias_attr=ParamAttr(initializer=fluid.initializer.Constant(0.0), name="cls_head_conv_bias") + ) + return x + +def ocr_module(input, last_inp_channels, num_classes, ocr_mid_channels, ocr_key_channels): + out_aux = aux_head(input, last_inp_channels, num_classes) + feats = conv3x3_ocr(input, ocr_mid_channels) + context = ocr_gather_head(feats, out_aux) + feats = ocr_distri_head(feats, context, ocr_mid_channels, ocr_key_channels) + out = cls_head(feats, num_classes) + return out, out_aux + +def high_resolution_ocr_net(input, num_classes): + + channels_2 = cfg.MODEL.HRNET.STAGE2.NUM_CHANNELS + channels_3 = cfg.MODEL.HRNET.STAGE3.NUM_CHANNELS + channels_4 = cfg.MODEL.HRNET.STAGE4.NUM_CHANNELS + + num_modules_2 = cfg.MODEL.HRNET.STAGE2.NUM_MODULES + num_modules_3 = cfg.MODEL.HRNET.STAGE3.NUM_MODULES + num_modules_4 = cfg.MODEL.HRNET.STAGE4.NUM_MODULES + + ocr_mid_channels = cfg.MODEL.OCR.OCR_MID_CHANNELS + ocr_key_channels = cfg.MODEL.OCR.OCR_KEY_CHANNELS + + last_inp_channels = sum(channels_4) + + x = conv_bn_layer(input=input, filter_size=3, num_filters=64, stride=2, if_act=True, name='layer1_1') + x = conv_bn_layer(input=x, filter_size=3, num_filters=64, stride=2, if_act=True, name='layer1_2') + + la1 = layer1(x, name='layer2') + tr1 = transition_layer([la1], [256], channels_2, name='tr1') + st2 = stage(tr1, num_modules_2, channels_2, name='st2') + tr2 = transition_layer(st2, channels_2, channels_3, name='tr2') + st3 = stage(tr2, num_modules_3, channels_3, name='st3') + tr3 = transition_layer(st3, channels_3, channels_4, name='tr3') + st4 = stage(tr3, num_modules_4, channels_4, name='st4') + + # upsample + shape = st4[0].shape + height, width = shape[-2], shape[-1] + st4[1] = fluid.layers.resize_bilinear( + st4[1], out_shape=[height, width]) + st4[2] = fluid.layers.resize_bilinear( + st4[2], out_shape=[height, width]) + st4[3] = fluid.layers.resize_bilinear( + st4[3], out_shape=[height, width]) + + feats = fluid.layers.concat(st4, axis=1) + + out, out_aux = ocr_module(feats, last_inp_channels, num_classes, ocr_mid_channels, ocr_key_channels) + + out = fluid.layers.resize_bilinear(out, input.shape[2:]) + out_aux = fluid.layers.resize_bilinear(out_aux, input.shape[2:]) + + return out, out_aux + +def ocnet(input, num_classes): + logit = high_resolution_ocr_net(input, num_classes) + return logit + +if __name__ == '__main__': + image_shape = [-1, 3, 769, 769] + image = fluid.data(name='image', shape=image_shape, dtype='float32') + logit = ocnet(image, 4) + print("logit:", logit.shape)