├── README.md ├── backBone ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── densenet.cpython-37.pyc │ ├── mobilenet.cpython-37.pyc │ ├── resnet.cpython-37.pyc │ ├── vgg.cpython-37.pyc │ └── xception.cpython-37.pyc ├── densenet.py ├── mobilenet.py ├── resnet.py ├── vgg.py └── xception.py ├── backBone_pretrained_weights └── densenet121_weights_tf_dim_ordering_tf_kernels_notop.h5 ├── builders ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ └── model_builder.cpython-37.pyc └── model_builder.py ├── config.py ├── data ├── image_A │ └── testfile ├── train │ ├── image │ │ └── 1.tif │ └── label │ │ └── 1.png └── val │ ├── image │ └── 11.tif │ └── label │ └── 11.png ├── models ├── Custom_net.py ├── __init__.py ├── __pycache__ │ ├── Custom_net.cpython-37.pyc │ ├── __init__.cpython-37.pyc │ ├── bisegnet.cpython-37.pyc │ ├── deeplab_v3.cpython-37.pyc │ ├── deeplab_v3_plus.cpython-37.pyc │ ├── denseaspp.cpython-37.pyc │ ├── fcn.cpython-37.pyc │ ├── network.cpython-37.pyc │ ├── pan.cpython-37.pyc │ ├── pspnet.cpython-37.pyc │ ├── refinenet.cpython-37.pyc │ ├── segnet.cpython-37.pyc │ └── unet.cpython-37.pyc ├── bisegnet.py ├── deeplab_v3.py ├── deeplab_v3_plus.py ├── denseaspp.py ├── fcn.py ├── network.py ├── pan.py ├── pspnet.py ├── refinenet.py ├── segnet.py └── unet.py ├── predict.py ├── tools ├── Block_tricks.py ├── __pycache__ │ ├── Block_tricks.cpython-37.pyc │ ├── callbacks.cpython-37.pyc │ ├── dataloader.cpython-37.pyc │ ├── layers.cpython-37.pyc │ ├── learning_rate.cpython-37.pyc │ └── metrics.cpython-37.pyc ├── callbacks.py ├── dataloader.py ├── layers.py ├── learning_rate.py ├── metrics.py └── split_val_data_from_train.py ├── train.py ├── train_custom.py └── train_unet.py /README.md: -------------------------------------------------------------------------------- 1 | # Remote-sensing-image-semantic-segmentation-tf2 2 | The remote sensing image semantic segmentation repository based on tf.keras includes backbone networks such as resnet, densenet, mobilenet, and segmentation networks such as deeplabv3+, pspnet, panet, and segnet. 3 | 4 | This repository has been used to participate in the remote sensing semantic image segmentation track of the 2020 National Artificial Intelligence Competition (NAIC). 5 | 6 | ---- 7 | ## Data description 8 | 9 | |class|label| 10 | | :------: | :------: | 11 | |Water|100| 12 | |Transportation|200| 13 | |Building|300| 14 | |Arable land|400| 15 | |Grassland|500| 16 | |Woodland|600| 17 | |Bare soil|700| 18 | |Others|800| 19 | 20 | ## Requirements 21 | - python 3.7 22 | - tensorflow-gpu 2.3 23 | - opencv-python 24 | - tqdm 25 | - numpy 26 | - argparse 27 | - matplotlib 28 | - Pillow 29 | 30 | ## Usage 31 | ### 1. Download dataset 32 | > data 33 | 34 | ### 2. Separate the validation set from the training data (optional) 35 | > `python split_val_data_from_train.py` 36 | 37 | ### 3. Complete the basic configuration of training and testing, such as data path, model path, etc. 38 | > Modify the config.py file 39 | 40 | ### 4. Train 41 | > `python train.py --model DeepLabV3Plus --backBone ResNet152 --lr_scheduler cosine_decay --lr_warmup True` 42 | 43 | ### 4. Download pre-trained weights 44 | > [Link (TBD)](#) 45 | 46 | ### 5. Inference 47 | > `python predict.py` 48 | 49 | ## Results 50 | 51 | > MIou, FWIou 52 | > The inference result image is in the results folder 53 | -------------------------------------------------------------------------------- /backBone/__init__.py: -------------------------------------------------------------------------------- 1 | from backBone.densenet import DenseNet 2 | from backBone.resnet import ResNet 3 | from backBone.vgg import VGG 4 | from backBone.xception import Xception 5 | from backBone.mobilenet import MobileNet -------------------------------------------------------------------------------- /backBone/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/backBone/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /backBone/__pycache__/densenet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/backBone/__pycache__/densenet.cpython-37.pyc -------------------------------------------------------------------------------- /backBone/__pycache__/mobilenet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/backBone/__pycache__/mobilenet.cpython-37.pyc -------------------------------------------------------------------------------- /backBone/__pycache__/resnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/backBone/__pycache__/resnet.cpython-37.pyc -------------------------------------------------------------------------------- /backBone/__pycache__/vgg.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/backBone/__pycache__/vgg.cpython-37.pyc -------------------------------------------------------------------------------- /backBone/__pycache__/xception.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/backBone/__pycache__/xception.cpython-37.pyc -------------------------------------------------------------------------------- /backBone/densenet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | layers = tf.keras.layers 3 | backend = tf.keras.backend 4 | 5 | 6 | class DenseNet(object): 7 | def __init__(self, version='DenseNet121', dilation=None, **kwargs): 8 | """ 9 | The implementation of DenseNet based on Tensorflow. 10 | :param version: 'DenseNet121', 'DenseNet169', 'DenseNet201' or 'DenseNet264'. 11 | :param dilation: Whether to use dilation strategy. 12 | :param kwargs: other parameters. 13 | """ 14 | super(DenseNet, self).__init__(**kwargs) 15 | params = {'DenseNet121': [6, 12, 24, 16], 16 | 'DenseNet169': [6, 12, 32, 32], 17 | 'DenseNet201': [6, 12, 48, 32], 18 | 'DenseNet264': [6, 12, 64, 48]} 19 | self.version = version 20 | assert version in params 21 | self.params = params[version] 22 | 23 | if dilation is None: 24 | self.dilation = [1, 1] 25 | else: 26 | self.dilation = dilation 27 | assert len(self.dilation) == 2 28 | 29 | def _dense_block(self, x, blocks, name, dilation=1): 30 | """A dense block. 31 | 32 | # Arguments 33 | x: input tensor. 34 | blocks: integer, the number of building blocks. 35 | name: string, block label. 36 | 37 | # Returns 38 | output tensor for the block. 39 | """ 40 | for i in range(blocks): 41 | x = self._conv_block(x, 32, name=name + '_block' + str(i + 1), dilation=dilation) 42 | return x 43 | 44 | def _transition_block(self, x, reduction, name, dilation=1): 45 | """A transition block. 46 | 47 | # Arguments 48 | x: input tensor. 49 | reduction: float, compression rate at transition layers. 50 | name: string, block label. 51 | 52 | # Returns 53 | output tensor for the block. 54 | """ 55 | bn_axis = 3 if backend.image_data_format() == 'channels_last' else 1 56 | x = layers.BatchNormalization(axis=bn_axis, epsilon=1.001e-5, 57 | name=name + '_bn')(x) 58 | x = layers.Activation('relu', name=name + '_relu')(x) 59 | x = layers.Conv2D(int(backend.int_shape(x)[bn_axis] * reduction), 1, 60 | use_bias=False, 61 | name=name + '_conv', 62 | dilation_rate=dilation)(x) 63 | if dilation == 1: 64 | x = layers.AveragePooling2D(2, strides=2, name=name + '_pool')(x) 65 | return x 66 | 67 | def _conv_block(self, x, growth_rate, name, dilation=1): 68 | """A building block for a dense block. 69 | 70 | # Arguments 71 | x: input tensor. 72 | growth_rate: float, growth rate at dense layers. 73 | name: string, block label. 74 | 75 | # Returns 76 | Output tensor for the block. 77 | """ 78 | _, h, w, _ = backend.int_shape(x) 79 | 80 | bn_axis = 3 if backend.image_data_format() == 'channels_last' else 1 81 | x1 = layers.BatchNormalization(axis=bn_axis, 82 | epsilon=1.001e-5, 83 | name=name + '_0_bn')(x) 84 | x1 = layers.Activation('relu', name=name + '_0_relu')(x1) 85 | x1 = layers.Conv2D(4 * growth_rate, 1, 86 | use_bias=False, 87 | name=name + '_1_conv')(x1) 88 | x1 = layers.BatchNormalization(axis=bn_axis, epsilon=1.001e-5, 89 | name=name + '_1_bn')(x1) 90 | x1 = layers.Activation('relu', name=name + '_1_relu')(x1) 91 | x1 = layers.Conv2D(growth_rate, 3, 92 | padding='same', 93 | use_bias=False, 94 | name=name + '_2_conv', 95 | dilation_rate=dilation)(x1) 96 | x = layers.Concatenate( axis=bn_axis, name=name + '_concat')([x, x1]) 97 | return x 98 | 99 | def __call__(self, inputs, output_stages='c5', **kwargs): 100 | """ 101 | call for DenseNet. 102 | :param inputs: a 4-D tensor. 103 | :param output_stages: str or a list of str containing the output stages. 104 | :param kwargs: other parameters. 105 | :return: the output of different stages. 106 | """ 107 | _, h, w, _ = backend.int_shape(inputs) 108 | 109 | blocks = self.params 110 | dilation = self.dilation 111 | bn_axis = 3 if backend.image_data_format() == 'channels_last' else 1 112 | 113 | x = layers.ZeroPadding2D(padding=((3, 3), (3, 3)))(inputs) 114 | x = layers.Conv2D(64, 7, strides=2, use_bias=False, name='conv1/conv')(x) 115 | x = layers.BatchNormalization( 116 | axis=bn_axis, epsilon=1.001e-5, name='conv1/bn')(x) 117 | x = layers.Activation('relu', name='conv1/relu')(x) 118 | x = layers.ZeroPadding2D(padding=((1, 1), (1, 1)))(x) 119 | x = layers.MaxPooling2D(3, strides=2, name='pool1')(x) 120 | c1 = x # 64 121 | 122 | x = self._dense_block(x, blocks[0], name='conv2') 123 | x = self._transition_block(x, 0.5, name='pool2') 124 | c2 = x # 32 125 | 126 | x = self._dense_block(x, blocks[1], name='conv3') 127 | x = self._transition_block(x, 0.5, name='pool3', dilation=dilation[0]) 128 | c3 = x # 16 129 | 130 | x = self._dense_block(x, blocks[2], name='conv4', dilation=dilation[0]) 131 | x = self._transition_block(x, 0.5, name='pool4', dilation=dilation[1]) # dilation[1] ==2 132 | c4 = x # 16 133 | 134 | x = self._dense_block(x, blocks[3], name='conv5', dilation=dilation[1]) # dilation[1] ==2 135 | x = layers.BatchNormalization( 136 | axis=bn_axis, epsilon=1.001e-5, name='bn')(x) 137 | x = layers.Activation('relu', name='relu')(x) 138 | c5 = x # 16 139 | 140 | self.outputs = {'c1': c1, 141 | 'c2': c2, 142 | 'c3': c3, 143 | 'c4': c4, 144 | 'c5': c5} 145 | 146 | if type(output_stages) is not list: 147 | return self.outputs[output_stages] 148 | else: 149 | return [self.outputs[ci] for ci in output_stages] 150 | -------------------------------------------------------------------------------- /backBone/mobilenet.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | 4 | layers = tf.keras.layers 5 | backend = tf.keras.backend 6 | 7 | 8 | def _make_divisible(v, divisor, min_value=None): 9 | if min_value is None: 10 | min_value = divisor 11 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 12 | # Make sure that round down does not go down by more than 10%. 13 | if new_v < 0.9 * v: 14 | new_v += divisor 15 | return new_v 16 | 17 | 18 | def _correct_pad(backend, inputs, kernel_size): 19 | """Returns a tuple for zero-padding for 2D convolution with downsampling. 20 | 21 | # Arguments 22 | input_size: An integer or tuple/list of 2 integers. 23 | kernel_size: An integer or tuple/list of 2 integers. 24 | 25 | # Returns 26 | A tuple. 27 | """ 28 | img_dim = 2 if backend.image_data_format() == 'channels_first' else 1 29 | input_size = backend.int_shape(inputs)[img_dim:(img_dim + 2)] 30 | 31 | if isinstance(kernel_size, int): 32 | kernel_size = (kernel_size, kernel_size) 33 | 34 | if input_size[0] is None: 35 | adjust = (1, 1) 36 | else: 37 | adjust = (1 - input_size[0] % 2, 1 - input_size[1] % 2) 38 | 39 | correct = (kernel_size[0] // 2, kernel_size[1] // 2) 40 | 41 | return ((correct[0] - adjust[0], correct[0]), 42 | (correct[1] - adjust[1], correct[1])) 43 | 44 | 45 | class MobileNet(object): 46 | def __init__(self, version='MobileNetV2', dilation=None, **kwargs): 47 | """ 48 | The implementation of MobileNetV1 and MobileNetV2 based on Tensorflow. 49 | :param version: 'MobileNetV1' or 'MobileNetV2' 50 | :param dilation: Whether to use dialtion strategy 51 | :param kwargs: other parameters 52 | """ 53 | super(MobileNet, self).__init__(**kwargs) 54 | self.version = version 55 | self.mobilenet = {'MobileNetV1': self._mobilenet_v1, 56 | 'MobileNetV2': self._mobilenet_v2} 57 | assert version in self.mobilenet 58 | 59 | if dilation is None: 60 | self.dilation = [1, 1] 61 | else: 62 | self.dilation = dilation 63 | assert len(self.dilation) == 2 64 | 65 | def __call__(self, inputs, output_stages='c5', **kwargs): 66 | """ 67 | call for MobileNetV1 or MobileNetV2. 68 | :param inputs: a 4-D tensor 69 | :param output_stages: str or a list of str indicating the output stages. 70 | :param kwargs: other parameters 71 | :return: a 4-D tensor 72 | """ 73 | net = self.mobilenet[self.version] 74 | c1, c2, c3, c4, c5 = net(inputs) 75 | 76 | self.outputs = {'c1': c1, 77 | 'c2': c2, 78 | 'c3': c3, 79 | 'c4': c4, 80 | 'c5': c5} 81 | 82 | if type(output_stages) is not list: 83 | return self.outputs[output_stages] 84 | else: 85 | return [self.outputs[ci] for ci in output_stages] 86 | 87 | def _inverted_res_block_v2(self, inputs, expansion, stride, alpha, filters, block_id, dilation=1): 88 | """ 89 | inverted residual block in MobileNetV2. 90 | :param inputs: a 4-D tensor 91 | :param expansion: the expansion rate. 92 | :param stride: stride for convolution 93 | :param alpha: controls the width of the network. 94 | :param filters: output filters 95 | :param block_id: block id 96 | :param dilation: dilation rate 97 | :return: a 4-D tensor 98 | """ 99 | channel_axis = 1 if backend.image_data_format() == 'channels_first' else -1 100 | 101 | in_channels = backend.int_shape(inputs)[channel_axis] 102 | pointwise_conv_filters = int(filters * alpha) 103 | pointwise_filters = _make_divisible(pointwise_conv_filters, 8) 104 | x = inputs 105 | prefix = 'block_{}_'.format(block_id) 106 | 107 | if block_id: 108 | # Expand 109 | x = layers.Conv2D(expansion * in_channels, 110 | kernel_size=1, 111 | padding='same', 112 | use_bias=False, 113 | activation=None, 114 | name=prefix + 'expand')(x) 115 | x = layers.BatchNormalization(axis=channel_axis, 116 | epsilon=1e-3, 117 | momentum=0.999, 118 | name=prefix + 'expand_BN')(x) 119 | x = layers.ReLU(6., name=prefix + 'expand_relu')(x) 120 | else: 121 | prefix = 'expanded_conv_' 122 | 123 | # Depthwise 124 | if stride == 2 and dilation == 1: 125 | x = layers.ZeroPadding2D(padding=_correct_pad(backend, x, 3), 126 | name=prefix + 'pad')(x) 127 | x = layers.DepthwiseConv2D(kernel_size=3, 128 | strides=stride if dilation == 1 else 1, 129 | activation=None, 130 | use_bias=False, 131 | padding='valid' if stride == 2 and dilation == 1 else 'same', 132 | name=prefix + 'depthwise', 133 | dilation_rate=dilation)(x) 134 | x = layers.BatchNormalization(axis=channel_axis, 135 | epsilon=1e-3, 136 | momentum=0.999, 137 | name=prefix + 'depthwise_BN')(x) 138 | 139 | x = layers.ReLU(6., name=prefix + 'depthwise_relu')(x) 140 | 141 | # Project 142 | x = layers.Conv2D(pointwise_filters, 143 | kernel_size=1, 144 | padding='same', 145 | use_bias=False, 146 | activation=None, 147 | name=prefix + 'project')(x) 148 | x = layers.BatchNormalization(axis=channel_axis, 149 | epsilon=1e-3, 150 | momentum=0.999, 151 | name=prefix + 'project_BN')(x) 152 | 153 | if in_channels == pointwise_filters and stride == 1: 154 | return layers.Add(name=prefix + 'add')([inputs, x]) 155 | return x 156 | 157 | def _conv_block_v1(self, inputs, filters, alpha, kernel=(3, 3), strides=(1, 1)): 158 | """Adds an initial convolution layer (with batch normalization and relu6). 159 | 160 | # Arguments 161 | inputs: Input tensor of shape `(rows, cols, 3)` 162 | (with `channels_last` data format) or 163 | (3, rows, cols) (with `channels_first` data format). 164 | It should have exactly 3 inputs channels, 165 | and width and height should be no smaller than 32. 166 | E.g. `(224, 224, 3)` would be one valid value. 167 | filters: Integer, the dimensionality of the output space 168 | (i.e. the number of output filters in the convolution). 169 | alpha: controls the width of the network. 170 | - If `alpha` < 1.0, proportionally decreases the number 171 | of filters in each layer. 172 | - If `alpha` > 1.0, proportionally increases the number 173 | of filters in each layer. 174 | - If `alpha` = 1, default number of filters from the paper 175 | are used at each layer. 176 | kernel: An integer or tuple/list of 2 integers, specifying the 177 | width and height of the 2D convolution window. 178 | Can be a single integer to specify the same value for 179 | all spatial dimensions. 180 | strides: An integer or tuple/list of 2 integers, 181 | specifying the strides of the convolution 182 | along the width and height. 183 | Can be a single integer to specify the same value for 184 | all spatial dimensions. 185 | Specifying any stride value != 1 is incompatible with specifying 186 | any `dilation_rate` value != 1. 187 | 188 | # Input shape 189 | 4D tensor with shape: 190 | `(samples, channels, rows, cols)` if data_format='channels_first' 191 | or 4D tensor with shape: 192 | `(samples, rows, cols, channels)` if data_format='channels_last'. 193 | 194 | # Output shape 195 | 4D tensor with shape: 196 | `(samples, filters, new_rows, new_cols)` 197 | if data_format='channels_first' 198 | or 4D tensor with shape: 199 | `(samples, new_rows, new_cols, filters)` 200 | if data_format='channels_last'. 201 | `rows` and `cols` values might have changed due to stride. 202 | 203 | # Returns 204 | Output tensor of block. 205 | """ 206 | channel_axis = 1 if backend.image_data_format() == 'channels_first' else -1 207 | filters = int(filters * alpha) 208 | x = layers.ZeroPadding2D(padding=((0, 1), (0, 1)), name='conv1_pad')(inputs) 209 | x = layers.Conv2D(filters, kernel, 210 | padding='valid', 211 | use_bias=False, 212 | strides=strides, 213 | name='conv1')(x) 214 | x = layers.BatchNormalization(axis=channel_axis, name='conv1_bn')(x) 215 | return layers.ReLU(6., name='conv1_relu')(x) 216 | 217 | def _depthwise_conv_block_v1(self, inputs, pointwise_conv_filters, alpha, 218 | depth_multiplier=1, strides=(1, 1), block_id=1, dilation=1): 219 | """Adds a depthwise convolution block. 220 | 221 | A depthwise convolution block consists of a depthwise conv, 222 | batch normalization, relu6, pointwise convolution, 223 | batch normalization and relu6 activation. 224 | 225 | # Arguments 226 | inputs: Input tensor of shape `(rows, cols, channels)` 227 | (with `channels_last` data format) or 228 | (channels, rows, cols) (with `channels_first` data format). 229 | pointwise_conv_filters: Integer, the dimensionality of the output space 230 | (i.e. the number of output filters in the pointwise convolution). 231 | alpha: controls the width of the network. 232 | - If `alpha` < 1.0, proportionally decreases the number 233 | of filters in each layer. 234 | - If `alpha` > 1.0, proportionally increases the number 235 | of filters in each layer. 236 | - If `alpha` = 1, default number of filters from the paper 237 | are used at each layer. 238 | depth_multiplier: The number of depthwise convolution output channels 239 | for each input channel. 240 | The total number of depthwise convolution output 241 | channels will be equal to `filters_in * depth_multiplier`. 242 | strides: An integer or tuple/list of 2 integers, 243 | specifying the strides of the convolution 244 | along the width and height. 245 | Can be a single integer to specify the same value for 246 | all spatial dimensions. 247 | Specifying any stride value != 1 is incompatible with specifying 248 | any `dilation_rate` value != 1. 249 | block_id: Integer, a unique identification designating 250 | the block number. 251 | 252 | # Input shape 253 | 4D tensor with shape: 254 | `(batch, channels, rows, cols)` if data_format='channels_first' 255 | or 4D tensor with shape: 256 | `(batch, rows, cols, channels)` if data_format='channels_last'. 257 | 258 | # Output shape 259 | 4D tensor with shape: 260 | `(batch, filters, new_rows, new_cols)` 261 | if data_format='channels_first' 262 | or 4D tensor with shape: 263 | `(batch, new_rows, new_cols, filters)` 264 | if data_format='channels_last'. 265 | `rows` and `cols` values might have changed due to stride. 266 | 267 | # Returns 268 | Output tensor of block. 269 | """ 270 | channel_axis = 1 if backend.image_data_format() == 'channels_first' else -1 271 | pointwise_conv_filters = int(pointwise_conv_filters * alpha) 272 | 273 | strides = (1, 1) if dilation > 1 else strides 274 | 275 | if strides == (1, 1): 276 | x = inputs 277 | else: 278 | x = layers.ZeroPadding2D(((0, 1), (0, 1)), 279 | name='conv_pad_%d' % block_id)(inputs) 280 | x = layers.DepthwiseConv2D((3, 3), 281 | padding='same' if strides == (1, 1) else 'valid', 282 | depth_multiplier=depth_multiplier, 283 | strides=strides, 284 | use_bias=False, 285 | name='conv_dw_%d' % block_id, 286 | dilation_rate=dilation)(x) 287 | x = layers.BatchNormalization( 288 | axis=channel_axis, name='conv_dw_%d_bn' % block_id)(x) 289 | x = layers.ReLU(6., name='conv_dw_%d_relu' % block_id)(x) 290 | 291 | x = layers.Conv2D(pointwise_conv_filters, (1, 1), 292 | padding='same', 293 | use_bias=False, 294 | strides=(1, 1), 295 | name='conv_pw_%d' % block_id)(x) 296 | x = layers.BatchNormalization(axis=channel_axis, 297 | name='conv_pw_%d_bn' % block_id)(x) 298 | return layers.ReLU(6., name='conv_pw_%d_relu' % block_id)(x) 299 | 300 | def _mobilenet_v1(self, inputs, alpha=1.0, depth_multiplier=1): 301 | """ 302 | call for MobileNetV1. 303 | :param inputs: a 4-D tensor. 304 | :param alpha: controls the width of the network. 305 | :param depth_multiplier: depth multiplier for depthwise convolution. 306 | :return: . 307 | """ 308 | dilation = self.dilation 309 | 310 | x = self._conv_block_v1(inputs, 32, alpha, strides=(2, 2)) 311 | x = self._depthwise_conv_block_v1(x, 64, alpha, depth_multiplier, block_id=1) 312 | c1 = x 313 | 314 | x = self._depthwise_conv_block_v1(x, 128, alpha, depth_multiplier, 315 | strides=(2, 2), block_id=2) 316 | x = self._depthwise_conv_block_v1(x, 128, alpha, depth_multiplier, block_id=3) 317 | c2 = x 318 | 319 | x = self._depthwise_conv_block_v1(x, 256, alpha, depth_multiplier, 320 | strides=(2, 2), block_id=4) 321 | x = self._depthwise_conv_block_v1(x, 256, alpha, depth_multiplier, block_id=5) 322 | c3 = x 323 | 324 | x = self._depthwise_conv_block_v1(x, 512, alpha, depth_multiplier, 325 | strides=(2, 2), block_id=6, dilation=dilation[0]) 326 | x = self._depthwise_conv_block_v1(x, 512, alpha, depth_multiplier, block_id=7, dilation=dilation[0]) 327 | x = self._depthwise_conv_block_v1(x, 512, alpha, depth_multiplier, block_id=8, dilation=dilation[0]) 328 | x = self._depthwise_conv_block_v1(x, 512, alpha, depth_multiplier, block_id=9, dilation=dilation[0]) 329 | x = self._depthwise_conv_block_v1(x, 512, alpha, depth_multiplier, block_id=10, dilation=dilation[0]) 330 | x = self._depthwise_conv_block_v1(x, 512, alpha, depth_multiplier, block_id=11, dilation=dilation[0]) 331 | c4 = x 332 | 333 | x = self._depthwise_conv_block_v1(x, 1024, alpha, depth_multiplier, 334 | strides=(2, 2), block_id=12, dilation=dilation[1]) 335 | x = self._depthwise_conv_block_v1(x, 1024, alpha, depth_multiplier, block_id=13, dilation=dilation[1]) 336 | c5 = x 337 | 338 | return c1, c2, c3, c4, c5 339 | 340 | def _mobilenet_v2(self, inputs, alpha=1.0): 341 | """ 342 | call for MobileNetV2. 343 | :param inputs: a 4-D tensor. 344 | :param alpha: controls the width of the network. 345 | :return: the output of different stages. 346 | """ 347 | dilation = self.dilation 348 | channel_axis = 1 if backend.image_data_format() == 'channels_first' else -1 349 | 350 | first_block_filters = _make_divisible(32 * alpha, 8) 351 | x = layers.ZeroPadding2D(padding=_correct_pad(backend, inputs, 3), 352 | name='Conv1_pad')(inputs) 353 | x = layers.Conv2D(first_block_filters, 354 | kernel_size=3, 355 | strides=(2, 2), 356 | padding='valid', 357 | use_bias=False, 358 | name='Conv1')(x) 359 | x = layers.BatchNormalization(axis=channel_axis, 360 | epsilon=1e-3, 361 | momentum=0.999, 362 | name='bn_Conv1')(x) 363 | x = layers.ReLU(6., name='Conv1_relu')(x) 364 | 365 | x = self._inverted_res_block_v2(x, filters=16, alpha=alpha, stride=1, 366 | expansion=1, block_id=0) 367 | c1 = x 368 | 369 | x = self._inverted_res_block_v2(x, filters=24, alpha=alpha, stride=2, 370 | expansion=6, block_id=1) 371 | x = self._inverted_res_block_v2(x, filters=24, alpha=alpha, stride=1, 372 | expansion=6, block_id=2) 373 | c2 = x 374 | 375 | x = self._inverted_res_block_v2(x, filters=32, alpha=alpha, stride=2, 376 | expansion=6, block_id=3) 377 | x = self._inverted_res_block_v2(x, filters=32, alpha=alpha, stride=1, 378 | expansion=6, block_id=4) 379 | x = self._inverted_res_block_v2(x, filters=32, alpha=alpha, stride=1, 380 | expansion=6, block_id=5) 381 | c3 = x 382 | 383 | x = self._inverted_res_block_v2(x, filters=64, alpha=alpha, stride=2, 384 | expansion=6, block_id=6, dilation=dilation[0]) 385 | x = self._inverted_res_block_v2(x, filters=64, alpha=alpha, stride=1, 386 | expansion=6, block_id=7, dilation=dilation[0]) 387 | x = self._inverted_res_block_v2(x, filters=64, alpha=alpha, stride=1, 388 | expansion=6, block_id=8, dilation=dilation[0]) 389 | x = self._inverted_res_block_v2(x, filters=64, alpha=alpha, stride=1, 390 | expansion=6, block_id=9, dilation=dilation[0]) 391 | 392 | x = self._inverted_res_block_v2(x, filters=96, alpha=alpha, stride=1, 393 | expansion=6, block_id=10, dilation=dilation[0]) 394 | x = self._inverted_res_block_v2(x, filters=96, alpha=alpha, stride=1, 395 | expansion=6, block_id=11, dilation=dilation[0]) 396 | x = self._inverted_res_block_v2(x, filters=96, alpha=alpha, stride=1, 397 | expansion=6, block_id=12, dilation=dilation[0]) 398 | c4 = x 399 | 400 | x = self._inverted_res_block_v2(x, filters=160, alpha=alpha, stride=2, 401 | expansion=6, block_id=13, dilation=dilation[1]) 402 | x = self._inverted_res_block_v2(x, filters=160, alpha=alpha, stride=1, 403 | expansion=6, block_id=14, dilation=dilation[1]) 404 | x = self._inverted_res_block_v2(x, filters=160, alpha=alpha, stride=1, 405 | expansion=6, block_id=15, dilation=dilation[1]) 406 | 407 | x = self._inverted_res_block_v2(x, filters=320, alpha=alpha, stride=1, 408 | expansion=6, block_id=16, dilation=dilation[1]) 409 | 410 | # no alpha applied to last conv as stated in the paper: 411 | # if the width multiplier is greater than 1 we 412 | # increase the number of output channels 413 | if alpha > 1.0: 414 | last_block_filters = _make_divisible(1280 * alpha, 8) 415 | else: 416 | last_block_filters = 1280 417 | 418 | x = layers.Conv2D(last_block_filters, 419 | kernel_size=1, 420 | use_bias=False, 421 | name='Conv_1')(x) 422 | x = layers.BatchNormalization(axis=channel_axis, 423 | epsilon=1e-3, 424 | momentum=0.999, 425 | name='Conv_1_bn')(x) 426 | x = layers.ReLU(6., name='out_relu')(x) 427 | c5 = x 428 | 429 | return c1, c2, c3, c4, c5 430 | -------------------------------------------------------------------------------- /backBone/resnet.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | 4 | layers = tf.keras.layers 5 | backend = tf.keras.backend 6 | 7 | 8 | class ResNet(object): 9 | def __init__(self, version='ResNet50', dilation=None, **kwargs): 10 | """ 11 | The implementation of ResNet based on Tensorflow. 12 | :param version: 'ResNet50', 'ResNet101' or 'ResNet152' 13 | :param dilation: Whether to use dilation strategy 14 | :param kwargs: other parameters. 15 | """ 16 | super(ResNet, self).__init__(**kwargs) 17 | params = {'ResNet50': [2, 3, 5, 2], 18 | 'ResNet101': [2, 3, 22, 2], 19 | 'ResNet152': [2, 7, 35, 2]} 20 | self.version = version 21 | assert version in params 22 | self.params = params[version] 23 | 24 | if dilation is None: 25 | self.dilation = [1, 1] 26 | else: 27 | self.dilation = dilation 28 | assert len(self.dilation) == 2 29 | 30 | def _identity_block(self, input_tensor, kernel_size, filters, stage, block, dilation=1): 31 | """The identity block is the block that has no conv layer at shortcut. 32 | 33 | # Arguments 34 | input_tensor: input tensor 35 | kernel_size: default 3, the kernel size of 36 | middle conv layer at main path 37 | filters: list of integers, the filters of 3 conv layer at main path 38 | stage: integer, current stage label, used for generating layer names 39 | block: 'a','b'..., current block label, used for generating layer names 40 | 41 | # Returns 42 | Output tensor for the block. 43 | """ 44 | filters1, filters2, filters3 = filters 45 | if backend.image_data_format() == 'channels_last': 46 | bn_axis = 3 47 | else: 48 | bn_axis = 1 49 | 50 | if block > 'z': 51 | block = chr(ord(block) - ord('z') + ord('A') - 1) 52 | 53 | conv_name_base = 'res' + str(stage) + block + '_branch' 54 | bn_name_base = 'bn' + str(stage) + block + '_branch' 55 | 56 | x = layers.Conv2D(filters1, (1, 1), 57 | kernel_initializer='he_normal', 58 | name=conv_name_base + '2a')(input_tensor) 59 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x) 60 | x = layers.Activation('relu')(x) 61 | 62 | x = layers.Conv2D(filters2, kernel_size, 63 | padding='same', 64 | kernel_initializer='he_normal', 65 | name=conv_name_base + '2b', 66 | dilation_rate=dilation)(x) 67 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x) 68 | x = layers.Activation('relu')(x) 69 | 70 | x = layers.Conv2D(filters3, (1, 1), 71 | kernel_initializer='he_normal', 72 | name=conv_name_base + '2c')(x) 73 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x) 74 | 75 | x = layers.add([x, input_tensor]) 76 | x = layers.Activation('relu')(x) 77 | return x 78 | 79 | def _conv_block(self, 80 | input_tensor, 81 | kernel_size, 82 | filters, 83 | stage, 84 | block, 85 | strides=(2, 2), 86 | dilation=1): 87 | """A block that has a conv layer at shortcut. 88 | 89 | # Arguments 90 | input_tensor: input tensor 91 | kernel_size: default 3, the kernel size of 92 | middle conv layer at main path 93 | filters: list of integers, the filters of 3 conv layer at main path 94 | stage: integer, current stage label, used for generating layer names 95 | block: 'a','b'..., current block label, used for generating layer names 96 | strides: Strides for the first conv layer in the block. 97 | 98 | # Returns 99 | Output tensor for the block. 100 | 101 | Note that from stage 3, 102 | the first conv layer at main path is with strides=(2, 2) 103 | And the shortcut should have strides=(2, 2) as well 104 | """ 105 | filters1, filters2, filters3 = filters 106 | if backend.image_data_format() == 'channels_last': 107 | bn_axis = 3 108 | else: 109 | bn_axis = 1 110 | conv_name_base = 'res' + str(stage) + block + '_branch' 111 | bn_name_base = 'bn' + str(stage) + block + '_branch' 112 | 113 | strides = (1, 1) if dilation > 1 else strides 114 | 115 | x = layers.Conv2D(filters1, (1, 1), strides=strides, 116 | kernel_initializer='he_normal', 117 | name=conv_name_base + '2a')(input_tensor) 118 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x) 119 | x = layers.Activation('relu')(x) 120 | 121 | x = layers.Conv2D(filters2, kernel_size, padding='same', 122 | kernel_initializer='he_normal', 123 | name=conv_name_base + '2b', 124 | dilation_rate=dilation)(x) 125 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x) 126 | x = layers.Activation('relu')(x) 127 | 128 | x = layers.Conv2D(filters3, (1, 1), 129 | kernel_initializer='he_normal', 130 | name=conv_name_base + '2c')(x) 131 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x) 132 | 133 | shortcut = layers.Conv2D(filters3, (1, 1), strides=strides, 134 | kernel_initializer='he_normal', 135 | name=conv_name_base + '1')(input_tensor) 136 | shortcut = layers.BatchNormalization( 137 | axis=bn_axis, name=bn_name_base + '1')(shortcut) 138 | 139 | x = layers.add([x, shortcut]) 140 | x = layers.Activation('relu')(x) 141 | return x 142 | 143 | def __call__(self, inputs, output_stages='c5', **kwargs): 144 | """ 145 | call for ResNet50, ResNet101 or ResNet152. 146 | :param inputs: a 4-D tensor. 147 | :param output_stages: str or a list of str containing the output stages. 148 | :param kwargs: other parameters. 149 | :return: the output of different stages. 150 | """ 151 | if backend.image_data_format() == 'channels_last': 152 | bn_axis = 3 153 | else: 154 | bn_axis = 1 155 | 156 | dilation = self.dilation 157 | 158 | x = layers.ZeroPadding2D(padding=(3, 3), name='conv1_pad')(inputs) 159 | x = layers.Conv2D(64, (7, 7), 160 | strides=(2, 2), 161 | padding='valid', 162 | kernel_initializer='he_normal', 163 | name='conv1')(x) 164 | x = layers.BatchNormalization(axis=bn_axis, name='bn_conv1')(x) 165 | x = layers.Activation('relu')(x) 166 | x = layers.ZeroPadding2D(padding=(1, 1), name='pool1_pad')(x) 167 | x = layers.MaxPooling2D((3, 3), strides=(2, 2))(x) 168 | c1 = x 169 | 170 | x = self._conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) 171 | for i in range(self.params[0]): 172 | x = self._identity_block(x, 3, [64, 64, 256], stage=2, block=chr(ord('b') + i)) 173 | c2 = x 174 | 175 | x = self._conv_block(x, 3, [128, 128, 512], stage=3, block='a') 176 | for i in range(self.params[1]): 177 | x = self._identity_block(x, 3, [128, 128, 512], stage=3, block=chr(ord('b') + i)) 178 | c3 = x 179 | 180 | x = self._conv_block(x, 3, [256, 256, 1024], stage=4, block='a', dilation=dilation[0]) 181 | for i in range(self.params[2]): 182 | x = self._identity_block(x, 3, [256, 256, 1024], stage=4, block=chr(ord('b') + i), dilation=dilation[0]) 183 | c4 = x 184 | 185 | x = self._conv_block(x, 3, [512, 512, 2048], stage=5, block='a', dilation=dilation[1]) 186 | for i in range(self.params[3]): 187 | x = self._identity_block(x, 3, [512, 512, 2048], stage=5, block=chr(ord('b') + i), dilation=dilation[1]) 188 | c5 = x 189 | 190 | self.outputs = {'c1': c1, 191 | 'c2': c2, 192 | 'c3': c3, 193 | 'c4': c4, 194 | 'c5': c5} 195 | 196 | if type(output_stages) is not list: 197 | return self.outputs[output_stages] 198 | else: 199 | return [self.outputs[ci] for ci in output_stages] 200 | -------------------------------------------------------------------------------- /backBone/vgg.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | 4 | layers = tf.keras.layers 5 | backend = tf.keras.backend 6 | 7 | 8 | class VGG(object): 9 | def __init__(self, version='VGG16', dilation=None, **kwargs): 10 | """ 11 | The implementation of VGG16 and VGG19 based on Tensorflow. 12 | :param version: 'VGG16' or 'VGG19' 13 | :param dilation: Whether to use dilation strategy 14 | :param kwargs: other parameters. 15 | """ 16 | super(VGG, self).__init__(**kwargs) 17 | params = {'VGG16': [2, 2, 3, 3, 3], 18 | 'VGG19': [2, 2, 4, 4, 4]} 19 | self.version = version 20 | assert version in params 21 | self.params = params[version] 22 | 23 | if dilation is None: 24 | self.dilation = [1, 1] 25 | else: 26 | self.dilation = dilation 27 | assert len(self.dilation) == 2 28 | 29 | def __call__(self, inputs, output_stages='c5', **kwargs): 30 | """ 31 | call for VGG16 or VGG19. 32 | :param inputs: a 4-D tensor. 33 | :param output_stages: str or a list of str containing the output stages. 34 | :param kwargs: other parameters. 35 | :return: the output of different stages. 36 | """ 37 | dilation = self.dilation 38 | _, h, w, _ = backend.int_shape(inputs) 39 | 40 | # Block 1 41 | for i in range(self.params[0]): 42 | x = layers.Conv2D(64, (3, 3), 43 | activation='relu', 44 | padding='same', 45 | name='block1_conv' + str(i + 1))(inputs) 46 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block1_pool')(x) 47 | c1 = x 48 | 49 | # Block 2 50 | for i in range(self.params[1]): 51 | x = layers.Conv2D(128, (3, 3), 52 | activation='relu', 53 | padding='same', 54 | name='block2_conv' + str(i + 1))(x) 55 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block2_pool')(x) 56 | c2 = x 57 | 58 | # Block 3 59 | for i in range(self.params[2]): 60 | x = layers.Conv2D(256, (3, 3), 61 | activation='relu', 62 | padding='same', 63 | name='block3_conv' + str(i + 1))(x) 64 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block3_pool')(x) 65 | c3 = x 66 | 67 | # Block 4 68 | for i in range(self.params[3]): 69 | x = layers.Conv2D(512, (3, 3), 70 | activation='relu', 71 | padding='same', 72 | name='block4_conv' + str(i + 1), 73 | dilation_rate=dilation[0])(x) 74 | if dilation[0] == 1: 75 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool')(x) 76 | c4 = x 77 | 78 | # Block 5 79 | for i in range(self.params[4]): 80 | x = layers.Conv2D(512, (3, 3), 81 | activation='relu', 82 | padding='same', 83 | name='block5_conv' + str(i + 1), 84 | dilation_rate=dilation[1])(x) 85 | if dilation[1] == 1: 86 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block5_pool')(x) 87 | c5 = x 88 | 89 | self.outputs = {'c1': c1, 90 | 'c2': c2, 91 | 'c3': c3, 92 | 'c4': c4, 93 | 'c5': c5} 94 | 95 | if type(output_stages) is not list: 96 | return self.outputs[output_stages] 97 | else: 98 | return [self.outputs[ci] for ci in output_stages] 99 | -------------------------------------------------------------------------------- /backBone/xception.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | 4 | layers = tf.keras.layers 5 | backend = tf.keras.backend 6 | 7 | 8 | class Xception(object): 9 | def __init__(self, version='Xception', dilation=None, **kwargs): 10 | """ 11 | The implementation of Xception and Xception in DeepLabV3Plus based on Tensorflow. 12 | :param version: 'Xception' or 'Xception-DeepLab' 13 | :param dilation: Whether to use dilation strategy 14 | :param kwargs: other parameters. 15 | """ 16 | super(Xception, self).__init__(**kwargs) 17 | self.version = version 18 | if dilation is None: 19 | self.strides = [2, 2] 20 | else: 21 | self.strides = [2 if dilation[0] == 1 else 1] + [2 if dilation[1] == 1 else 1] 22 | assert len(self.strides) == 2 23 | assert version in ['Xception', 'Xception-DeepLab'] 24 | 25 | def __call__(self, inputs, output_stages='c5', **kwargs): 26 | """ 27 | call for Xception or Xception-DeepLab. 28 | :param inputs: a 4-D tensor. 29 | :param output_stages: str or a list of str containing the output stages. 30 | :param kwargs: other parameters. 31 | :return: the output of different stages. 32 | """ 33 | strides = self.strides 34 | if self.version == 'Xception-DeepLab': 35 | rm_pool = True 36 | num_middle_flow = 16 37 | else: 38 | rm_pool = False 39 | num_middle_flow = 8 40 | 41 | channel_axis = 1 if backend.image_data_format() == 'channels_first' else -1 42 | 43 | x = layers.Conv2D(32, (3, 3), 44 | strides=(2, 2), 45 | use_bias=False, 46 | padding='same', 47 | name='block1_conv1')(inputs) 48 | x = layers.BatchNormalization(axis=channel_axis, name='block1_conv1_bn')(x) 49 | x = layers.Activation('relu', name='block1_conv1_act')(x) 50 | x = layers.Conv2D(64, (3, 3), use_bias=False, padding='same', name='block1_conv2')(x) 51 | x = layers.BatchNormalization(axis=channel_axis, name='block1_conv2_bn')(x) 52 | x = layers.Activation('relu', name='block1_conv2_act')(x) 53 | 54 | residual = layers.Conv2D(128, (1, 1), 55 | strides=(2, 2), 56 | padding='same', 57 | use_bias=False)(x) 58 | residual = layers.BatchNormalization(axis=channel_axis)(residual) 59 | 60 | x = layers.SeparableConv2D(128, (3, 3), 61 | padding='same', 62 | use_bias=False, 63 | name='block2_sepconv1')(x) 64 | x = layers.BatchNormalization(axis=channel_axis, name='block2_sepconv1_bn')(x) 65 | x = layers.Activation('relu', name='block2_sepconv2_act')(x) 66 | x = layers.SeparableConv2D(128, (3, 3), 67 | padding='same', 68 | use_bias=False, 69 | name='block2_sepconv2')(x) 70 | x = layers.BatchNormalization(axis=channel_axis, name='block2_sepconv2_bn')(x) 71 | 72 | x = layers.MaxPooling2D((3, 3), 73 | strides=(2, 2), 74 | padding='same', 75 | name='block2_pool')(x) 76 | x = layers.add([x, residual]) 77 | c1 = x 78 | 79 | residual = layers.Conv2D(256, (1, 1), strides=(2, 2), 80 | padding='same', use_bias=False)(x) 81 | residual = layers.BatchNormalization(axis=channel_axis)(residual) 82 | 83 | x = layers.Activation('relu', name='block3_sepconv1_act')(x) 84 | x = layers.SeparableConv2D(256, (3, 3), 85 | padding='same', 86 | use_bias=False, 87 | name='block3_sepconv1')(x) 88 | x = layers.BatchNormalization(axis=channel_axis, name='block3_sepconv1_bn')(x) 89 | x = layers.Activation('relu', name='block3_sepconv2_act')(x) 90 | x = layers.SeparableConv2D(256, (3, 3), 91 | padding='same', 92 | use_bias=False, 93 | name='block3_sepconv2')(x) 94 | x = layers.BatchNormalization(axis=channel_axis, name='block3_sepconv2_bn')(x) 95 | 96 | if rm_pool: 97 | x = layers.Activation('relu', name='block3_sepconv3_act')(x) 98 | x = layers.SeparableConv2D(256, (3, 3), 99 | strides=(2, 2), 100 | padding='same', 101 | use_bias=False, 102 | name='block3_sepconv3')(x) 103 | x = layers.BatchNormalization(axis=channel_axis, name='block3_sepconv3_bn')(x) 104 | else: 105 | x = layers.MaxPooling2D((3, 3), strides=(2, 2), 106 | padding='same', 107 | name='block3_pool')(x) 108 | x = layers.add([x, residual]) 109 | c2 = x 110 | 111 | residual = layers.Conv2D(728, (1, 1), 112 | strides=strides[0], 113 | padding='same', 114 | use_bias=False)(x) 115 | residual = layers.BatchNormalization(axis=channel_axis)(residual) 116 | 117 | x = layers.Activation('relu', name='block4_sepconv1_act')(x) 118 | x = layers.SeparableConv2D(728, (3, 3), 119 | padding='same', 120 | use_bias=False, 121 | name='block4_sepconv1')(x) 122 | x = layers.BatchNormalization(axis=channel_axis, name='block4_sepconv1_bn')(x) 123 | x = layers.Activation('relu', name='block4_sepconv2_act')(x) 124 | x = layers.SeparableConv2D(728, (3, 3), 125 | padding='same', 126 | use_bias=False, 127 | name='block4_sepconv2')(x) 128 | x = layers.BatchNormalization(axis=channel_axis, name='block4_sepconv2_bn')(x) 129 | 130 | if rm_pool: 131 | x = layers.Activation('relu', name='block4_sepconv3_act')(x) 132 | x = layers.SeparableConv2D(728, (3, 3), 133 | strides=strides[0], 134 | padding='same', 135 | use_bias=False, 136 | name='block4_sepconv3')(x) 137 | x = layers.BatchNormalization(axis=channel_axis, name='block4_sepconv3_bn')(x) 138 | else: 139 | x = layers.MaxPooling2D((3, 3), strides=(2, 2), 140 | padding='same', 141 | name='block4_pool')(x) 142 | x = layers.add([x, residual]) 143 | c3 = x 144 | 145 | for i in range(num_middle_flow): 146 | residual = x 147 | prefix = 'block' + str(i + 5) 148 | 149 | x = layers.Activation('relu', name=prefix + '_sepconv1_act')(x) 150 | x = layers.SeparableConv2D(728, (3, 3), 151 | padding='same', 152 | use_bias=False, 153 | name=prefix + '_sepconv1')(x) 154 | x = layers.BatchNormalization(axis=channel_axis, 155 | name=prefix + '_sepconv1_bn')(x) 156 | x = layers.Activation('relu', name=prefix + '_sepconv2_act')(x) 157 | x = layers.SeparableConv2D(728, (3, 3), 158 | padding='same', 159 | use_bias=False, 160 | name=prefix + '_sepconv2')(x) 161 | x = layers.BatchNormalization(axis=channel_axis, 162 | name=prefix + '_sepconv2_bn')(x) 163 | x = layers.Activation('relu', name=prefix + '_sepconv3_act')(x) 164 | x = layers.SeparableConv2D(728, (3, 3), 165 | padding='same', 166 | use_bias=False, 167 | name=prefix + '_sepconv3')(x) 168 | x = layers.BatchNormalization(axis=channel_axis, 169 | name=prefix + '_sepconv3_bn')(x) 170 | 171 | x = layers.add([x, residual]) 172 | c4 = x 173 | 174 | residual = layers.Conv2D(1024, (1, 1), strides=strides[1], 175 | padding='same', use_bias=False)(x) 176 | residual = layers.BatchNormalization(axis=channel_axis)(residual) 177 | 178 | id = 5 + num_middle_flow 179 | x = layers.Activation('relu', name='block{id}_sepconv1_act'.format(id=id))(x) 180 | x = layers.SeparableConv2D(728, (3, 3), 181 | padding='same', 182 | use_bias=False, 183 | name='block{id}_sepconv1'.format(id=id))(x) 184 | x = layers.BatchNormalization(axis=channel_axis, name='block{id}_sepconv1_bn'.format(id=id))(x) 185 | x = layers.Activation('relu', name='block{id}_sepconv2_act'.format(id=id))(x) 186 | x = layers.SeparableConv2D(1024, (3, 3), 187 | padding='same', 188 | use_bias=False, 189 | name='block{id}_sepconv2'.format(id=id))(x) 190 | x = layers.BatchNormalization(axis=channel_axis, name='block{id}_sepconv2_bn'.format(id=id))(x) 191 | 192 | if rm_pool: 193 | x = layers.Activation('relu', name='block{id}_sepconv3_act'.format(id=id))(x) 194 | x = layers.SeparableConv2D(1024, (3, 3), 195 | strides=strides[1], 196 | padding='same', 197 | use_bias=False, 198 | name='block{id}_sepconv3'.format(id=id))(x) 199 | x = layers.BatchNormalization(axis=channel_axis, name='block{id}_sepconv3_bn'.format(id=id))(x) 200 | else: 201 | x = layers.MaxPooling2D((3, 3), 202 | strides=(2, 2), 203 | padding='same', 204 | name='block{id}_pool'.format(id=id))(x) 205 | x = layers.add([x, residual]) 206 | 207 | x = layers.SeparableConv2D(1536, (3, 3), 208 | padding='same', 209 | use_bias=False, 210 | name='block{id}_sepconv1'.format(id=id + 1))(x) 211 | x = layers.BatchNormalization(axis=channel_axis, name='block{id}_sepconv1_bn'.format(id=id + 1))(x) 212 | x = layers.Activation('relu', name='block{id}_sepconv1_act'.format(id=id + 1))(x) 213 | 214 | if self.version == 'Xception-DeepLab': 215 | x = layers.SeparableConv2D(1536, (3, 3), 216 | padding='same', 217 | use_bias=False, 218 | name='block{id}_sepconv1_1'.format(id=id + 1))(x) 219 | x = layers.BatchNormalization(axis=channel_axis, name='block{id}_sepconv1_1_bn'.format(id=id + 1))(x) 220 | x = layers.Activation('relu', name='block{id}_sepconv1_1_act'.format(id=id + 1))(x) 221 | 222 | x = layers.SeparableConv2D(2048, (3, 3), 223 | padding='same', 224 | use_bias=False, 225 | name='block{id}_sepconv2'.format(id=id + 1))(x) 226 | x = layers.BatchNormalization(axis=channel_axis, name='block{id}_sepconv2_bn'.format(id=id + 1))(x) 227 | x = layers.Activation('relu', name='block{id}_sepconv2_act'.format(id=id + 1))(x) 228 | 229 | c5 = x 230 | 231 | self.outputs = {'c1': c1, 232 | 'c2': c2, 233 | 'c3': c3, 234 | 'c4': c4, 235 | 'c5': c5} 236 | 237 | if type(output_stages) is not list: 238 | return self.outputs[output_stages] 239 | else: 240 | return [self.outputs[ci] for ci in output_stages] 241 | -------------------------------------------------------------------------------- /backBone_pretrained_weights/densenet121_weights_tf_dim_ordering_tf_kernels_notop.h5: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/backBone_pretrained_weights/densenet121_weights_tf_dim_ordering_tf_kernels_notop.h5 -------------------------------------------------------------------------------- /builders/__init__.py: -------------------------------------------------------------------------------- 1 | from builders.model_builder import builder -------------------------------------------------------------------------------- /builders/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/builders/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /builders/__pycache__/model_builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/builders/__pycache__/model_builder.cpython-37.pyc -------------------------------------------------------------------------------- /builders/model_builder.py: -------------------------------------------------------------------------------- 1 | 2 | from models import * 3 | import tensorflow as tf 4 | 5 | layers = tf.keras.layers 6 | 7 | 8 | def builder(num_classes, input_size=(256, 256), model='SegNet', base_model=None): 9 | models = {'FCN-8s': FCN, 10 | 'FCN-16s': FCN, 11 | 'FCN-32s': FCN, 12 | 'SegNet': SegNet, 13 | 'Bayesian-SegNet': SegNet, 14 | 'PAN': PAN, 15 | 'PSPNet': PSPNet, 16 | 'RefineNet': RefineNet, 17 | 'DenseASPP': DenseASPP, 18 | 'DeepLabV3': DeepLabV3, 19 | 'DeepLabV3Plus': DeepLabV3Plus, 20 | 'BiSegNet': BiSegNet} 21 | 22 | assert model in models 23 | 24 | net = models[model](num_classes, model, base_model) 25 | 26 | inputs = layers.Input(shape=input_size+(3,)) 27 | 28 | return net(inputs), net.get_base_model() 29 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | class Config(): 4 | 5 | def __init__(self, flag): 6 | self.n_classes = 8 7 | self.get_attr(flag) 8 | 9 | def get_attr(self,flag): 10 | 11 | if flag.lower() == 'train': 12 | self.epoch = 150 13 | self.batch_size = 16 14 | self.lr = 0.008 15 | self.train_data_path = 'data/train/image' 16 | self.train_label_path = 'data/train/label' 17 | self.val_data_path = 'data/val/image' 18 | self.val_label_path = 'data/val/label' 19 | self.data_augment = False 20 | self.weight_path = 'weights/custom' 21 | self.train_number = self.get_train_number() 22 | self.val_number = self.get_val_number() 23 | self.steps_per_epoch = self.train_number//self.batch_size if self.train_number % self.batch_size ==0 else self.train_number//self.batch_size +1 24 | self.validation_steps= self.val_number//self.batch_size if self.val_number % self.batch_size ==0 else self.val_number// self.batch_size +1 25 | 26 | if flag.lower() == 'test': 27 | self.batch_size = 32 28 | self.data_path = 'data/image_A' 29 | # self.weight_path = 'weights/DenseNet121_DeepLabV3Plus---/weights-219-0.4049-0.7466.h5' 30 | self.weight_path = 'weights/custom-2/weights-130-0.5092-0.6668.h5' 31 | self.output_path= 'results' 32 | self.image_number = self.get_test_number() 33 | self.steps = self.image_number//self.batch_size if self.image_number % self.batch_size ==0 else self.image_number//self.batch_size +1 34 | 35 | def check_folder(self,dir): 36 | if not os.path.exists(dir): 37 | os.makedirs(dir) 38 | 39 | def get_train_number(self): 40 | res = 0 41 | for dir_entry in os.listdir(self.train_data_path): 42 | if os.path.isfile(os.path.join(self.train_data_path, dir_entry)): 43 | res += 1 44 | return res 45 | 46 | def get_val_number(self): 47 | res = 0 48 | for dir_entry in os.listdir(self.val_data_path): 49 | if os.path.isfile(os.path.join(self.val_data_path, dir_entry)): 50 | res += 1 51 | return res 52 | 53 | def get_test_number(self): 54 | res = 0 55 | for dir_entry in os.listdir(self.data_path): 56 | if os.path.isfile(os.path.join(self.data_path, dir_entry)): 57 | res += 1 58 | return res 59 | 60 | 61 | 62 | 63 | 64 | 65 | -------------------------------------------------------------------------------- /data/image_A/testfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/data/image_A/testfile -------------------------------------------------------------------------------- /data/train/image/1.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/data/train/image/1.tif -------------------------------------------------------------------------------- /data/train/label/1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/data/train/label/1.png -------------------------------------------------------------------------------- /data/val/image/11.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/data/val/image/11.tif -------------------------------------------------------------------------------- /data/val/label/11.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/data/val/label/11.png -------------------------------------------------------------------------------- /models/Custom_net.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | layers = tf.keras.layers 3 | backend = tf.keras.backend 4 | from tensorflow.keras.regularizers import l2 5 | from tools import layers as custom_layers 6 | from tools.Block_tricks import apRelu, global_context_block 7 | 8 | class DenseNet(object): 9 | def __init__(self, params, growth_rate, reduction): 10 | assert len(params) == 4 11 | self.params = params 12 | self.reduction = reduction 13 | self.growth_rate =growth_rate 14 | 15 | def _conv_bn_relu(self, x, filters, kernel_size, strides=1): 16 | x = layers.Conv2D(filters, kernel_size, strides=strides, use_bias=False, kernel_initializer='he_normal', padding='same')(x) 17 | x = layers.BatchNormalization()(x) 18 | x = apRelu(x) 19 | return x 20 | 21 | def _bottleneck_block(self, x, growth_rate, dilation, name): 22 | 23 | x1 = layers.BatchNormalization(name=name + '_1_bn')(x) 24 | x1 = layers.Activation('swish')(x1) 25 | x1 = layers.Conv2D(2 * growth_rate, 1, use_bias=False, kernel_initializer='he_normal', name=name + '_1_conv')(x1) 26 | x1 = layers.BatchNormalization(name=name + '_2_bn')(x1) 27 | x1 = apRelu(x1) 28 | x1 = layers.Conv2D(growth_rate, 3, padding='same', use_bias=False, kernel_initializer='he_normal', dilation_rate=dilation, name=name + '_2_conv')(x1) 29 | x = layers.Concatenate(axis=-1, name=name + '_concat')([x, x1]) 30 | return x 31 | 32 | def _dense_block(self, x, growth_rate, n_blocks, name, dilation=1): 33 | 34 | for i in range(n_blocks): 35 | x = self._bottleneck_block(x, growth_rate, dilation, name=name + '_block_' + str(i + 1)) 36 | return x 37 | 38 | def _transition_block(self, x, reduction, name, dilation=1): 39 | 40 | x = layers.BatchNormalization(name=name + '_bn')(x) 41 | x = layers.Activation('swish')(x) 42 | x = layers.Conv2D(int(backend.int_shape(x)[-1] * reduction), 1, use_bias=False, kernel_initializer='he_normal', name=name + '_conv')(x) 43 | 44 | if dilation == 1: 45 | x = layers.AveragePooling2D(2, strides=2, name=name + '_pool')(x) 46 | return x 47 | 48 | 49 | def __call__(self, inputs, output_stages='c5'): 50 | """ 51 | call for DenseNet. 52 | :param inputs: a 4-D tensor. 53 | :param output_stages: str or a list of str containing the output stages. 54 | :return: the output of different stages. 55 | """ 56 | x = layers.ZeroPadding2D(padding=((3, 3), (3, 3)))(inputs) 57 | x = layers.Conv2D(64, 7, strides=2, use_bias=False, name='conv1/conv')(x) 58 | x = layers.BatchNormalization( 59 | axis=-1, epsilon=1.001e-5, name='conv1/bn')(x) 60 | x = layers.Activation('swish', name='conv1/swish')(x) 61 | x = layers.ZeroPadding2D(padding=((1, 1), (1, 1)))(x) 62 | x = layers.MaxPooling2D(3, strides=2, name='pool1')(x) 63 | 64 | c1 = x # 64 65 | x = self._dense_block(x, self.growth_rate, self.params[0], name='D1') 66 | x = self._transition_block(x, self.reduction, name='T1') 67 | c2 = x # 32 68 | 69 | x = self._dense_block(x, self.growth_rate, self.params[1], name='D2') 70 | x = self._transition_block(x, self.reduction, name='T2') 71 | c3 = x # 16 72 | 73 | x = self._dense_block(x, self.growth_rate, self.params[2], name='D3') 74 | x = self._transition_block(x, self.reduction, name='T3', dilation=2) 75 | c4 = x # 16 76 | 77 | x = self._dense_block(x, self.growth_rate, self.params[3], name='D4', dilation=2) 78 | x = layers.BatchNormalization(name='bn')(x) 79 | x = apRelu(x) 80 | c5 = x # 16 81 | 82 | self.outputs = {'c1': c1, 'c2': c2, 'c3': c3, 'c4': c4, 'c5': c5} 83 | 84 | if type(output_stages) is not list: 85 | return self.outputs[output_stages] 86 | else: 87 | return [self.outputs[ci] for ci in output_stages] 88 | 89 | 90 | 91 | class DeepLabV3Plus(): 92 | def __init__(self, num_classes, encoder): 93 | 94 | dilation = [1, 2] 95 | 96 | self.dilation = dilation 97 | self.num_classes = num_classes 98 | self.encoder = encoder 99 | 100 | def __call__(self, inputs): 101 | 102 | return self._deeplab_v3_plus(inputs) 103 | 104 | def _deeplab_v3_plus(self, inputs): 105 | 106 | _, h, w, _ = backend.int_shape(inputs) 107 | self.aspp_size = (16, 16) 108 | 109 | c2, c5 = self.encoder(inputs, output_stages=['c1', 'c5']) 110 | # 64 ; 16 111 | 112 | x = self._aspp(c5, 256) 113 | x = layers.Dropout(rate=0.2)(x) 114 | 115 | x = layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(x) 116 | x = self._conv_bn_relu(x, 64, 1, strides=1) # 64 117 | 118 | c2 = self._conv_bn_relu(c2, 64, 1, strides=1) # 64 119 | 120 | x = layers.concatenate([x, c2]) 121 | x = self._conv_bn_relu(x, 256, 3, 1) 122 | x = layers.Dropout(rate=0.5)(x) 123 | 124 | x = self._conv_bn_relu(x, 256, 3, 1) 125 | x = layers.Dropout(rate=0.1)(x) 126 | 127 | x = layers.Conv2D(self.num_classes, 1, strides=1)(x) 128 | x = layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(x) # 256 129 | 130 | outputs = layers.Activation('softmax', name='softmax_out')(x) 131 | return outputs 132 | 133 | def _conv_bn_relu(self, x, filters, kernel_size, strides=1): 134 | x = layers.Conv2D(filters, kernel_size, strides=strides, use_bias=False, kernel_initializer='he_normal', padding='same')(x) 135 | x = layers.BatchNormalization()(x) 136 | x = apRelu(x) 137 | return x 138 | 139 | def _aspp(self, x, out_filters): 140 | xs = list() 141 | x1 = layers.Conv2D(out_filters, 1, strides=1, use_bias=False, kernel_initializer='he_normal')(x) 142 | x1 = global_context_block(x1, out_filters // 4) 143 | xs.append(x1) 144 | 145 | for i in range(3): 146 | xi = layers.Conv2D(out_filters, 3, strides=1, padding='same',kernel_initializer='he_normal', dilation_rate=6 * (i + 1))(x) 147 | xi = global_context_block(xi, out_filters//4) 148 | xs.append(xi) 149 | img_pool = custom_layers.GlobalAveragePooling2D(keep_dims=True)(x) 150 | img_pool = layers.Conv2D(out_filters, 1, 1, kernel_initializer='he_normal')(img_pool) 151 | img_pool = layers.UpSampling2D(size=self.aspp_size, interpolation='bilinear')(img_pool) 152 | xs.append(img_pool) 153 | 154 | x = layers.Concatenate()(xs) 155 | x = layers.Conv2D(out_filters, 1, strides=1, kernel_initializer='he_normal')(x) 156 | x = layers.BatchNormalization()(x) 157 | x = apRelu(x) 158 | 159 | return x 160 | 161 | 162 | def Net(inputs, params, growth_rate, reduction, num_classes): 163 | encoder = DenseNet(params,growth_rate,reduction) 164 | v3P = DeepLabV3Plus(num_classes, encoder) 165 | output = v3P(inputs) 166 | return output 167 | 168 | 169 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | from models.network import Network 2 | from models.fcn import FCN 3 | from models.pspnet import PSPNet 4 | from models.segnet import SegNet 5 | 6 | from models.pan import PAN 7 | from models.deeplab_v3 import DeepLabV3 8 | from models.deeplab_v3_plus import DeepLabV3Plus 9 | from models.refinenet import RefineNet 10 | from models.denseaspp import DenseASPP 11 | from models.bisegnet import BiSegNet 12 | -------------------------------------------------------------------------------- /models/__pycache__/Custom_net.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/Custom_net.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/bisegnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/bisegnet.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/deeplab_v3.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/deeplab_v3.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/deeplab_v3_plus.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/deeplab_v3_plus.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/denseaspp.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/denseaspp.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/fcn.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/fcn.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/network.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/network.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/pan.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/pan.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/pspnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/pspnet.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/refinenet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/refinenet.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/segnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/segnet.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/unet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/models/__pycache__/unet.cpython-37.pyc -------------------------------------------------------------------------------- /models/bisegnet.py: -------------------------------------------------------------------------------- 1 | 2 | from tools import layers as custom_layers 3 | from models import Network 4 | import tensorflow as tf 5 | 6 | layers = tf.keras.layers 7 | models = tf.keras.models 8 | backend = tf.keras.backend 9 | 10 | 11 | class BiSegNet(Network): 12 | def __init__(self, num_classes, version='BiSegNet', base_model='Xception', **kwargs): 13 | """ 14 | The initialization of BiSegNet. 15 | :param num_classes: the number of predicted classes. 16 | :param version: 'BiSegNet' 17 | :param base_model: the backbone model 18 | :param kwargs: other parameters 19 | """ 20 | base_model = 'Xception' if base_model is None else base_model 21 | 22 | assert version == 'BiSegNet' 23 | super(BiSegNet, self).__init__(num_classes, version, base_model, **kwargs) 24 | 25 | def __call__(self, inputs=None, input_size=None, **kwargs): 26 | assert inputs is not None or input_size is not None 27 | 28 | if inputs is None: 29 | assert isinstance(input_size, tuple) 30 | inputs = layers.Input(shape=input_size+(3,)) 31 | return self._bisegnet(inputs) 32 | 33 | def _conv_block(self, x, filters, kernel_size=3, strides=1): 34 | x = layers.Conv2D(filters, kernel_size, strides, padding='same', kernel_initializer='he_normal')(x) 35 | x = layers.BatchNormalization()(x) 36 | x = layers.ReLU()(x) 37 | return x 38 | 39 | def _attention_refinement_module(self, x): 40 | # Global average pooling 41 | _, _, _, c = backend.int_shape(x) 42 | 43 | glb = custom_layers.GlobalAveragePooling2D(keep_dims=True)(x) 44 | glb = layers.Conv2D(c, 1, strides=1, kernel_initializer='he_normal')(glb) 45 | glb = layers.BatchNormalization()(glb) 46 | glb = layers.Activation(activation='sigmoid')(glb) 47 | 48 | x = layers.Multiply()([x, glb]) 49 | 50 | return x 51 | 52 | def _feature_fusion_module(self, input_1, input_2, filters): 53 | inputs = layers.Concatenate()([input_1, input_2]) 54 | inputs = self._conv_block(inputs, filters=filters, kernel_size=3) 55 | 56 | # Global average pooling 57 | _, _, _, c = backend.int_shape(inputs) 58 | 59 | glb = custom_layers.GlobalAveragePooling2D(keep_dims=True)(inputs) 60 | glb = layers.Conv2D(filters, 1, strides=1, activation='relu', kernel_initializer='he_normal')(glb) 61 | glb = layers.Conv2D(filters, 1, strides=1, activation='sigmoid', kernel_initializer='he_normal')(glb) 62 | 63 | x = layers.Multiply()([inputs, glb]) 64 | 65 | return x 66 | 67 | def _bisegnet(self, inputs): 68 | num_classes = self.num_classes 69 | 70 | # the spatial path 71 | sx = self._conv_block(inputs, 64, 3, 2) 72 | sx = self._conv_block(sx, 128, 3, 2) 73 | sx = self._conv_block(sx, 256, 3, 2) 74 | 75 | # the context path 76 | if self.base_model in ['VGG16', 77 | 'VGG19', 78 | 'ResNet50', 79 | 'ResNet101', 80 | 'ResNet152', 81 | 'MobileNetV1', 82 | 'MobileNetV2', 83 | 'Xception', 84 | 'Xception-DeepLab']: 85 | c4, c5 = self.encoder(inputs, output_stages=['c4', 'c5']) 86 | else: 87 | c4, c5 = self.encoder(inputs, output_stages=['c3', 'c5']) 88 | 89 | c4 = self._attention_refinement_module(c4) 90 | c5 = self._attention_refinement_module(c5) 91 | 92 | glb = custom_layers.GlobalAveragePooling2D(keep_dims=True)(c5) 93 | c5 = layers.Multiply()([c5, glb]) 94 | 95 | # combining the paths 96 | c4 = layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(c4) 97 | c5 = layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(c5) 98 | 99 | cx = layers.Concatenate()([c4, c5]) 100 | 101 | x = self._feature_fusion_module(sx, cx, num_classes) 102 | 103 | x = layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(x) 104 | x = layers.Conv2D(num_classes, 1, 1, kernel_initializer='he_normal')(x) 105 | 106 | outputs = x 107 | 108 | return models.Model(inputs, outputs, name=self.version) 109 | -------------------------------------------------------------------------------- /models/deeplab_v3.py: -------------------------------------------------------------------------------- 1 | 2 | from tools import layers as custom_layers 3 | from models import Network 4 | import tensorflow as tf 5 | 6 | layers = tf.keras.layers 7 | models = tf.keras.models 8 | backend = tf.keras.backend 9 | 10 | 11 | class DeepLabV3(Network): 12 | def __init__(self, num_classes, version='DeepLabV3', base_model='ResNet50', **kwargs): 13 | """ 14 | The initialization of DeepLabV3. 15 | :param num_classes: the number of predicted classes. 16 | :param version: 'DeepLabV3' 17 | :param base_model: the backbone model 18 | :param kwargs: other parameters 19 | """ 20 | dilation = [1, 2] 21 | base_model = 'ResNet50' if base_model is None else base_model 22 | 23 | assert version == 'DeepLabV3' 24 | assert base_model in ['VGG16', 25 | 'VGG19', 26 | 'ResNet50', 27 | 'ResNet101', 28 | 'ResNet152', 29 | 'DenseNet121', 30 | 'DenseNet169', 31 | 'DenseNet201', 32 | 'DenseNet264', 33 | 'MobileNetV1', 34 | 'MobileNetV2', 35 | 'Xception-DeepLab'] 36 | super(DeepLabV3, self).__init__(num_classes, version, base_model, dilation, **kwargs) 37 | self.dilation = dilation 38 | 39 | def __call__(self, inputs=None, input_size=None, **kwargs): 40 | assert inputs is not None or input_size is not None 41 | 42 | if inputs is None: 43 | assert isinstance(input_size, tuple) 44 | inputs = layers.Input(shape=input_size + (3,)) 45 | return self._deeplabv3(inputs) 46 | 47 | def _deeplabv3(self, inputs): 48 | multi_grid = [1, 2, 4] 49 | num_classes = self.num_classes 50 | dilation = self.dilation 51 | 52 | _, h, w, _ = backend.int_shape(inputs) 53 | self.aspp_size = (h // 16, w // 16) 54 | 55 | x = self.encoder(inputs, output_stages='c4') 56 | 57 | x = self._conv_block(x, 3, [512, 512, 2048], stage=5, block='a', dilation=dilation[1]) 58 | for i in range(2): 59 | x = self._identity_block(x, 3, [512, 512, 2048], 60 | stage=5, 61 | block=chr(ord('b') + i), 62 | dilation=dilation[1] * multi_grid[i]) 63 | x = self._aspp(x, 256) 64 | x = layers.Conv2D(num_classes, 1, strides=1, kernel_initializer='he_normal')(x) 65 | x = layers.UpSampling2D(size=(16, 16), interpolation='bilinear')(x) 66 | 67 | outputs = x 68 | return models.Model(inputs, outputs, name=self.version) 69 | 70 | def _aspp(self, x, out_filters): 71 | xs = list() 72 | x1 = layers.Conv2D(out_filters, 1, strides=1, kernel_initializer='he_normal')(x) 73 | xs.append(x1) 74 | 75 | for i in range(3): 76 | xi = layers.Conv2D(out_filters, 3, 77 | strides=1, 78 | padding='same', 79 | dilation_rate=6 * (i + 1))(x) 80 | xs.append(xi) 81 | img_pool = custom_layers.GlobalAveragePooling2D(keep_dims=True)(x) 82 | img_pool = layers.Conv2D(out_filters, 1, 1, kernel_initializer='he_normal')(img_pool) 83 | img_pool = layers.UpSampling2D(size=self.aspp_size, interpolation='bilinear')(img_pool) 84 | xs.append(img_pool) 85 | 86 | x = layers.Concatenate()(xs) 87 | x = layers.Conv2D(out_filters, 1, strides=1, kernel_initializer='he_normal')(x) 88 | x = layers.BatchNormalization()(x) 89 | 90 | return x 91 | 92 | def _identity_block(self, input_tensor, kernel_size, filters, stage, block, dilation=1): 93 | """The identity block is the block that has no conv layer at shortcut. 94 | 95 | # Arguments 96 | input_tensor: input tensor 97 | kernel_size: default 3, the kernel size of 98 | middle conv layer at main path 99 | filters: list of integers, the filters of 3 conv layer at main path 100 | stage: integer, current stage label, used for generating layer names 101 | block: 'a','b'..., current block label, used for generating layer names 102 | 103 | # Returns 104 | Output tensor for the block. 105 | """ 106 | filters1, filters2, filters3 = filters 107 | if backend.image_data_format() == 'channels_last': 108 | bn_axis = 3 109 | else: 110 | bn_axis = 1 111 | conv_name_base = 'res' + str(stage) + block + '_branch' 112 | bn_name_base = 'bn' + str(stage) + block + '_branch' 113 | 114 | x = layers.Conv2D(filters1, (1, 1), 115 | kernel_initializer='he_normal', 116 | name=conv_name_base + '2a')(input_tensor) 117 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x) 118 | x = layers.Activation('relu')(x) 119 | 120 | x = layers.Conv2D(filters2, kernel_size, 121 | padding='same', 122 | kernel_initializer='he_normal', 123 | name=conv_name_base + '2b', 124 | dilation_rate=dilation)(x) 125 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x) 126 | x = layers.Activation('relu')(x) 127 | 128 | x = layers.Conv2D(filters3, (1, 1), 129 | kernel_initializer='he_normal', 130 | name=conv_name_base + '2c')(x) 131 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x) 132 | 133 | x = layers.add([x, input_tensor]) 134 | x = layers.Activation('relu')(x) 135 | return x 136 | 137 | def _conv_block(self, 138 | input_tensor, 139 | kernel_size, 140 | filters, 141 | stage, 142 | block, 143 | strides=(2, 2), 144 | dilation=1): 145 | """A block that has a conv layer at shortcut. 146 | 147 | # Arguments 148 | input_tensor: input tensor 149 | kernel_size: default 3, the kernel size of 150 | middle conv layer at main path 151 | filters: list of integers, the filters of 3 conv layer at main path 152 | stage: integer, current stage label, used for generating layer names 153 | block: 'a','b'..., current block label, used for generating layer names 154 | strides: Strides for the first conv layer in the block. 155 | 156 | # Returns 157 | Output tensor for the block. 158 | 159 | Note that from stage 3, 160 | the first conv layer at main path is with strides=(2, 2) 161 | And the shortcut should have strides=(2, 2) as well 162 | """ 163 | filters1, filters2, filters3 = filters 164 | if backend.image_data_format() == 'channels_last': 165 | bn_axis = 3 166 | else: 167 | bn_axis = 1 168 | conv_name_base = 'res' + str(stage) + block + '_branch' 169 | bn_name_base = 'bn' + str(stage) + block + '_branch' 170 | 171 | strides = (1, 1) if dilation > 1 else strides 172 | 173 | x = layers.Conv2D(filters1, (1, 1), 174 | strides=strides, 175 | name=conv_name_base + '2a', 176 | kernel_initializer='he_normal')(input_tensor) 177 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x) 178 | x = layers.Activation('relu')(x) 179 | 180 | x = layers.Conv2D(filters2, kernel_size, 181 | padding='same', 182 | name=conv_name_base + '2b', 183 | kernel_initializer='he_normal', 184 | dilation_rate=dilation)(x) 185 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x) 186 | x = layers.Activation('relu')(x) 187 | 188 | x = layers.Conv2D(filters3, (1, 1), 189 | name=conv_name_base + '2c', 190 | kernel_initializer='he_normal')(x) 191 | x = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x) 192 | 193 | shortcut = layers.Conv2D(filters3, (1, 1), 194 | strides=strides, 195 | name=conv_name_base + '1', 196 | kernel_initializer='he_normal')(input_tensor) 197 | shortcut = layers.BatchNormalization(axis=bn_axis, name=bn_name_base + '1')(shortcut) 198 | 199 | x = layers.add([x, shortcut]) 200 | x = layers.Activation('relu')(x) 201 | return x 202 | -------------------------------------------------------------------------------- /models/deeplab_v3_plus.py: -------------------------------------------------------------------------------- 1 | 2 | from tools import layers as custom_layers 3 | from models import Network 4 | import tensorflow as tf 5 | from tensorflow.keras import regularizers 6 | layers = tf.keras.layers 7 | models = tf.keras.models 8 | backend = tf.keras.backend 9 | 10 | 11 | class DeepLabV3Plus(Network): 12 | def __init__(self, num_classes, version='DeepLabV3Plus', base_model='Xception-DeepLab', **kwargs): 13 | """ 14 | The initialization of DeepLabV3Plus. 15 | :param num_classes: the number of predicted classes. 16 | :param version: 'DeepLabV3Plus' 17 | :param base_model: the backbone model 18 | :param kwargs: other parameters 19 | """ 20 | dilation = [1, 2] 21 | base_model = 'Xception-DeepLab' if base_model is None else base_model 22 | 23 | assert version == 'DeepLabV3Plus' 24 | assert base_model in ['VGG16', 25 | 'VGG19', 26 | 'ResNet50', 27 | 'ResNet101', 28 | 'ResNet152', 29 | 'DenseNet121', 30 | 'DenseNet169', 31 | 'DenseNet201', 32 | 'DenseNet264', 33 | 'MobileNetV1', 34 | 'MobileNetV2', 35 | 'Xception-DeepLab'] 36 | super(DeepLabV3Plus, self).__init__(num_classes, version, base_model, dilation, **kwargs) 37 | self.dilation = dilation 38 | 39 | def __call__(self, inputs=None, input_size=None, **kwargs): 40 | assert inputs is not None or input_size is not None 41 | 42 | if inputs is None: 43 | assert isinstance(input_size, tuple) 44 | inputs = layers.Input(shape=input_size + (3,)) 45 | return self._deeplab_v3_plus(inputs) 46 | 47 | def _deeplab_v3_plus(self, inputs): 48 | num_classes = self.num_classes 49 | _, h, w, _ = backend.int_shape(inputs) 50 | self.aspp_size = (h // 16, w // 16) 51 | 52 | if self.base_model in ['VGG16', 53 | 'VGG19', 54 | 'ResNet50', 55 | 'ResNet101', 56 | 'ResNet152', 57 | 'MobileNetV1', 58 | 'MobileNetV2']: 59 | c2, c5 = self.encoder(inputs, output_stages=['c2', 'c5']) 60 | else: 61 | c2, c5 = self.encoder(inputs, output_stages=['c1', 'c5']) 62 | 63 | x = self._aspp(c5, 256) 64 | x = layers.Dropout(rate=0.5)(x) 65 | 66 | x = layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(x) 67 | x = self._conv_bn_relu(x, 48, 1, strides=1) 68 | 69 | x = layers.Concatenate()([x, c2]) 70 | x = self._conv_bn_relu(x, 256, 3, 1) 71 | x = layers.Dropout(rate=0.5)(x) 72 | 73 | x = self._conv_bn_relu(x, 256, 3, 1) 74 | x = layers.Dropout(rate=0.1)(x) 75 | 76 | x = layers.Conv2D(num_classes, 1, strides=1)(x) 77 | x = layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(x) 78 | 79 | outputs = layers.Activation('softmax', name='softmax_1')(x) 80 | return models.Model(inputs, outputs, name=self.version) 81 | 82 | def _conv_bn_relu(self, x, filters, kernel_size, strides=1): 83 | x = layers.Conv2D(filters, kernel_size, strides=strides, padding='same')(x) 84 | x = layers.BatchNormalization()(x) 85 | x = layers.ReLU()(x) 86 | return x 87 | 88 | def _aspp(self, x, out_filters): 89 | xs = list() 90 | x1 = layers.Conv2D(out_filters, 1, strides=1)(x) 91 | xs.append(x1) 92 | 93 | for i in range(3): 94 | xi = layers.Conv2D(out_filters, 3, strides=1, padding='same', dilation_rate=6 * (i + 1))(x) 95 | xs.append(xi) 96 | img_pool = custom_layers.GlobalAveragePooling2D(keep_dims=True)(x) 97 | img_pool = layers.Conv2D(out_filters, 1, 1, kernel_initializer='he_normal')(img_pool) 98 | img_pool = layers.UpSampling2D(size=self.aspp_size, interpolation='bilinear')(img_pool) 99 | xs.append(img_pool) 100 | 101 | x = layers.Concatenate()(xs) 102 | x = layers.Conv2D(out_filters, 1, strides=1, kernel_initializer='he_normal')(x) 103 | x = layers.BatchNormalization()(x) 104 | 105 | return x 106 | -------------------------------------------------------------------------------- /models/denseaspp.py: -------------------------------------------------------------------------------- 1 | 2 | from tools import layers as custom_layers 3 | from models import Network 4 | import tensorflow as tf 5 | 6 | layers = tf.keras.layers 7 | models = tf.keras.models 8 | backend = tf.keras.backend 9 | 10 | 11 | class DenseASPP(Network): 12 | def __init__(self, num_classes, version='DenseASPP', base_model='DenseNet121', **kwargs): 13 | """ 14 | The initialization of DenseASPP based. 15 | :param num_classes: the number of predicted classes. 16 | :param version: 'DenseASPP' 17 | :param base_model: the backbone model 18 | :param kwargs: other parameters 19 | """ 20 | dilation = [2, 4] 21 | base_model = 'DenseNet121' if base_model is None else base_model 22 | 23 | assert version == 'DenseASPP' 24 | assert base_model in ['VGG16', 25 | 'VGG19', 26 | 'ResNet50', 27 | 'ResNet101', 28 | 'ResNet152', 29 | 'DenseNet121', 30 | 'DenseNet169', 31 | 'DenseNet201', 32 | 'DenseNet264', 33 | 'MobileNetV1', 34 | 'MobileNetV2', 35 | 'Xception-DeepLab'] 36 | super(DenseASPP, self).__init__(num_classes, version, base_model, dilation, **kwargs) 37 | 38 | def __call__(self, inputs=None, input_size=None, **kwargs): 39 | assert inputs is not None or input_size is not None 40 | 41 | if inputs is None: 42 | assert isinstance(input_size, tuple) 43 | inputs = layers.Input(shape=input_size + (3,)) 44 | return self._denseaspp(inputs) 45 | 46 | def _dilated_conv_block(self, inputs, filters, kernel_size=3, rate=1): 47 | x = layers.BatchNormalization()(inputs) 48 | x = layers.ReLU()(x) 49 | x = layers.Conv2D(filters, kernel_size, 50 | padding='same', 51 | dilation_rate=rate, 52 | kernel_initializer='he_normal')(x) 53 | return x 54 | 55 | def _denseaspp(self, inputs): 56 | _, inputs_h, inputs_w, _ = backend.int_shape(inputs) 57 | aspp_size = inputs_h // 8, inputs_w // 8 58 | num_classes = self.num_classes 59 | 60 | c5 = self.encoder(inputs, output_stages='c5') 61 | 62 | # First block rate=3 63 | d3 = self._dilated_conv_block(c5, 256, 1) 64 | d3 = self._dilated_conv_block(d3, 64, 3, rate=3) 65 | 66 | # Second block rate=6 67 | d4 = layers.Concatenate()([c5, d3]) 68 | d4 = self._dilated_conv_block(d4, 256, 1) 69 | d4 = self._dilated_conv_block(d4, 64, 3, rate=6) 70 | 71 | # Third block rate=12 72 | d5 = layers.Concatenate()([c5, d3, d4]) 73 | d5 = self._dilated_conv_block(d5, 256, 1) 74 | d5 = self._dilated_conv_block(d5, 64, 3, rate=12) 75 | 76 | # Forth block rate=18 77 | d6 = layers.Concatenate()([c5, d3, d4, d5]) 78 | d6 = self._dilated_conv_block(d6, 256, 1) 79 | d6 = self._dilated_conv_block(d6, 64, 3, rate=18) 80 | 81 | # Fifth block rate=24 82 | d7 = layers.Concatenate()([c5, d3, d4, d5, d6]) 83 | d7 = self._dilated_conv_block(d7, 256, 1) 84 | d7 = self._dilated_conv_block(d7, 64, 3, rate=24) 85 | 86 | x = layers.Concatenate()([c5, d3, d4, d5, d6, d7]) 87 | x = layers.Conv2D(num_classes, 1, strides=1, kernel_initializer='he_normal')(x) 88 | x = layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(x) 89 | 90 | outputs = x 91 | return models.Model(inputs, outputs, name=self.version) 92 | -------------------------------------------------------------------------------- /models/fcn.py: -------------------------------------------------------------------------------- 1 | 2 | from models import Network 3 | import tensorflow as tf 4 | 5 | layers = tf.keras.layers 6 | models = tf.keras.models 7 | backend = tf.keras.backend 8 | 9 | 10 | class FCN(Network): 11 | def __init__(self, num_classes, version='FCN-8s', base_model='VGG16', **kwargs): 12 | """ 13 | The initialization of FCN-8s/16s/32s. 14 | :param num_classes: the number of predicted classes. 15 | :param version: 'FCN-8s', 'FCN-16s' or 'FCN-32s'. 16 | :param base_model: the backbone model 17 | :param kwargs: other parameters 18 | """ 19 | fcn = {'FCN-8s': self._fcn_8s, 20 | 'FCN-16s': self._fcn_16s, 21 | 'FCN-32s': self._fcn_32s} 22 | base_model = 'VGG16' if base_model is None else base_model 23 | 24 | assert version in fcn 25 | self.fcn = fcn[version] 26 | super(FCN, self).__init__(num_classes, version, base_model, **kwargs) 27 | 28 | def __call__(self, inputs=None, input_size=None, **kwargs): 29 | assert inputs is not None or input_size is not None 30 | 31 | if inputs is None: 32 | assert isinstance(input_size, tuple) 33 | inputs = layers.Input(shape=input_size + (3,)) 34 | return self.fcn(inputs) 35 | 36 | def _conv_relu(self, x, filters, kernel_size=1): 37 | x = layers.Conv2D(filters, kernel_size, padding='same', kernel_initializer='he_normal')(x) 38 | x = layers.ReLU()(x) 39 | return x 40 | 41 | def _fcn_32s(self, inputs): 42 | num_classes = self.num_classes 43 | 44 | x = self.encoder(inputs) 45 | x = self._conv_relu(x, 4096, 7) 46 | x = layers.Dropout(rate=0.5)(x) 47 | x = self._conv_relu(x, 4096, 1) 48 | x = layers.Dropout(rate=0.5)(x) 49 | 50 | x = layers.Conv2D(num_classes, 1, kernel_initializer='he_normal')(x) 51 | x = layers.Conv2DTranspose(num_classes, 64, strides=32, padding='same', kernel_initializer='he_normal')(x) 52 | 53 | outputs = x 54 | return models.Model(inputs, outputs, name=self.version) 55 | 56 | def _fcn_16s(self, inputs): 57 | num_classes = self.num_classes 58 | 59 | if self.base_model in ['DenseNet121', 60 | 'DenseNet169', 61 | 'DenseNet201', 62 | 'DenseNet264', 63 | 'Xception', 64 | 'Xception-DeepLab']: 65 | c4, c5 = self.encoder(inputs, output_stages=['c3', 'c5']) 66 | else: 67 | c4, c5 = self.encoder(inputs, output_stages=['c4', 'c5']) 68 | 69 | x = self._conv_relu(c5, 4096, 7) 70 | x = layers.Dropout(rate=0.5)(x) 71 | x = self._conv_relu(x, 4096, 1) 72 | x = layers.Dropout(rate=0.5)(x) 73 | 74 | x = layers.Conv2D(num_classes, 1, kernel_initializer='he_normal')(x) 75 | x = layers.Conv2DTranspose(num_classes, 4, 76 | strides=2, 77 | padding='same', 78 | kernel_initializer='he_normal')(x) 79 | c4 = layers.Conv2D(num_classes, 1, kernel_initializer='he_normal')(c4) 80 | x = layers.Add()([x, c4]) 81 | 82 | x = layers.Conv2DTranspose(num_classes, 32, 83 | strides=16, 84 | padding='same', 85 | kernel_initializer='he_normal')(x) 86 | 87 | outputs = x 88 | return models.Model(inputs, outputs, name=self.version) 89 | 90 | def _fcn_8s(self, inputs): 91 | num_classes = self.num_classes 92 | 93 | if self.base_model in ['VGG16', 94 | 'VGG19', 95 | 'ResNet50', 96 | 'ResNet101', 97 | 'ResNet152', 98 | 'MobileNetV1', 99 | 'MobileNetV2']: 100 | c3, c4, c5 = self.encoder(inputs, output_stages=['c3', 'c4', 'c5']) 101 | else: 102 | c3, c4, c5 = self.encoder(inputs, output_stages=['c2', 'c3', 'c5']) 103 | 104 | x = self._conv_relu(c5, 4096, 7) 105 | x = layers.Dropout(rate=0.5)(x) 106 | x = self._conv_relu(x, 4096, 1) 107 | x = layers.Dropout(rate=0.5)(x) 108 | 109 | x = layers.Conv2D(num_classes, 1, kernel_initializer='he_normal')(x) 110 | x = layers.Conv2DTranspose(num_classes, 4, 111 | strides=2, 112 | padding='same', 113 | kernel_initializer='he_normal')(x) 114 | c4 = layers.Conv2D(num_classes, 1)(c4) 115 | x = layers.Add()([x, c4]) 116 | 117 | x = layers.Conv2DTranspose(num_classes, 4, 118 | strides=2, 119 | padding='same', 120 | kernel_initializer='he_normal')(x) 121 | c3 = layers.Conv2D(num_classes, 1)(c3) 122 | x = layers.Add()([x, c3]) 123 | 124 | x = layers.Conv2DTranspose(num_classes, 16, 125 | strides=8, 126 | padding='same', 127 | kernel_initializer='he_normal')(x) 128 | 129 | outputs = x 130 | return models.Model(inputs, outputs, name=self.version) 131 | -------------------------------------------------------------------------------- /models/network.py: -------------------------------------------------------------------------------- 1 | 2 | from backBone import * 3 | 4 | 5 | class Network(object): 6 | def __init__(self, num_classes, version='PAN', base_model='ResNet50', dilation=None, **kwargs): 7 | super(Network, self).__init__(**kwargs) 8 | if base_model in ['VGG16', 'VGG19']: 9 | self.encoder = VGG(base_model, dilation=dilation) 10 | elif base_model in ['ResNet50', 'ResNet101', 'ResNet152']: 11 | self.encoder = ResNet(base_model, dilation=dilation) 12 | elif base_model in ['DenseNet121', 'DenseNet169', 'DenseNet201', 'DenseNet264']: 13 | self.encoder = DenseNet(base_model, dilation=dilation) 14 | elif base_model in ['Xception', 'Xception-DeepLab']: 15 | self.encoder = Xception(base_model, dilation=dilation) 16 | elif base_model in ['MobileNetV1', 'MobileNetV2']: 17 | self.encoder = MobileNet(base_model, dilation=dilation) 18 | else: 19 | raise ValueError('The base model {model} is not in the ' 20 | 'supported model list!!!'.format(model=base_model)) 21 | 22 | self.num_classes = num_classes 23 | self.version = version 24 | self.base_model = base_model 25 | 26 | def __call__(self, inputs, **kwargs): 27 | return inputs 28 | 29 | def get_version(self): 30 | return self.version 31 | 32 | def get_base_model(self): 33 | return self.base_model 34 | -------------------------------------------------------------------------------- /models/pan.py: -------------------------------------------------------------------------------- 1 | 2 | from tools import layers as custom_layers 3 | from models import Network 4 | import tensorflow as tf 5 | 6 | layers = tf.keras.layers 7 | models = tf.keras.models 8 | backend = tf.keras.backend 9 | 10 | 11 | class PAN(Network): 12 | def __init__(self, num_classes, version='PAN', base_model='ResNet50', **kwargs): 13 | """ 14 | The initialization of PAN. 15 | :param num_classes: the number of predicted classes. 16 | :param version: 'PAN' 17 | :param base_model: the backbone model 18 | :param kwargs: other parameters 19 | """ 20 | base_model = 'ResNet50' if base_model is None else base_model 21 | assert version == 'PAN' 22 | 23 | dilation = [2, 4] 24 | if base_model in ['VGG16', 25 | 'VGG19', 26 | 'MobileNetV1', 27 | 'MobileNetV2', 28 | 'ResNet50', 29 | 'ResNet101', 30 | 'ResNet152']: 31 | self.up_size = [(1, 1), (1, 1), (2, 2), (4, 4)] 32 | elif base_model in ['DenseNet121', 33 | 'DenseNet169', 34 | 'DenseNet201', 35 | 'DenseNet264', 36 | 'Xception-DeepLab']: 37 | self.up_size = [(1, 1), (1, 1), (1, 1), (8, 8)] 38 | else: 39 | raise ValueError('The base model \'{model}\' is not ' 40 | 'supported in PAN.'.format(model=base_model)) 41 | 42 | super(PAN, self).__init__(num_classes, version, base_model, dilation, **kwargs) 43 | 44 | def __call__(self, inputs=None, input_size=None, **kwargs): 45 | assert inputs is not None or input_size is not None 46 | 47 | if inputs is None: 48 | assert isinstance(input_size, tuple) 49 | inputs = layers.Input(shape=input_size + (3,)) 50 | return self._pan(inputs) 51 | 52 | def _conv_bn_relu(self, x, filters, kernel_size, strides=1): 53 | x = layers.Conv2D(filters, kernel_size, strides, padding='same', kernel_initializer='he_normal')(x) 54 | x = layers.BatchNormalization()(x) 55 | x = layers.ReLU()(x) 56 | return x 57 | 58 | def _fpa(self, x, out_filters): 59 | _, h, w, _ = backend.int_shape(x) 60 | 61 | # global average pooling 62 | glb = custom_layers.GlobalAveragePooling2D(keep_dims=True)(x) 63 | glb = layers.Conv2D(out_filters, 1, strides=1, kernel_initializer='he_normal')(glb) 64 | 65 | # down 66 | down1 = layers.AveragePooling2D(pool_size=(2, 2))(x) 67 | down1 = self._conv_bn_relu(down1, 1, 7, 1) 68 | 69 | down2 = layers.AveragePooling2D(pool_size=(2, 2))(down1) 70 | down2 = self._conv_bn_relu(down2, 1, 5, 1) 71 | 72 | down3 = layers.AveragePooling2D(pool_size=(2, 2))(down2) 73 | down3 = self._conv_bn_relu(down3, 1, 3, 1) 74 | 75 | down1 = self._conv_bn_relu(down1, 1, 7, 1) 76 | down2 = self._conv_bn_relu(down2, 1, 5, 1) 77 | down3 = self._conv_bn_relu(down3, 1, 3, 1) 78 | 79 | # up 80 | up2 = layers.UpSampling2D(size=(2, 2))(down3) 81 | up2 = layers.Add()([up2, down2]) 82 | 83 | up1 = layers.UpSampling2D(size=(2, 2))(up2) 84 | up1 = layers.Add()([up1, down1]) 85 | 86 | up = layers.UpSampling2D(size=(2, 2))(up1) 87 | 88 | x = layers.Conv2D(out_filters, 1, strides=1, kernel_initializer='he_normal')(x) 89 | x = layers.BatchNormalization()(x) 90 | 91 | # multiply 92 | x = layers.Multiply()([x, up]) 93 | 94 | # add 95 | x = layers.Add()([x, glb]) 96 | 97 | return x 98 | 99 | def _gau(self, x, y, out_filters, up_size=(2, 2)): 100 | glb = custom_layers.GlobalAveragePooling2D(keep_dims=True)(y) 101 | glb = layers.Conv2D(out_filters, 1, strides=1, activation='sigmoid', kernel_initializer='he_normal')(glb) 102 | 103 | x = self._conv_bn_relu(x, out_filters, 3, 1) 104 | x = layers.Multiply()([x, glb]) 105 | 106 | y = layers.UpSampling2D(size=up_size, interpolation='bilinear')(y) 107 | 108 | y = layers.Add()([x, y]) 109 | 110 | return y 111 | 112 | def _pan(self, inputs): 113 | num_classes = self.num_classes 114 | up_size = self.up_size 115 | 116 | c2, c3, c4, c5 = self.encoder(inputs, output_stages=['c2', 'c3', 'c4', 'c5']) 117 | 118 | y = self._fpa(c5, num_classes) 119 | 120 | y = self._gau(c4, y, num_classes, up_size[0]) 121 | y = self._gau(c3, y, num_classes, up_size[1]) 122 | y = self._gau(c2, y, num_classes, up_size[2]) 123 | 124 | y = layers.UpSampling2D(size=up_size[3], interpolation='bilinear')(y) 125 | 126 | outputs = y 127 | 128 | return models.Model(inputs, outputs, name=self.version) 129 | -------------------------------------------------------------------------------- /models/pspnet.py: -------------------------------------------------------------------------------- 1 | 2 | from tools import layers as custom_layers 3 | from models import Network 4 | import tensorflow as tf 5 | 6 | layers = tf.keras.layers 7 | models = tf.keras.models 8 | backend = tf.keras.backend 9 | 10 | 11 | class PSPNet(Network): 12 | def __init__(self, num_classes, version='PSPNet', base_model='ResNet50', **kwargs): 13 | """ 14 | The initialization of PSPNet. 15 | :param num_classes: the number of predicted classes. 16 | :param version: 'PSPNet' 17 | :param base_model: the backbone model 18 | :param kwargs: other parameters 19 | """ 20 | dilation = [2, 4] 21 | base_model = 'ResNet50' if base_model is None else base_model 22 | 23 | assert version == 'PSPNet' 24 | assert base_model in ['VGG16', 25 | 'VGG19', 26 | 'ResNet50', 27 | 'ResNet101', 28 | 'ResNet152', 29 | 'DenseNet121', 30 | 'DenseNet169', 31 | 'DenseNet201', 32 | 'DenseNet264', 33 | 'MobileNetV1', 34 | 'MobileNetV2', 35 | 'Xception-DeepLab'] 36 | super(PSPNet, self).__init__(num_classes, version, base_model, dilation, **kwargs) 37 | 38 | def __call__(self, inputs=None, input_size=None, **kwargs): 39 | assert inputs is not None or input_size is not None 40 | 41 | if inputs is None: 42 | assert isinstance(input_size, tuple) 43 | inputs = layers.Input(shape=input_size + (3,)) 44 | return self._pspnet(inputs) 45 | 46 | def _pspnet(self, inputs): 47 | num_classes = self.num_classes 48 | _, inputs_h, inputs_w, _ = backend.int_shape(inputs) 49 | 50 | h, w = inputs_h // 8, inputs_w // 8 51 | x = self.encoder(inputs) 52 | 53 | if not (h % 6 == 0 and w % 6 == 0): 54 | raise ValueError('\'pyramid pooling\' size must be divided by 6, but received {size}'.format(size=(h, w))) 55 | pool_size = [(h, w), 56 | (h // 2, w // 2), 57 | (h // 3, w // 3), 58 | (h // 6, w // 6)] 59 | 60 | # pyramid pooling 61 | x1 = custom_layers.GlobalAveragePooling2D(keep_dims=True)(x) 62 | x1 = layers.Conv2D(512, 1, strides=1, kernel_initializer='he_normal')(x1) 63 | x1 = layers.BatchNormalization()(x1) 64 | x1 = layers.ReLU()(x1) 65 | x1 = layers.UpSampling2D(size=pool_size[0])(x1) 66 | 67 | x2 = layers.AveragePooling2D(pool_size=pool_size[1])(x) 68 | x2 = layers.Conv2D(512, 1, strides=1, kernel_initializer='he_normal')(x2) 69 | x2 = layers.BatchNormalization()(x2) 70 | x2 = layers.ReLU()(x2) 71 | x2 = layers.UpSampling2D(size=pool_size[1])(x2) 72 | 73 | x3 = layers.AveragePooling2D(pool_size=pool_size[2])(x) 74 | x3 = layers.Conv2D(512, 1, strides=1, kernel_initializer='he_normal')(x3) 75 | x3 = layers.BatchNormalization()(x3) 76 | x3 = layers.ReLU()(x3) 77 | x3 = layers.UpSampling2D(size=pool_size[2])(x3) 78 | 79 | x6 = layers.AveragePooling2D(pool_size=pool_size[3])(x) 80 | x6 = layers.Conv2D(512, 1, strides=1, kernel_initializer='he_normal')(x6) 81 | x6 = layers.BatchNormalization()(x6) 82 | x6 = layers.ReLU()(x6) 83 | x6 = layers.UpSampling2D(size=pool_size[3])(x6) 84 | 85 | x = layers.Concatenate()([x, x1, x2, x3, x6]) 86 | 87 | x = layers.Conv2D(512, 3, strides=1, padding='same', kernel_initializer='he_normal')(x) 88 | x = layers.BatchNormalization()(x) 89 | x = layers.ReLU()(x) 90 | 91 | x = layers.Conv2D(num_classes, 1, strides=1, kernel_initializer='he_normal')(x) 92 | x = layers.BatchNormalization()(x) 93 | 94 | x = layers.UpSampling2D(size=(8, 8), interpolation='bilinear')(x) 95 | 96 | outputs = x 97 | 98 | return models.Model(inputs, outputs, name=self.version) 99 | -------------------------------------------------------------------------------- /models/refinenet.py: -------------------------------------------------------------------------------- 1 | 2 | from models import Network 3 | import tensorflow as tf 4 | 5 | layers = tf.keras.layers 6 | models = tf.keras.models 7 | backend = tf.keras.backend 8 | 9 | 10 | class RefineNet(Network): 11 | def __init__(self, num_classes, version='RefineNet', base_model='ResNet50', **kwargs): 12 | """ 13 | The initialization of RefineNet. 14 | :param num_classes: the number of predicted classes. 15 | :param version: 'RefineNet' 16 | :param base_model: the backbone model 17 | :param kwargs: other parameters 18 | """ 19 | base_model = 'ResNet50' if base_model is None else base_model 20 | 21 | assert version == 'RefineNet' 22 | assert base_model in ['VGG16', 23 | 'VGG19', 24 | 'ResNet50', 25 | 'ResNet101', 26 | 'ResNet152', 27 | 'MobileNetV1', 28 | 'MobileNetV2'] 29 | super(RefineNet, self).__init__(num_classes, version, base_model, **kwargs) 30 | 31 | def __call__(self, inputs=None, input_size=None, **kwargs): 32 | assert inputs is not None or input_size is not None 33 | 34 | if inputs is None: 35 | assert isinstance(input_size, tuple) 36 | inputs = layers.Input(shape=input_size + (3,)) 37 | return self._refinenet(inputs) 38 | 39 | def _refinenet(self, inputs): 40 | num_classes = self.num_classes 41 | 42 | xs = self.encoder(inputs, output_stages=['c2', 'c3', 'c4', 'c5'])[::-1] 43 | 44 | g = [None, None, None, None] 45 | h = [None, None, None, None] 46 | 47 | for i in range(4): 48 | h[i] = layers.Conv2D(256, 1, strides=1, kernel_initializer='he_normal')(xs[i]) 49 | 50 | g[0] = self._refine_block(high_inputs=None, low_inputs=h[0]) 51 | g[1] = self._refine_block(g[0], h[1]) 52 | g[2] = self._refine_block(g[1], h[2]) 53 | g[3] = self._refine_block(g[2], h[3]) 54 | 55 | x = layers.Conv2D(num_classes, 1, strides=1, kernel_initializer='he_normal')(g[3]) 56 | x = layers.UpSampling2D(size=(4, 4), interpolation='bilinear')(x) 57 | 58 | outputs = x 59 | 60 | return models.Model(inputs, outputs, name=self.version) 61 | 62 | def _residual_conv_unit(self, inputs, features=256, kernel_size=3): 63 | x = layers.ReLU()(inputs) 64 | x = layers.Conv2D(features, kernel_size, padding='same', kernel_initializer='he_normal')(x) 65 | x = layers.ReLU()(x) 66 | x = layers.Conv2D(features, kernel_size, padding='same', kernel_initializer='he_normal')(x) 67 | x = layers.Add()([inputs, x]) 68 | return x 69 | 70 | def _chained_residual_pooling(self, inputs, features=256): 71 | x_relu = layers.ReLU()(inputs) 72 | x = layers.MaxPool2D((5, 5), strides=1, padding='same')(x_relu) 73 | x = layers.Conv2D(features, 3, padding='same', kernel_initializer='he_normal')(x) 74 | x_sum_1 = layers.Add()([x, x_relu]) 75 | 76 | x = layers.MaxPool2D((5, 5), strides=1, padding='same')(x_relu) 77 | x = layers.Conv2D(features, 3, padding='same', kernel_initializer='he_normal')(x) 78 | x_sum_2 = layers.Add()([x, x_sum_1]) 79 | 80 | return x_sum_2 81 | 82 | def _multi_resolution_fusion(self, high_inputs=None, low_inputs=None, features=256): 83 | 84 | if high_inputs is None: # refineNet block 4 85 | rcu_low_1 = low_inputs[0] 86 | rcu_low_2 = low_inputs[1] 87 | 88 | rcu_low_1 = layers.Conv2D(features, 3, padding='same', kernel_initializer='he_normal')(rcu_low_1) 89 | rcu_low_2 = layers.Conv2D(features, 3, padding='same', kernel_initializer='he_normal')(rcu_low_2) 90 | 91 | return layers.Add()([rcu_low_1, rcu_low_2]) 92 | 93 | else: 94 | rcu_low_1 = low_inputs[0] 95 | rcu_low_2 = low_inputs[1] 96 | 97 | rcu_low_1 = layers.Conv2D(features, 3, padding='same', kernel_initializer='he_normal')(rcu_low_1) 98 | rcu_low_2 = layers.Conv2D(features, 3, padding='same', kernel_initializer='he_normal')(rcu_low_2) 99 | 100 | rcu_low = layers.Add()([rcu_low_1, rcu_low_2]) 101 | 102 | rcu_high_1 = high_inputs[0] 103 | rcu_high_2 = high_inputs[1] 104 | 105 | rcu_high_1 = layers.UpSampling2D(size=(2, 2), interpolation='bilinear')( 106 | layers.Conv2D(features, 3, padding='same', kernel_initializer='he_normal')(rcu_high_1)) 107 | rcu_high_2 = layers.UpSampling2D(size=(2, 2), interpolation='bilinear')( 108 | layers.Conv2D(features, 3, padding='same', kernel_initializer='he_normal')(rcu_high_2)) 109 | 110 | rcu_high = layers.Add()([rcu_high_1, rcu_high_2]) 111 | 112 | return layers.Add()([rcu_low, rcu_high]) 113 | 114 | def _refine_block(self, high_inputs=None, low_inputs=None): 115 | 116 | if high_inputs is None: # block 4 117 | rcu_low_1 = self._residual_conv_unit(low_inputs, features=256) 118 | rcu_low_2 = self._residual_conv_unit(low_inputs, features=256) 119 | rcu_low = [rcu_low_1, rcu_low_2] 120 | 121 | fuse = self._multi_resolution_fusion(high_inputs=None, low_inputs=rcu_low, features=256) 122 | fuse_pooling = self._chained_residual_pooling(fuse, features=256) 123 | output = self._residual_conv_unit(fuse_pooling, features=256) 124 | return output 125 | else: 126 | rcu_low_1 = self._residual_conv_unit(low_inputs, features=256) 127 | rcu_low_2 = self._residual_conv_unit(low_inputs, features=256) 128 | rcu_low = [rcu_low_1, rcu_low_2] 129 | 130 | rcu_high_1 = self._residual_conv_unit(high_inputs, features=256) 131 | rcu_high_2 = self._residual_conv_unit(high_inputs, features=256) 132 | rcu_high = [rcu_high_1, rcu_high_2] 133 | 134 | fuse = self._multi_resolution_fusion(rcu_high, rcu_low, features=256) 135 | fuse_pooling = self._chained_residual_pooling(fuse, features=256) 136 | output = self._residual_conv_unit(fuse_pooling, features=256) 137 | return output 138 | -------------------------------------------------------------------------------- /models/segnet.py: -------------------------------------------------------------------------------- 1 | 2 | from models import Network 3 | import tensorflow as tf 4 | 5 | layers = tf.keras.layers 6 | models = tf.keras.models 7 | backend = tf.keras.backend 8 | 9 | 10 | class SegNet(Network): 11 | def __init__(self, num_classes, version='SegNet', base_model='VGG16', **kwargs): 12 | """ 13 | The initialization of SegNet or Bayesian-SegNet. 14 | :param num_classes: the number of predicted classes. 15 | :param version: 'SegNet' or 'Bayesian-SegNet'. 16 | :param base_model: the backbone model 17 | :param kwargs: other parameters 18 | """ 19 | base_model = 'VGG16' if base_model is None else base_model 20 | assert version in ['SegNet', 'Bayesian-SegNet'] 21 | assert base_model in ['VGG16', 22 | 'VGG19', 23 | 'ResNet50', 24 | 'ResNet101', 25 | 'ResNet152', 26 | 'DenseNet121', 27 | 'DenseNet169', 28 | 'DenseNet201', 29 | 'DenseNet269', 30 | 'MobileNetV1', 31 | 'MobileNetV2', 32 | 'Xception', 33 | 'Xception-DeepLab'] 34 | super(SegNet, self).__init__(num_classes, version, base_model, **kwargs) 35 | 36 | def __call__(self, inputs=None, input_size=None, **kwargs): 37 | assert inputs is not None or input_size is not None 38 | 39 | if inputs is None: 40 | assert isinstance(input_size, tuple) 41 | inputs = layers.Input(shape=input_size + (3,)) 42 | return self._segnet(inputs) 43 | 44 | def _conv_bn_relu(self, x, filters, kernel_size=1, strides=1): 45 | x = layers.Conv2D(filters, kernel_size, 46 | strides=strides, 47 | padding='same', 48 | kernel_initializer='he_normal')(x) 49 | x = layers.BatchNormalization()(x) 50 | x = layers.ReLU()(x) 51 | return x 52 | 53 | def _segnet(self, inputs): 54 | num_classes = self.num_classes 55 | dropout = True if self.version == 'Bayesian-SegNet' else False 56 | 57 | x = self.encoder(inputs) 58 | 59 | if dropout: 60 | x = layers.Dropout(rate=0.5)(x) 61 | x = layers.UpSampling2D(size=(2, 2))(x) 62 | x = self._conv_bn_relu(x, 512, 3, strides=1) 63 | x = self._conv_bn_relu(x, 512, 3, strides=1) 64 | x = self._conv_bn_relu(x, 512, 3, strides=1) 65 | 66 | if dropout: 67 | x = layers.Dropout(rate=0.5)(x) 68 | x = layers.UpSampling2D(size=(2, 2))(x) 69 | x = self._conv_bn_relu(x, 512, 3, strides=1) 70 | x = self._conv_bn_relu(x, 512, 3, strides=1) 71 | x = self._conv_bn_relu(x, 256, 3, strides=1) 72 | 73 | if dropout: 74 | x = layers.Dropout(rate=0.5)(x) 75 | x = layers.UpSampling2D(size=(2, 2))(x) 76 | x = self._conv_bn_relu(x, 256, 3, strides=1) 77 | x = self._conv_bn_relu(x, 256, 3, strides=1) 78 | x = self._conv_bn_relu(x, 128, 3, strides=1) 79 | 80 | if dropout: 81 | x = layers.Dropout(rate=0.5)(x) 82 | x = layers.UpSampling2D(size=(2, 2))(x) 83 | x = self._conv_bn_relu(x, 128, 3, strides=1) 84 | x = self._conv_bn_relu(x, 64, 3, strides=1) 85 | 86 | if dropout: 87 | x = layers.Dropout(rate=0.5)(x) 88 | x = layers.UpSampling2D(size=(2, 2))(x) 89 | x = self._conv_bn_relu(x, 64, 3, strides=1) 90 | x = layers.Conv2D(num_classes, 1, 91 | strides=1, 92 | kernel_initializer='he_normal')(x) 93 | x = layers.BatchNormalization()(x) 94 | 95 | outputs = x 96 | return models.Model(inputs, outputs, name=self.version) 97 | -------------------------------------------------------------------------------- /models/unet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.keras import regularizers 3 | from tensorflow.keras.layers import Conv2D, MaxPooling2D, UpSampling2D, BatchNormalization, Reshape, Permute, Activation, Convolution2DTranspose, Concatenate 4 | from tensorflow.keras import backend as K 5 | 6 | 7 | def conv_block(input_tensor, filters, kernel_size, name, strides, padding='same', dila=1): 8 | x = Conv2D(filters, kernel_size, strides=strides, name=name, padding=padding, 9 | kernel_regularizer=regularizers.l1_l2(l1=1e-5, l2=1e-5), dilation_rate=dila)(input_tensor) 10 | x = BatchNormalization(name='bn_' + name)(x) 11 | x = Activation('relu')(x) 12 | return x 13 | 14 | # ----- 15 | 16 | 17 | def Net(n_classes, input): 18 | 19 | # ---------left branch ----- 20 | x = conv_block(input, 32, (3, 3), strides=1, name='L_conv1-1') 21 | L1 = conv_block(x, 32, (3, 3), strides=1, name='L_conv1-2') 22 | x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(L1) 23 | # 256 -> 128 24 | 25 | x = conv_block(x, 64, (3, 3), strides=1, name='L_conv2-1') 26 | L2 = conv_block(x, 64, (3, 3), strides=1, name='L_conv2-2') 27 | x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(L2) 28 | # 128 -> 64 29 | 30 | x = conv_block(x, 128, (3, 3), strides=1, name='L_conv3-1') 31 | L3 = conv_block(x, 128, (3, 3), strides=1, name='L_conv3-2') 32 | x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(L3) 33 | # 64 -> 32 34 | 35 | x = conv_block(x, 256, (3, 3), strides=1, name='L_conv4-1') 36 | L4 = conv_block(x, 256, (3, 3), strides=1, name='L_conv4-2') 37 | x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(L4) 38 | # 32 -> 16 39 | 40 | x = conv_block(x, 512, (3, 3), strides=1, name='bottom-1') 41 | x = conv_block(x, 512, (3, 3), strides=1, dila=2, name='bottom-2') 42 | L5 = conv_block(x, 512, (3, 3), strides=1, name='bottom-3') 43 | # 16 44 | 45 | # ---------Right branch ----- 46 | 47 | # 16 -> 32 48 | x = Convolution2DTranspose( 49 | 256, kernel_size=2, strides=2, padding='same', name='R_conv1-1')(L5) 50 | x = BatchNormalization(name='R_conv1-1_' + 'bn')(x) 51 | x = conv_block(Concatenate(axis=-1)([x, L4]), 256, (3, 3), strides=1, name='R_conv1-2') 52 | x = conv_block(x, 256, (3, 3), strides=1, name='R_conv1-3') 53 | 54 | # 32 -> 64 55 | x = Convolution2DTranspose( 56 | 128, kernel_size=2, strides=2, padding='same', name='R_conv2-1')(x) 57 | x = BatchNormalization(name='R_conv2-1_' + 'bn')(x) 58 | x = conv_block(Concatenate(axis=-1)([x, L3]), 128, (3, 3), strides=1, name='R_conv2-2') 59 | x = conv_block(x, 128, (3, 3), strides=1, name='R_conv2-3') 60 | 61 | # 64 -> 128 62 | x = Convolution2DTranspose( 63 | 64, kernel_size=2, strides=2, padding='same', name='R_conv3-1')(x) 64 | x = BatchNormalization(name='R_conv3-1_' + 'bn')(x) 65 | x = conv_block(Concatenate(axis=-1) ([x, L2]), 64, (3, 3), strides=1, name='R_conv3-2') 66 | x = conv_block(x, 64, (3, 3), strides=1, name='R_conv3-3') 67 | 68 | # 128 -> 256 69 | x = Convolution2DTranspose( 70 | 32, kernel_size=2, strides=2, padding='same', name='R_conv4-1')(x) 71 | x = BatchNormalization(name='R_conv4-1_' + 'bn')(x) 72 | x = conv_block(Concatenate(axis=-1) ([x, L1]), 32, (3, 3), strides=1, name='R_conv4-2') 73 | x = conv_block(x, 32, (3, 3), strides=1, name='R_conv4-3') 74 | 75 | final = Conv2D(n_classes, (1, 1), name='final_out')(x) 76 | 77 | final = Activation('softmax', name='softmax_1')(final) 78 | 79 | return final 80 | -------------------------------------------------------------------------------- /predict.py: -------------------------------------------------------------------------------- 1 | from tensorflow.keras.models import load_model 2 | import cv2 3 | import numpy as np 4 | import os, sys 5 | from tqdm import tqdm 6 | from tools.dataloader import test_data_generator 7 | import config 8 | from tools.layers import GlobalAveragePooling2D 9 | 10 | os.environ['CUDA_VISIBLE_DEVICES'] = '0' 11 | cfg = config.Config('test') 12 | 13 | 14 | def predict_single(input_path, output_path, model, n_class): 15 | 16 | for image in tqdm(os.listdir(input_path)): 17 | index, _ = os.path.splitext(image) 18 | img = cv2.imread(os.path.join(input_path, image), cv2.IMREAD_UNCHANGED) 19 | img = np.float32(img) / 127.5 - 1 20 | pr = model.predict(np.expand_dims(img, axis=0), verbose=1)[0] 21 | pr = pr.reshape((256, 256, n_class)).argmax(axis=2) 22 | seg_img = np.zeros((256, 256), dtype=np.uint16) 23 | for c in range(n_class): 24 | seg_img[pr[:, :] == c] = int((c + 1) * 100) 25 | cv2.imwrite(os.path.join(output_path, index + ".png"), seg_img) 26 | 27 | 28 | def predict_batch(input_path, output_path, model, n_class): 29 | 30 | g = test_data_generator(input_path, cfg.batch_size) 31 | for x, r in g: 32 | out = model.predict(x, verbose=0) 33 | for i in range(out.shape[0]): 34 | pr = out[i].reshape((256, 256, n_class)).argmax(axis=2) 35 | seg_img = np.zeros((256, 256), dtype=np.uint16) 36 | for c in range(n_class): 37 | seg_img[pr[:, :] == c] = int((c + 1) * 100) 38 | cv2.imwrite(os.path.join(output_path, r[i].replace('tif', 'png')), seg_img) 39 | 40 | 41 | if __name__ == "__main__": 42 | weights_path = cfg.weight_path 43 | input_path = cfg.data_path 44 | output_path = cfg.output_path 45 | n_class = cfg.n_classes 46 | cfg.check_folder(output_path) 47 | 48 | if weights_path is None: 49 | print('weights_path ERROR!') 50 | sys.exit() 51 | print(f'loaded : {weights_path}') 52 | co = {'GlobalAveragePooling2D' : GlobalAveragePooling2D} 53 | model = load_model(weights_path, custom_objects= co) 54 | predict_batch(input_path, output_path, model, n_class) -------------------------------------------------------------------------------- /tools/Block_tricks.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | layers = tf.keras.layers 3 | backend = tf.keras.backend 4 | from tensorflow.keras.regularizers import l2 5 | from tools import layers as custom_layers 6 | 7 | 8 | # An adaptively parametric rectifier linear unit (APReLU) 9 | def apRelu(inputs): 10 | # get the number of channels 11 | channels = inputs.get_shape().as_list()[-1] 12 | # get a zero feature map 13 | zeros_input = layers.subtract([inputs, inputs]) 14 | # get a feature map with only positive features 15 | pos_input = layers.Activation('relu')(inputs) 16 | # get a feature map with only negative features 17 | neg_input = layers.Minimum()([inputs,zeros_input]) 18 | # define a network to obtain the scaling coefficients 19 | scales_p = layers.GlobalAveragePooling2D()(pos_input) 20 | scales_n = layers.GlobalAveragePooling2D()(neg_input) 21 | scales = layers.Concatenate()([scales_n, scales_p]) 22 | scales = layers.Dense(channels//4, activation='linear', kernel_initializer='he_normal')(scales) 23 | scales = layers.BatchNormalization(momentum=0.9)(scales) 24 | scales = layers.Activation('relu')(scales) 25 | scales = layers.Dense(channels, activation='linear', kernel_initializer='he_normal')(scales) 26 | scales = layers.BatchNormalization(momentum=0.9)(scales) 27 | scales = layers.Activation('sigmoid')(scales) 28 | scales = layers.Reshape((1,1,channels))(scales) 29 | # apply a paramtetric relu 30 | neg_part = layers.multiply([scales, neg_input]) 31 | return layers.add([pos_input, neg_part]) 32 | 33 | def _conv_bn_relu(x, filters, kernel_size, strides=1): 34 | x = layers.Conv2D(filters, kernel_size, strides=strides, use_bias=False, kernel_initializer='he_normal', 35 | padding='same')(x) 36 | x = layers.BatchNormalization()(x) 37 | x = apRelu(x) 38 | return x 39 | 40 | def backend_expand_dims_1(x): 41 | return backend.expand_dims(x, axis=1) 42 | 43 | def backend_expand_dims_last( x): 44 | return backend.expand_dims(x, axis=-1) 45 | 46 | def backend_dot(x): 47 | return backend.batch_dot(x[0], x[1]) 48 | 49 | def global_context_block(x, channels): 50 | bs, h, w, c = x.shape.as_list() 51 | input_x = x 52 | input_x = layers.Reshape((-1, c))(input_x) # [N, H*W, C] 53 | input_x = layers.Permute((2, 1))(input_x) # [N, C, H*W] 54 | input_x = layers.Lambda(backend_expand_dims_1)(input_x) # [N, 1, C, H*W] 55 | 56 | context_mask = layers.Conv2D(1, (1, 1))(x) 57 | context_mask = layers.Reshape((-1, 1))(context_mask) # [N, H*W, 1] 58 | context_mask = layers.Softmax(axis=1)(context_mask) # [N, H*W, 1] 59 | context_mask = layers.Permute((2, 1))(context_mask) # [N, 1, H*W] 60 | context_mask = layers.Lambda(backend_expand_dims_last)(context_mask) # [N, 1, H*W, 1] 61 | 62 | context = layers.Lambda(backend_dot)([input_x, context_mask]) 63 | context = layers.Reshape((1, 1, c))(context) # [N, 1, 1, C] 64 | 65 | context_transform = _conv_bn_relu(context, channels, 1, strides=1) 66 | context_transform = layers.Conv2D(c, (1, 1))(context_transform) 67 | context_transform = layers.Activation('sigmoid')(context_transform) 68 | x = layers.Multiply()([x, context_transform]) 69 | 70 | context_transform = _conv_bn_relu(context, channels, 1, strides=1) 71 | context_transform = layers.Conv2D(c, (1, 1))(context_transform) 72 | x = layers.add([x, context_transform]) 73 | 74 | return x 75 | # ----- -------------------------------------------------------------------------------- /tools/__pycache__/Block_tricks.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/tools/__pycache__/Block_tricks.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/callbacks.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/tools/__pycache__/callbacks.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/dataloader.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/tools/__pycache__/dataloader.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/layers.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/tools/__pycache__/layers.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/learning_rate.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/tools/__pycache__/learning_rate.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/metrics.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TachibanaYoshino/Remote-sensing-image-semantic-segmentation-tf2/3a0acb862556ce5cc194caba74b5a5794d6c3e79/tools/__pycache__/metrics.cpython-37.pyc -------------------------------------------------------------------------------- /tools/callbacks.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | import numpy as np 4 | 5 | callbacks = tf.keras.callbacks 6 | backend = tf.keras.backend 7 | 8 | 9 | class LearningRateScheduler(callbacks.Callback): 10 | def __init__(self, 11 | schedule, 12 | learning_rate=None, 13 | warmup=False, 14 | steps_per_epoch=None, 15 | num_epochs = 5, 16 | verbose=0): 17 | super(LearningRateScheduler, self).__init__() 18 | self.lr = learning_rate 19 | self.schedule = schedule 20 | self.verbose = verbose 21 | self.warmup_epochs = 5 if warmup else 0 22 | self.warmup_steps = int(steps_per_epoch) * self.warmup_epochs if warmup else 0 23 | self.global_batch = 0 24 | 25 | # lr schedule strategy 26 | if warmup and num_epochs - 5 <= 0: 27 | raise ValueError('num_epochs must be larger than 5 if lr warm up is used.') 28 | if warmup and learning_rate is None: 29 | raise ValueError('learning_rate cannot be None if warmup is used.') 30 | if warmup and steps_per_epoch is None: 31 | raise ValueError('steps_per_epoch cannot be None if warmup is used.') 32 | 33 | def on_train_batch_begin(self, batch, logs=None): 34 | self.global_batch += 1 35 | if self.global_batch < self.warmup_steps: 36 | if not hasattr(self.model.optimizer, 'lr'): 37 | raise ValueError('Optimizer must have a "lr" attribute.') 38 | lr = self.learning_rate * self.global_batch / self.warmup_steps 39 | backend.set_value(self.model.optimizer.lr, lr) 40 | if self.verbose > 0: 41 | print('\nBatch %05d: LearningRateScheduler warming up learning ' 42 | 'rate to %s.' % (self.global_batch, lr)) 43 | 44 | def on_epoch_begin(self, epoch, logs=None): 45 | if not hasattr(self.model.optimizer, 'lr'): 46 | raise ValueError('Optimizer must have a "lr" attribute.') 47 | lr = float(backend.get_value(self.model.optimizer.lr)) 48 | 49 | if epoch >= self.warmup_epochs: 50 | try: # new API 51 | lr = self.schedule(epoch - self.warmup_epochs, lr) 52 | except TypeError: # Support for old API for backward compatibility 53 | lr = self.schedule(epoch - self.warmup_epochs) 54 | if not isinstance(lr, (float, np.float32, np.float64)): 55 | raise ValueError('The output of the "schedule" function ' 56 | 'should be float.') 57 | backend.set_value(self.model.optimizer.lr, lr) 58 | 59 | if self.verbose > 0: 60 | print('\nEpoch %05d: LearningRateScheduler reducing learning ' 61 | 'rate to %s.' % (epoch + 1, lr)) 62 | 63 | def on_epoch_end(self, epoch, logs=None): 64 | logs = logs or {} 65 | logs['lr'] = backend.get_value(self.model.optimizer.lr) 66 | -------------------------------------------------------------------------------- /tools/dataloader.py: -------------------------------------------------------------------------------- 1 | import random 2 | import os 3 | import cv2 4 | import numpy as np 5 | from PIL import Image 6 | from tqdm import tqdm 7 | 8 | # 类别对应 9 | class_values = [100, 200, 300, 400, 500, 600, 700, 800] 10 | 11 | 12 | def rotate_bound(img, angle): 13 | if angle == 90: 14 | out = Image.fromarray(img).transpose(Image.ROTATE_90) 15 | return np.array(out) 16 | if angle == 180: 17 | out = Image.fromarray(img).transpose(Image.ROTATE_180) 18 | return np.array(out) 19 | if angle == 270: 20 | out = Image.fromarray(img).transpose(Image.ROTATE_270) 21 | return np.array(out) 22 | 23 | 24 | def data_augment(x, y): 25 | flag = random.choice([1, 2, 3, 4, 5, 6]) 26 | if flag == 1: 27 | x, y = cv2.flip(x, 1), cv2.flip(y, 1) # Horizontal mirror 28 | if flag == 2: 29 | x, y = cv2.flip(x, 0), cv2.flip(y, 0) # Vertical mirror 30 | if flag == 3: 31 | x, y = rotate_bound(x, 90), rotate_bound(y, 90) 32 | if flag == 4: 33 | x, y = rotate_bound(x, 180), rotate_bound(y, 180) 34 | if flag == 5: 35 | x, y = rotate_bound(x, 270), rotate_bound(y, 270) 36 | else: 37 | pass 38 | return x, y 39 | 40 | 41 | def get_img_label_paths(images_path, labels_path): 42 | res = [] 43 | for dir_entry in os.listdir(images_path): 44 | if os.path.isfile(os.path.join(images_path, dir_entry)): 45 | file_name, _ = os.path.splitext(dir_entry) 46 | res.append((os.path.join(images_path, file_name+".tif"), 47 | os.path.join(labels_path, file_name+".png"))) 48 | return res 49 | 50 | 51 | def get_image_array(img): 52 | return np.float32(img) / 127.5 - 1 53 | 54 | 55 | def get_segmentation_array(img, nClasses): 56 | assert len(img.shape) == 2 57 | seg_labels = np.zeros((256, 256, nClasses)) 58 | 59 | for p in class_values: 60 | img[img == p] = class_values.index(p) 61 | 62 | for c in range(nClasses): 63 | seg_labels[:, :, c] = (img == c).astype(int) 64 | 65 | seg_labels = np.reshape(seg_labels, (256, 256, nClasses)) 66 | return seg_labels 67 | 68 | 69 | def train_data_generator(images_path, labels_path, batch_size, num_class, use_augment): 70 | img_seg_pairs = get_img_label_paths(images_path, labels_path) 71 | pairs_number = len(img_seg_pairs) 72 | 73 | while True: 74 | X, Y = [], [] 75 | random.shuffle(img_seg_pairs) 76 | for i in range(pairs_number): 77 | img = cv2.imread(img_seg_pairs[i][0], cv2.IMREAD_UNCHANGED) 78 | seg = cv2.imread(img_seg_pairs[i][1], cv2.IMREAD_UNCHANGED) 79 | 80 | if use_augment: 81 | img, seg = data_augment(img, seg) 82 | 83 | X.append(get_image_array(img)) 84 | Y.append(get_segmentation_array(seg, num_class)) 85 | 86 | if i == pairs_number - 1: 87 | yield np.array(X), np.array(Y) 88 | X, Y = [], [] 89 | 90 | if len(X) == batch_size: 91 | assert len(X) == len(Y) 92 | yield np.array(X), np.array(Y) 93 | X, Y = [], [] 94 | 95 | 96 | def val_data_generator(images_path, labels_path, batch_size, num_class): 97 | img_seg_pairs = get_img_label_paths(images_path, labels_path) 98 | pairs_number = len(img_seg_pairs) 99 | 100 | while True: 101 | X, Y = [], [] 102 | for i in range(pairs_number): 103 | img = cv2.imread(img_seg_pairs[i][0], cv2.IMREAD_UNCHANGED) 104 | seg = cv2.imread(img_seg_pairs[i][1], cv2.IMREAD_UNCHANGED) 105 | 106 | X.append(get_image_array(img)) 107 | Y.append(get_segmentation_array(seg, num_class)) 108 | 109 | if i == pairs_number - 1: 110 | yield np.array(X), np.array(Y) 111 | X, Y = [], [] 112 | if len(X) == batch_size: 113 | assert len(X) == len(Y) 114 | yield np.array(X), np.array(Y) 115 | X, Y = [], [] 116 | 117 | 118 | def test_data_generator(images_path, batch_size): 119 | images = [] 120 | for dir_entry in os.listdir(images_path): 121 | if os.path.isfile(os.path.join(images_path, dir_entry)): 122 | assert dir_entry.endswith('tif') 123 | images.append(os.path.join(images_path, dir_entry)) 124 | pairs_number = len(images) 125 | 126 | X, R= [], [] 127 | for i in tqdm(range(pairs_number)): 128 | img = cv2.imread(images[i], cv2.IMREAD_UNCHANGED) 129 | X.append(get_image_array(img)) 130 | R.append(os.path.basename(images[i])) 131 | if i == pairs_number - 1: 132 | yield np.array(X), R 133 | X, R = [], [] 134 | if len(X) == batch_size: 135 | yield np.array(X), R 136 | X, R = [], [] 137 | 138 | 139 | 140 | if __name__ == "__main__": 141 | data = [[1, 2, 3, 4], [5, 6, 7, 8], [3, 5, 1, 8]] 142 | arr = np.array(data) 143 | print(arr.shape, arr) 144 | print(arr == 1) 145 | print((arr == 1).astype(int)) 146 | print(arr) 147 | arr[arr == 1] = 100 148 | print(arr) 149 | -------------------------------------------------------------------------------- /tools/layers.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | 4 | layers = tf.keras.layers 5 | backend = tf.keras.backend 6 | 7 | 8 | class GlobalAveragePooling2D(layers.GlobalAveragePooling2D): 9 | def __init__(self, keep_dims=False, **kwargs): 10 | super(GlobalAveragePooling2D, self).__init__(**kwargs) 11 | self.keep_dims = keep_dims 12 | 13 | def call(self, inputs): 14 | if self.keep_dims is False: 15 | return super(GlobalAveragePooling2D, self).call(inputs) 16 | else: 17 | return backend.mean(inputs, axis=[1, 2], keepdims=True) 18 | 19 | def compute_output_shape(self, input_shape): 20 | if self.keep_dims is False: 21 | return super(GlobalAveragePooling2D, self).compute_output_shape(input_shape) 22 | else: 23 | input_shape = tf.TensorShape(input_shape).as_list() 24 | return tf.TensorShape([input_shape[0], 1, 1, input_shape[3]]) 25 | 26 | def get_config(self): 27 | config = super(GlobalAveragePooling2D, self).get_config() 28 | config['keep_dims'] = self.keep_dims 29 | return config 30 | 31 | 32 | class PixelShuffle(layers.Layer): 33 | def __init__(self, block_size=2, **kwargs): 34 | super(PixelShuffle, self).__init__(**kwargs) 35 | if isinstance(block_size, int): 36 | self.block_size = block_size 37 | elif isinstance(block_size, (list, tuple)): 38 | assert len(block_size) == 2 and block_size[0] == block_size[1] 39 | self.block_size = block_size[0] 40 | else: 41 | raise ValueError('error \'block_size\'.') 42 | 43 | def build(self, input_shape): 44 | pass 45 | 46 | def call(self, inputs, **kwargs): 47 | return tf.nn.depth_to_space(inputs, self.block_size) 48 | 49 | def compute_output_shape(self, input_shape): 50 | input_shape = tf.TensorShape(input_shape).as_list() 51 | 52 | _, h, w, c = input_shape 53 | 54 | new_h = h * self.block_size 55 | new_w = w * self.block_size 56 | new_c = c / self.block_size ** 2 57 | 58 | return tf.TensorShape([input_shape[0], new_h, new_w, new_c]) 59 | 60 | def get_config(self): 61 | config = super(PixelShuffle, self).get_config() 62 | config['block_size'] = self.block_size 63 | return config 64 | -------------------------------------------------------------------------------- /tools/learning_rate.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def step_decay(lr=3e-4, max_epochs=100, warmup=False): 5 | """ 6 | step decay. 7 | :param lr: initial lr 8 | :param max_epochs: max epochs 9 | :param warmup: warm up or not 10 | :return: current lr 11 | """ 12 | drop = 0.1 13 | max_epochs = max_epochs - 5 if warmup else max_epochs 14 | 15 | def decay(epoch): 16 | lrate = lr * np.power(drop, np.floor((1 + epoch) / max_epochs)) 17 | return lrate 18 | 19 | return decay 20 | 21 | 22 | def poly_decay(lr=3e-4, max_epochs=100, warmup=False): 23 | """ 24 | poly decay. 25 | :param lr: initial lr 26 | :param max_epochs: max epochs 27 | :param warmup: warm up or not 28 | :return: current lr 29 | """ 30 | max_epochs = max_epochs - 5 if warmup else max_epochs 31 | 32 | def decay(epoch): 33 | lrate = lr * (1 - np.power(epoch / max_epochs, 0.9)) 34 | return lrate 35 | 36 | return decay 37 | 38 | 39 | def cosine_decay(max_epochs, max_lr, min_lr=1e-7, warmup=False): 40 | """ 41 | cosine annealing scheduler. 42 | :param max_epochs: max epochs 43 | :param max_lr: max lr 44 | :param min_lr: min lr 45 | :param warmup: warm up or not 46 | :return: current lr 47 | """ 48 | max_epochs = max_epochs - 5 if warmup else max_epochs 49 | 50 | def decay(epoch): 51 | lrate = min_lr + (max_lr - min_lr) * ( 52 | 1 + np.cos(np.pi * epoch / max_epochs)) / 2 53 | return lrate 54 | 55 | return decay 56 | 57 | def lr_decays_func(lr_decay, learning_rate, num_epochs, lr_warmup ): 58 | 59 | func = { 60 | 'step_decay': step_decay(learning_rate, num_epochs - 5 if lr_warmup else num_epochs,warmup=lr_warmup), 61 | 'poly_decay': poly_decay(learning_rate, num_epochs - 5 if lr_warmup else num_epochs,warmup=lr_warmup), 62 | 'cosine_decay': cosine_decay(num_epochs - 5 if lr_warmup else num_epochs, learning_rate, warmup=lr_warmup)} 63 | 64 | return func[lr_decay] 65 | -------------------------------------------------------------------------------- /tools/metrics.py: -------------------------------------------------------------------------------- 1 | import tensorflow.keras.backend as K 2 | import tensorflow as tf 3 | 4 | 5 | def f1_score(y_true, y_pred): 6 | def recall(y_true, y_pred): 7 | """Recall metric. 8 | Only computes a batch-wise average of recall. 9 | Computes the recall, a metric for multi-label classification of 10 | how many relevant items are selected. 11 | """ 12 | true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1))) 13 | possible_positives = K.sum(K.round(K.clip(y_true, 0, 1))) 14 | recall = true_positives / (possible_positives + K.epsilon()) 15 | return recall 16 | 17 | def precision(y_true, y_pred): 18 | """Precision metric. 19 | Only computes a batch-wise average of precision. 20 | Computes the precision, a metric for multi-label classification of 21 | how many selected items are relevant. 22 | """ 23 | true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1))) 24 | predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1))) 25 | precision = true_positives / (predicted_positives + K.epsilon()) 26 | return precision 27 | 28 | precision = precision(y_true, y_pred) 29 | recall = recall(y_true, y_pred) 30 | 31 | return 2 * ((precision * recall) / (precision + recall + K.epsilon())) 32 | 33 | 34 | class MeanIoU(tf.keras.metrics.MeanIoU): 35 | def __init__(self, num_labels, name = 'mean_iou'): 36 | super(MeanIoU, self).__init__(num_labels, name= name) 37 | 38 | def update_state(self, y_true, y_pred, sample_weight=None): 39 | # sparse code 40 | y_true = tf.argmax(y_true, axis=-1) 41 | y_pred = tf.argmax(y_pred, axis=-1) 42 | return super(MeanIoU, self).update_state(y_true, y_pred, sample_weight) 43 | -------------------------------------------------------------------------------- /tools/split_val_data_from_train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | random.seed(1) 4 | import shutil 5 | from tqdm import tqdm 6 | 7 | 8 | def check_folder(dir): 9 | if not os.path.exists(dir): 10 | os.makedirs(dir) 11 | 12 | def move_file(old_path, new_path): 13 | print(old_path) 14 | print(new_path) 15 | filelist = os.listdir(old_path) #列出该目录下的所有文件,listdir返回的文件列表是不包含路径的。 16 | print(filelist) 17 | for file in filelist: 18 | src = os.path.join(old_path, file) 19 | dst = os.path.join(new_path, file) 20 | shutil.move(src, dst) 21 | 22 | # get val data from train data 23 | def split(dir): 24 | data =[] 25 | for x in os.listdir(dir): 26 | if x.endswith('tif'): 27 | data.append(os.path.join(dir,x)) 28 | return data 29 | 30 | if __name__=='__main__': 31 | train_dir = '../data/train/image' 32 | train_label_dir = '../data/train/label' 33 | val_dir = '../data/val/image' 34 | val_label_dir = '../data/val/label' 35 | check_folder(val_label_dir) 36 | check_folder(val_dir) 37 | val_ratio = 0.1 38 | res = split(train_dir) 39 | random.shuffle(res) 40 | print(len(res)) 41 | print(res) 42 | for i in tqdm(range(int(len(res) * val_ratio))): 43 | shutil.move(res[i],res[i].replace(train_dir,val_dir)) # tif images 44 | # png images 45 | shutil.move(res[i].replace(train_dir,train_label_dir).replace('tif','png'),res[i].replace(train_dir,val_label_dir).replace('tif','png')) 46 | 47 | 48 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | from tools import dataloader 2 | from tensorflow.keras import optimizers 3 | from tools.callbacks import LearningRateScheduler 4 | from tools.learning_rate import lr_decays_func 5 | from tools.metrics import MeanIoU 6 | from tensorflow.keras.callbacks import ModelCheckpoint,ReduceLROnPlateau 7 | from builders import builder 8 | import tensorflow as tf 9 | import argparse 10 | import os 11 | from config import Config 12 | 13 | os.environ["CUDA_VISIBLE_DEVICES"] = "1" 14 | cfg = Config('train') 15 | 16 | # config = tf.compat.v1.ConfigProto() 17 | # config.gpu_options.allow_growth = True 18 | # tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config)) 19 | 20 | def args_parse(): 21 | parser = argparse.ArgumentParser() 22 | parser.add_argument('--model', help='Choose the semantic segmentation methods.', type=str, default='DeepLabV3Plus') 23 | parser.add_argument('--backBone', help='Choose the backbone model.', type=str, default='DenseNet121') 24 | parser.add_argument('--num_epochs', help='num_epochs', type=int, default=cfg.epoch) 25 | 26 | parser.add_argument('--weights', help='The path of weights to be loaded.', type=str, default='weights') 27 | parser.add_argument('--lr_scheduler', help='The strategy to schedule learning rate.', type=str, 28 | default='cosine_decay', 29 | choices=['step_decay', 'poly_decay', 'cosine_decay']) 30 | parser.add_argument('--lr_warmup', help='Whether to use lr warm up.', type=bool, default=False) 31 | parser.add_argument('--learning_rate', help='learning_rate.', type=float, default=cfg.lr) 32 | 33 | args = parser.parse_args() 34 | return args 35 | 36 | 37 | def train(args): 38 | filepath = "weights-{epoch:03d}-{val_loss:.4f}-{val_mean_iou:.4f}.h5" 39 | weights_dir = os.path.join(args.weights, args.backBone + '_' + args.model) 40 | cfg.check_folder(weights_dir) 41 | model_weights = os.path.join(weights_dir, filepath) 42 | 43 | # build the model 44 | model, base_model = builder(cfg.n_classes, (256, 256), args.model, args.backBone) 45 | model.summary() 46 | 47 | # compile the model 48 | sgd = optimizers.SGD(lr=0.0001, momentum=0.9) 49 | adam = optimizers.Adam(lr=cfg.lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8) 50 | model.compile(optimizer=sgd, loss='categorical_crossentropy', metrics=[MeanIoU(cfg.n_classes)]) 51 | 52 | # checkpoint setting 53 | model_checkpoint = ModelCheckpoint(model_weights, monitor='val_loss', save_best_only=True, mode='auto') 54 | 55 | # learning rate scheduler setting 56 | lr_decay = lr_decays_func(args.lr_scheduler, args.learning_rate, args.num_epochs, args.lr_warmup) 57 | learning_rate_scheduler = LearningRateScheduler(lr_decay, args.learning_rate, args.lr_warmup, cfg.steps_per_epoch, 58 | num_epochs=args.num_epochs, verbose=1) 59 | Reduce_LR = ReduceLROnPlateau(monitor='val_mean_iou', mode='max', patience=2, verbose=1, factor=0.2, min_lr=1e-7) 60 | # callbacks = [model_checkpoint] 61 | callbacks = [model_checkpoint, Reduce_LR] 62 | 63 | 64 | # training... 65 | train_set = dataloader.train_data_generator(cfg.train_data_path, cfg.train_label_path, cfg.batch_size, 66 | cfg.n_classes, cfg.data_augment) 67 | val_set = dataloader.val_data_generator(cfg.val_data_path, cfg.val_label_path, cfg.batch_size, cfg.n_classes) 68 | 69 | start_epoch = 0 70 | if os.path.exists(weights_dir) and os.listdir(weights_dir): 71 | a = sorted(file for file in os.listdir(weights_dir)) 72 | model.load_weights(weights_dir + '/' + a[-1], by_name=True) 73 | # if load success, output info 74 | print('loaded :' + '-' * 8 + weights_dir + '/' + a[-1]) 75 | start_epoch = int(a[-1][8:11]) 76 | 77 | for layer in model.layers: 78 | layer.trainable=False 79 | for i in range(-1, -27, -1): 80 | model.layers[i].trainable = True 81 | model.summary() 82 | 83 | print("start_epoch: ", start_epoch) 84 | if start_epoch == 0: 85 | backbone_pretrained_path = "backBone_pretrained_weights/" + 'densenet121_weights_tf_dim_ordering_tf_kernels_notop.h5' 86 | model.load_weights(backbone_pretrained_path, by_name=True) 87 | print(f"loaded : {backbone_pretrained_path}" ) 88 | 89 | print(len(model.layers)) 90 | for layer in model.layers: 91 | layer.trainable=False 92 | for i in range(-1, -27, -1): 93 | model.layers[i].trainable = True 94 | model.summary() 95 | 96 | model.fit(train_set, 97 | steps_per_epoch=cfg.steps_per_epoch, 98 | epochs=args.num_epochs, 99 | callbacks=callbacks, 100 | validation_data=val_set, 101 | validation_steps=cfg.validation_steps, 102 | max_queue_size= cfg.batch_size, 103 | initial_epoch=start_epoch) 104 | 105 | if __name__ == '__main__': 106 | args = args_parse() 107 | train(args) 108 | 109 | -------------------------------------------------------------------------------- /train_custom.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import tensorflow as tf 4 | from tensorflow.keras.layers import Input 5 | from tensorflow.keras.models import Model 6 | from tensorflow.keras import optimizers 7 | from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau 8 | from tools import dataloader 9 | from config import Config 10 | from tools.learning_rate import lr_decays_func 11 | from tools.metrics import MeanIoU 12 | from tools.callbacks import LearningRateScheduler 13 | from models.Custom_net import Net 14 | 15 | 16 | os.environ["CUDA_VISIBLE_DEVICES"] = "1" 17 | cfg = Config(flag='train') 18 | 19 | config = tf.compat.v1.ConfigProto() 20 | config.gpu_options.allow_growth = True 21 | tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config)) 22 | 23 | def args_parse(): 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--num_epochs', help='num_epochs', type=int, default=cfg.epoch) 26 | parser.add_argument('--weights', help='The path of weights to be loaded.', type=str, default=cfg.weight_path) 27 | parser.add_argument('--lr_scheduler', help='The strategy to schedule learning rate.', type=str, 28 | default='cosine_decay', 29 | choices=['step_decay', 'poly_decay', 'cosine_decay']) 30 | parser.add_argument('--lr_warmup', help='Whether to use lr warm up.', type=bool, default=False) 31 | parser.add_argument('--learning_rate', help='learning_rate.', type=float, default=cfg.lr) 32 | 33 | args = parser.parse_args() 34 | return args 35 | 36 | def train(args): 37 | 38 | filepath = "weights-{epoch:03d}-{val_loss:.4f}-{val_mean_iou:.4f}.h5" 39 | model_weights = os.path.join(args.weights, filepath) 40 | 41 | input = Input(shape=(None, None, 3)) 42 | pred = Net(input, [10,12,16,14], 32, 0.5, cfg.n_classes) 43 | model = Model(input, pred) 44 | model.summary() 45 | 46 | # compile the model 47 | sgd = optimizers.SGD(lr=cfg.lr, momentum= 0.9) 48 | # adam = optimizers.Adam(lr=cfg.lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8) 49 | # nadam = optimizers.Nadam(lr=cfg.lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8) 50 | model.compile(optimizer=sgd, loss='categorical_crossentropy',metrics=[MeanIoU(cfg.n_classes)]) 51 | 52 | modelcheck = ModelCheckpoint(model_weights, monitor='val_loss', save_best_only=True, mode='auto') 53 | # learning rate scheduler setting 54 | lr_decay = lr_decays_func(args.lr_scheduler, args.learning_rate,args.num_epochs, args.lr_warmup) 55 | learning_rate_scheduler = LearningRateScheduler(lr_decay, args.learning_rate, args.lr_warmup, cfg.steps_per_epoch, 56 | num_epochs=args.num_epochs, verbose=1) 57 | 58 | Reduce_LR = ReduceLROnPlateau(monitor='val_mean_iou', mode='max', patience=2, verbose=1, factor=0.2, min_lr=1e-7) 59 | 60 | # callable = [modelcheck, learning_rate_scheduler] 61 | callable = [modelcheck, Reduce_LR] 62 | # callable = [modelcheck] 63 | 64 | 65 | train_set = dataloader.train_data_generator( 66 | cfg.train_data_path, cfg.train_label_path, cfg.batch_size, cfg.n_classes, cfg.data_augment) 67 | val_set = dataloader.val_data_generator( 68 | cfg.val_data_path, cfg.val_label_path, cfg.batch_size, cfg.n_classes) 69 | 70 | start_epoch = 0 71 | if os.path.exists(args.weights) and os.listdir(args.weights): 72 | a = sorted(file for file in os.listdir(args.weights)) 73 | model.load_weights(args.weights + '/' + a[-1], by_name=True) 74 | # if load success, output info 75 | print('loaded :' + '-'*8 + args.weights + '/' + a[-1]) 76 | start_epoch = int(a[-1][8:11]) 77 | 78 | H = model.fit(x=train_set, steps_per_epoch=cfg.steps_per_epoch, epochs=cfg.epoch, 79 | verbose=1, validation_data=val_set, validation_steps=cfg.validation_steps, 80 | callbacks=callable, max_queue_size= 2*cfg.batch_size, initial_epoch=start_epoch) 81 | 82 | if __name__ == '__main__': 83 | args = args_parse() 84 | cfg.check_folder(args.weights) 85 | train(args) 86 | -------------------------------------------------------------------------------- /train_unet.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import tensorflow as tf 4 | from tensorflow.keras.layers import Input 5 | from tensorflow.keras.models import Model 6 | from tensorflow.keras import optimizers 7 | from tensorflow.keras.callbacks import ModelCheckpoint 8 | from tools import dataloader 9 | from config import Config 10 | from tools.learning_rate import lr_decays_func 11 | from tools.metrics import MeanIoU 12 | from tools.callbacks import LearningRateScheduler 13 | from models.unet import Net 14 | 15 | 16 | os.environ["CUDA_VISIBLE_DEVICES"] = "1" 17 | cfg = Config(flag='train') 18 | 19 | config = tf.compat.v1.ConfigProto() 20 | config.gpu_options.allow_growth = True 21 | tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config)) 22 | 23 | def args_parse(): 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--num_epochs', help='num_epochs', type=int, default=cfg.epoch) 26 | parser.add_argument('--weights', help='The path of weights to be loaded.', type=str, default=cfg.weight_path) 27 | parser.add_argument('--lr_scheduler', help='The strategy to schedule learning rate.', type=str, 28 | default='cosine_decay', 29 | choices=['step_decay', 'poly_decay', 'cosine_decay']) 30 | parser.add_argument('--lr_warmup', help='Whether to use lr warm up.', type=bool, default=False) 31 | parser.add_argument('--learning_rate', help='learning_rate.', type=float, default=cfg.lr) 32 | 33 | args = parser.parse_args() 34 | return args 35 | 36 | def train(args): 37 | 38 | filepath = "weights-{epoch:03d}-{val_loss:.4f}-{val_mean_iou:.4f}.h5" 39 | model_weights = os.path.join(args.weights, filepath) 40 | 41 | input = Input(shape=(None, None, 3)) 42 | pred = Net(cfg.n_classes, input) 43 | model = Model(input, pred) 44 | model.summary() 45 | 46 | # compile the model 47 | # sgd = optimizers.SGD(lr=cfg.lr, momentum= 0.9) 48 | #adam = optimizers.Adam(lr=cfg.lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8) 49 | nadam = optimizers.Nadam(lr=cfg.lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8) 50 | model.compile(optimizer=nadam, loss='categorical_crossentropy',metrics=[MeanIoU(cfg.n_classes)]) 51 | 52 | modelcheck = ModelCheckpoint(model_weights, monitor='val_loss', save_best_only=True, mode='auto') 53 | # learning rate scheduler setting 54 | lr_decay = lr_decays_func(args.lr_scheduler, args.learning_rate,args.num_epochs, args.lr_warmup) 55 | learning_rate_scheduler = LearningRateScheduler(lr_decay, args.learning_rate, args.lr_warmup, cfg.steps_per_epoch, 56 | num_epochs=args.num_epochs, verbose=1) 57 | # callable = [modelcheck, learning_rate_scheduler] 58 | callable = [modelcheck] # adam optimizer 59 | 60 | train_set = dataloader.train_data_generator( 61 | cfg.train_data_path, cfg.train_label_path, cfg.batch_size, cfg.n_classes, cfg.data_augment) 62 | val_set = dataloader.val_data_generator( 63 | cfg.val_data_path, cfg.val_label_path, cfg.batch_size, cfg.n_classes) 64 | 65 | start_epoch = 0 66 | if os.path.exists(args.weights) and os.listdir(args.weights): 67 | a = sorted(file for file in os.listdir(args.weights)) 68 | model.load_weights(args.weights + '/' + a[-1], by_name=True) 69 | # if load success, output info 70 | print('loaded :' + '-'*8 + args.weights + '/' + a[-1]) 71 | start_epoch = int(a[-1][8:11]) 72 | 73 | H = model.fit(x=train_set, steps_per_epoch=cfg.steps_per_epoch, epochs=cfg.epoch, 74 | verbose=1, validation_data=val_set, validation_steps=cfg.validation_steps, 75 | callbacks=callable, max_queue_size= cfg.batch_size, initial_epoch=start_epoch) 76 | 77 | if __name__ == '__main__': 78 | args = args_parse() 79 | cfg.check_folder(args.weights) 80 | train(args) 81 | --------------------------------------------------------------------------------