├── HRNet v2_Train.py ├── LICENSE ├── Model.py ├── README.md ├── logs └── HRNet V2 │ ├── train │ └── events.out.tfevents.1559205568.CXS │ └── val │ └── events.out.tfevents.1559205573.CXS └── utils ├── HighResolutionModule.py ├── Models_Block.py ├── Models_Config.py ├── __pycache__ ├── HighResolutionModule.cpython-35.pyc ├── Models_Block.cpython-35.pyc ├── Models_Config.cpython-35.pyc └── data_utils.cpython-35.pyc ├── color_utils.py └── data_utils.py /HRNet v2_Train.py: -------------------------------------------------------------------------------- 1 | from Model import HighResolutionNet 2 | from utils.data_utils import DataSet 3 | from sklearn.metrics import accuracy_score 4 | from sklearn.model_selection import train_test_split 5 | 6 | import tensorflow as tf 7 | import pandas as pd 8 | import numpy as np 9 | import os 10 | 11 | class args: 12 | batch_size = 2 13 | lr = 0.0002 14 | # display = 500 15 | display = 10 16 | weight_decay = 0.00001 17 | num_class = 6 18 | model_name = 'HRNet V2' 19 | 20 | data_path_df = pd.read_csv('dataset/path_list.csv') 21 | 22 | tf.logging.set_verbosity(tf.logging.INFO) 23 | #将 TensorFlow 日志信息输出到屏幕.TensorFlow将输出与该级别相对应的所有日志消息以及更高程度严重性的所有级别的日志信息。 24 | train_path, val_path = train_test_split(data_path_df, test_size=0.25, shuffle=True) 25 | #train_test_split是交叉验证中常用的函数，功能是从样本中随机的按比例选取train_data和test_data 26 | dataset_tr = DataSet(image_path=train_path['image'].values, label_path=train_path['label'].values) 27 | dataset_val = DataSet(image_path=val_path['image'].values, label_path=val_path['label'].values) 28 | 29 | model = HighResolutionNet(args.num_class) 30 | 31 | image = tf.placeholder(tf.float32, [None, 256, 256, 3], name='input_x') 32 | label = tf.placeholder(tf.int32, [None, 256, 256]) 33 | lr = tf.placeholder(tf.float32,) 34 | 35 | logits = model.forword(image) 36 | print(logits) 37 | predicts = tf.argmax(logits, axis=-1, name='predicts') 38 | print(predicts) 39 | 40 | # cross_entropy 41 | cross_entropy = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=label)) 42 | # l2_norm l2正则化 43 | l2_loss = args.weight_decay * tf.add_n( 44 | [tf.nn.l2_loss(tf.cast(v, tf.float32)) for v in tf.trainable_variables()]) 45 | 46 | optimizer = tf.train.AdamOptimizer(learning_rate=lr) 47 | loss = cross_entropy + l2_loss 48 | 49 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 50 | 51 | with tf.control_dependencies(update_ops): 52 | train_op = optimizer.minimize(loss=loss) 53 | 54 | total_acc_tr, total_acc_val, total_loss_tr, total_loss_val, total_l2_loss = 0, 0, 0, 0, 0 55 | # 我们要保存所有的参数 56 | saver = tf.train.Saver(tf.all_variables()) 57 | 58 | best_val_acc = 0.60 59 | 60 | with tf.Session() as sess: 61 | sess.run(tf.local_variables_initializer()) 62 | sess.run(tf.global_variables_initializer()) 63 | # finetune resnet_v2_50参数 64 | # restorer.restore(sess, 'ckpts/resnet_v2_50/resnet_v2_50.ckpt') 65 | 66 | log_path = 'logs/%s/' % args.model_name 67 | model_path = 'ckpts/%s/' % args.model_name 68 | 69 | if not os.path.exists(model_path): os.makedirs(model_path) 70 | if not os.path.exists('./logs'): os.makedirs('./logs') 71 | if not os.path.exists(log_path): os.makedirs(log_path) 72 | 73 | train_summary_writer = tf.summary.FileWriter('%s/train' % log_path, sess.graph) 74 | val_summary_writer = tf.summary.FileWriter('%s/val' % log_path, sess.graph) 75 | 76 | learning_rate = args.lr 77 | 78 | ckpt_dir = tf.train.get_checkpoint_state('ckpts/deeplab_v3/') 79 | if ckpt_dir and ckpt_dir.model_checkpoint_path: 80 | saver.restore(sess, ckpt_dir.model_checkpoint_path) 81 | #learning_rate = args.lr / 10 82 | else: 83 | print('not find ckpts') 84 | 85 | for step in range(1, 9550): 86 | if step == 5000 or step == 12000: 87 | learning_rate = learning_rate / 10 88 | x_tr, y_tr = dataset_tr.next_batch(args.batch_size) 89 | x_val, y_val = dataset_val.next_batch(args.batch_size) 90 | loss_tr, l2_loss_tr, predicts_tr, _ = sess.run( 91 | fetches=[cross_entropy, l2_loss, predicts, train_op], 92 | feed_dict={ 93 | image: x_tr, 94 | label: y_tr, 95 | lr: learning_rate}) 96 | loss_val, predicts_val = sess.run( 97 | fetches=[cross_entropy, predicts], 98 | feed_dict={ 99 | image: x_val, 100 | label: y_val,}) 101 | #计数 102 | total_loss_tr += loss_tr 103 | total_loss_val += loss_val 104 | total_l2_loss += l2_loss_tr 105 | 106 | acc_tr = accuracy_score(np.reshape(y_tr, [-1]), np.reshape(predicts_tr, [-1])) 107 | acc_val = accuracy_score(np.reshape(y_val, [-1]), np.reshape(predicts_val, [-1])) 108 | total_acc_tr += acc_tr 109 | total_acc_val += acc_val 110 | 111 | # 每隔多少步打印日志 112 | if step % args.display == 0: 113 | tf.logging.info("Iter:%d, lr:%.6f, loss_tr:%.4f, acc_tr:%.4f, loss_val:%.4f, acc_val:%.4f" % 114 | (step, 115 | learning_rate, 116 | total_loss_tr / args.display, 117 | total_acc_tr / args.display, 118 | total_loss_val / args.display, 119 | total_acc_val / args.display)) 120 | 121 | train_summary = tf.Summary( 122 | value=[tf.Summary.Value(tag='loss', simple_value=total_loss_tr / args.display), 123 | tf.Summary.Value(tag='accuracy', simple_value=total_acc_tr / args.display), 124 | tf.Summary.Value(tag='l2_loss', simple_value=total_l2_loss / args.display)]) 125 | 126 | val_summary = tf.Summary( 127 | value=[tf.Summary.Value(tag='loss', simple_value=total_loss_val / args.display), 128 | tf.Summary.Value(tag='accuracy', simple_value=total_acc_val / args.display)]) 129 | # 记录summary 130 | train_summary_writer.add_summary(train_summary, step) 131 | train_summary_writer.flush() 132 | val_summary_writer.add_summary(val_summary, step) 133 | val_summary_writer.flush() 134 | 135 | 136 | # 保存模型 137 | if (total_acc_val / args.display) > best_val_acc: 138 | # 保存acc最高的模型 139 | saver.save(sess, model_path + 'model.ckpt') 140 | best_val_acc = total_acc_val / args.display 141 | tf.logging.info('Save model successfully to "%s"!' % (model_path + 'model.ckpt')) 142 | 143 | if(step>9450): 144 | if(total_acc_val / args.display>temp_best): 145 | saver.save(sess, 'ckpts/'+ 'model.ckpt') 146 | temp_best = total_acc_val / args.display 147 | tf.logging.info('Save model successfully to "%s"!' % (model_path + 'model.ckpt')) 148 | # 累计归零 149 | total_acc_tr, total_acc_val, total_loss_tr, total_loss_val, total_l2_loss = 0, 0, 0, 0, 0 150 | 151 | 152 | 153 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 AI-Chen 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Model.py: -------------------------------------------------------------------------------- 1 | """ 2 | Related papers: 3 | High-Resolution Representations for Labeling Pixels and Regions.higharXiv:1904.04514v1 [cs.CV] 9 Apr 2019 4 | """ 5 | from __future__ import absolute_import 6 | from __future__ import division 7 | from __future__ import print_function 8 | from tensorflow.python.training import moving_averages 9 | 10 | import tensorflow as tf 11 | import logging 12 | import cv2 13 | import numpy as np 14 | from utils import Models_Block as block_Model 15 | from utils import Models_Config as config 16 | from utils import HighResolutionModule as HM 17 | 18 | 19 | 20 | blocks_dict = { 21 | 'BASIC': block_Model.BasicBlock, 22 | 'BOTTLENECK': block_Model.Bottleneck 23 | } 24 | 25 | class HighResolutionNet(): 26 | def __init__(self,num_class): 27 | self._is_training = tf.cast(True, tf.bool) 28 | self.layer = [] 29 | self.nua_class = num_class 30 | def forword(self, x): 31 | #..............................stage1..................................# 32 | with tf.variable_scope('stage1'): 33 | with tf.variable_scope('stage1_conv1'): 34 | temp = block_Model.conv(x, 3, 64, 'stage1_conv1', stride=2) 35 | temp = block_Model.batch_norm(temp,self._is_training) 36 | with tf.variable_scope('stage1_conv2'): 37 | temp = block_Model.conv(temp, 3, 64, 'stage1_conv2', stride=2) 38 | temp = block_Model.batch_norm(temp, self._is_training) 39 | layer1 = self._make_layer(temp,block_Model.Bottleneck, 64, 64, 4) 40 | self.layer.append(layer1) 41 | #......................................................................# 42 | 43 | # ..............................stage2..................................# 44 | with tf.variable_scope('stage2'): 45 | num_channels = config.HIGH_RESOLUTION_NET[1]['NUM_CHANNELS'] 46 | block = blocks_dict[config.HIGH_RESOLUTION_NET[1]['BLOCK']] 47 | num_channels = [num_channels[i] * 1 for i in range(len(num_channels))] 48 | self.layer = self._make_transition_layer('stage2', self.layer,[256], num_channels) 49 | self.layer, pre_stage_channels = self._make_stage('stage2',self.layer,config.HIGH_RESOLUTION_NET[1], 50 | num_channels) 51 | # ......................................................................# 52 | 53 | # ..............................stage3..................................# 54 | with tf.variable_scope('stage3'): 55 | num_channels = config.HIGH_RESOLUTION_NET[2]['NUM_CHANNELS'] 56 | block = blocks_dict[config.HIGH_RESOLUTION_NET[2]['BLOCK']] 57 | num_channels = [num_channels[i] * 1 for i in range(len(num_channels))] 58 | self.layer = self._make_transition_layer('stage3', self.layer,pre_stage_channels, num_channels) 59 | self.layer, pre_stage_channels = self._make_stage('stage3', self.layer, config.HIGH_RESOLUTION_NET[2], 60 | num_channels) 61 | # ......................................................................# 62 | 63 | # ..............................stage4..................................# 64 | with tf.variable_scope('stage4'): 65 | num_channels = config.HIGH_RESOLUTION_NET[3]['NUM_CHANNELS'] 66 | block = blocks_dict[config.HIGH_RESOLUTION_NET[3]['BLOCK']] 67 | num_channels = [num_channels[i] * 1 for i in range(len(num_channels))] 68 | self.layer = self._make_transition_layer('stage4', self.layer, pre_stage_channels, num_channels) 69 | self.layer, pre_stage_channels = self._make_stage('stage4', self.layer, config.HIGH_RESOLUTION_NET[3], 70 | num_channels) 71 | # ......................................................................# 72 | 73 | # ..............................concat..................................# 74 | with tf.variable_scope('concat'): 75 | for i in range(len(self.layer)): 76 | self.layer[i] = tf.image.resize_bilinear(self.layer[i], [256,256]) 77 | feature_map = tf.concat((self.layer[0], self.layer[1], self.layer[2], self.layer[3]), axis=3) 78 | feature_map = block_Model.conv(feature_map,1,self.nua_class,'concat_conv',1) 79 | return feature_map 80 | 81 | # ......................................................................# 82 | 83 | def _make_layer(self, x, block, inplanes, planes, blocks, stride=1): 84 | '''作用：make layer1 85 | 在stage1中调用''' 86 | downsample = None 87 | if stride != 1 or inplanes != planes * 4: 88 | downsample = block_Model.conv 89 | 90 | x = block(x, is_training=self._is_training, block_name='layer1_Botteleneck0', outplanes=planes, 91 | stride=stride,downsample=downsample) 92 | for i in range(1, blocks): 93 | x = block(x, is_training=self._is_training, block_name='layer1_Botteleneck'+ str(i), outplanes=planes) 94 | return x 95 | 96 | def _make_transition_layer(self, stage,layer, num_channels_pre_layer, num_channels_cur_layer): 97 | 98 | num_branches_cur = len(num_channels_cur_layer) 99 | num_branches_pre = len(num_channels_pre_layer) 100 | 101 | transition_layers = [] 102 | for i in range(num_branches_cur): 103 | if i < num_branches_pre: 104 | with tf.variable_scope(stage + '_transition_layers_1' + str(i)): 105 | if num_channels_cur_layer[i] != num_channels_pre_layer[i]: 106 | layer[i] = block_Model.conv(layer[i],3,num_channels_cur_layer[i],'transition_layers',1) 107 | layer[i] = block_Model.batch_norm(layer[i],self._is_training) 108 | transition_layers.append(layer[i]) 109 | else: 110 | transition_layers.append(layer[i]) 111 | else: 112 | for j in range(i + 1 - num_branches_pre): 113 | inchannels = num_channels_pre_layer[-1] 114 | outchannels = num_channels_cur_layer[i] \ 115 | if j == i - num_branches_pre else inchannels 116 | with tf.variable_scope(stage + 'stage1_transition_layers_2'+str(j)): 117 | temp = block_Model.conv(layer[i-1], 3, outchannels, 'transition_layers', 2) 118 | temp = block_Model.batch_norm(temp,self._is_training) 119 | transition_layers.append(temp) 120 | return transition_layers 121 | 122 | def _make_stage(self, stage,layer,layer_config, num_inchannels, 123 | multi_scale_output=True): 124 | num_modules = layer_config['NUM_MODULES'] 125 | num_branches = layer_config['NUM_BRANCHES'] 126 | num_blocks = layer_config['NUM_BLOCKS'] 127 | num_channels = layer_config['NUM_CHANNELS'] 128 | block = blocks_dict[layer_config['BLOCK']] 129 | fuse_method = layer_config['FUSE_METHOD'] 130 | 131 | make_stage = [] 132 | channel = [] 133 | 134 | for i in range(num_modules): 135 | # multi_scale_output is only used last module 136 | if not multi_scale_output and i == num_modules - 1: 137 | reset_multi_scale_output = False 138 | else: 139 | reset_multi_scale_output = True 140 | hm= HM.HighResolutionModule('stage2', self._is_training,num_branches,block,num_blocks,num_inchannels, 141 | num_channels,fuse_method,reset_multi_scale_output) 142 | make_stage,channel = hm.forword(layer,num_branches,block,num_blocks,num_inchannels, 143 | num_channels,fuse_method,reset_multi_scale_output) 144 | return make_stage,channel 145 | 146 | # img1 = cv2.imread('dataset/top_mosaic_09cm_area1.tif',1) 147 | # img_pre = img1[0:256,0:256] 148 | # inputX = [] 149 | # inputX.append(img_pre) 150 | # 151 | # 152 | # image = tf.placeholder(tf.float32, [None, 256, 256, 3], name='input_x') 153 | # 154 | # 155 | # HRNet = HighResolutionNet() 156 | # 157 | # pred = HRNet.forword(image) 158 | # with tf.Session() as sess: 159 | # sess.run(tf.local_variables_initializer()) 160 | # sess.run(tf.global_variables_initializer()) 161 | # 162 | # sess.run(fetches=[pred], 163 | # feed_dict={ 164 | # image: inputX, 165 | # }) 166 | 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # HRNet-V2进行语义分割的tensorflow实现 2 | ## 文件简介 3 | + **utils文件夹**中包含*color_utils.py、data_utils.py、HighResolutionModule.py、Models_Block.py、Models_Config.py*. 4 | 1. *color_utils.py、data_utils.py*是进行数据预处理的功能函数； 5 | 2. *HighResolutionModule.py*是HRNet v2中的进行并行卷积融合的功能模块； 6 | 3. *Models_Block.py*是对卷积、批归一化、BasicBlock、Bottleneck的封装； 7 | 4. *Models_Config.py*包含了HRNet中每一个阶段的超参数设置； 8 | + **__pycache__文件夹**是PyCharm编译时产生的文件. 9 | + **Model.py**是HRNet v2模型的实现 10 | + **HRNet v2_Train.py**是运用模型进行训练的文件。 11 | ## Getting started 12 | 1. 数据集使用已经经过预处理适合本模型进行训练的ISPRS 2D Vaihingen数据集。下载地址[百度网盘](链接：https://pan.baidu.com/s/1RqjWXTZOCPO4cRkW6SVglA 13 | 提取码：cslj)。下载后解压至文件根目录； 14 | 2. 在根目录新建文件夹，命名为：ckpts。该文件夹用于存储cpkt文件； 15 | 3. 运行HRNet v2_Train.py即可开始模型的训练。 16 | ## Dependency 17 | + tensorflow==1.8.0 18 | + numpy==1.15.2 19 | + opencv3==3.1.0 20 | -------------------------------------------------------------------------------- /logs/HRNet V2/train/events.out.tfevents.1559205568.CXS: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Chen/HRNet-V2/341e5f57714f04f8eedb4962a49adbd60c1bb47e/logs/HRNet V2/train/events.out.tfevents.1559205568.CXS -------------------------------------------------------------------------------- /logs/HRNet V2/val/events.out.tfevents.1559205573.CXS: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Chen/HRNet-V2/341e5f57714f04f8eedb4962a49adbd60c1bb47e/logs/HRNet V2/val/events.out.tfevents.1559205573.CXS -------------------------------------------------------------------------------- /utils/HighResolutionModule.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | from tensorflow.python.training import moving_averages 5 | 6 | import tensorflow as tf 7 | import logging 8 | from utils import Models_Block as block_Model 9 | from utils import Models_Config as config 10 | 11 | logger = logging.getLogger(__name__) 12 | 13 | class HighResolutionModule(): 14 | def __init__(self,stage,_is_training,num_branches, blocks, num_blocks, num_inchannels, 15 | num_channels, fuse_method, multi_scale_output=True): 16 | self.stage = stage 17 | self._is_training = _is_training 18 | self.num_inchannels = num_inchannels 19 | self.fuse_method = fuse_method 20 | self.num_branches = num_branches 21 | self.multi_scale_output = multi_scale_output 22 | 23 | def forword(self,layer,num_branches, blocks, num_blocks, num_inchannels, 24 | num_channels, fuse_method, multi_scale_output=True): 25 | self._check_branches(num_branches,num_blocks, num_inchannels, num_channels) 26 | Result = self._make_branches(layer,num_branches, blocks, num_blocks, num_channels) 27 | Result = self._make_fuse_layers(Result) 28 | channel = [] 29 | for i in range(len(Result)): 30 | channel.append(Result[i].get_shape().as_list()[3]) 31 | return Result,channel 32 | 33 | 34 | 35 | def _check_branches(self, num_branches,num_blocks, 36 | num_inchannels, num_channels): 37 | if num_branches != len(num_blocks): 38 | error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format( 39 | num_branches, len(num_blocks)) 40 | logger.error(error_msg) 41 | raise ValueError(error_msg) 42 | 43 | if num_branches != len(num_channels): 44 | error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format( 45 | num_branches, len(num_channels)) 46 | logger.error(error_msg) 47 | raise ValueError(error_msg) 48 | 49 | if num_branches != len(num_inchannels): 50 | error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format( 51 | num_branches, len(num_inchannels)) 52 | logger.error(error_msg) 53 | raise ValueError(error_msg) 54 | 55 | def _make_branches(self, layer,num_branches, block, num_blocks, num_channels): 56 | branches = [] 57 | for i in range(num_branches): 58 | branches.append( 59 | self._make_one_branch(layer[i], i, block, num_blocks, num_channels)) 60 | return branches 61 | 62 | def _make_one_branch(self,x,branch_index, block, num_blocks, num_channels, 63 | stride=1): 64 | downsample = None 65 | if stride != 1 or self.num_inchannels[branch_index] != num_channels[branch_index] * 1: 66 | downsample = block_Model.conv 67 | x = block(x, is_training=self._is_training, block_name=self.stage + '_Botteleneck0' + str(branch_index), outplanes=num_channels[branch_index] * 1, 68 | stride=stride, downsample=downsample) 69 | for i in range(1, num_blocks[branch_index]): 70 | x = block(x, is_training=self._is_training, block_name=self.stage + 'Botteleneck' + str(i) + str(branch_index), outplanes=num_channels[branch_index]) 71 | return x 72 | 73 | def _make_fuse_layers(self,x): 74 | if self.num_branches == 1: 75 | return x 76 | num_branches = self.num_branches 77 | num_inchannels = self.num_inchannels 78 | fuse_layers = [] 79 | for i in range(num_branches if self.multi_scale_output else 1): 80 | temp_layer = [] 81 | for j in range(num_branches): 82 | if j > i: 83 | with tf.variable_scope(self.stage + '_fuse_layers' + str(i)+str(j)): 84 | temp = block_Model.conv(x[j], 1, num_inchannels[i], 'conv_', stride=1) 85 | temp = block_Model.batch_norm(temp, self._is_training) 86 | temp = tf.image.resize_bilinear(temp, x[i].get_shape().as_list()[1:3]) 87 | temp_layer.append(temp) 88 | elif j == i: 89 | temp_layer.append(x[j]) 90 | else: 91 | with tf.variable_scope(self.stage + '_fuse_layers' + str(i) + str(j)): 92 | temp = block_Model.conv(x[j], 3, num_inchannels[i], 'conv_', stride=2) 93 | temp = block_Model.batch_norm(temp, self._is_training) 94 | temp_layer.append(temp) 95 | for k in range(len(temp_layer)): 96 | fuse = 0 97 | fuse = fuse + temp_layer[k] 98 | fuse_layers.append(fuse) 99 | return fuse_layers 100 | -------------------------------------------------------------------------------- /utils/Models_Block.py: -------------------------------------------------------------------------------- 1 | """ 2 | Related papers: 3 | High-Resolution Representations for Labeling Pixels and Regions.higharXiv:1904.04514v1 [cs.CV] 9 Apr 2019 4 | """ 5 | from __future__ import absolute_import 6 | from __future__ import division 7 | from __future__ import print_function 8 | from tensorflow.python.training import moving_averages 9 | 10 | import tensorflow as tf 11 | 12 | 13 | 14 | 15 | 16 | def conv(x, kernel_size, out_planes, name,stride=1): 17 | """3x3 convolution with padding""" 18 | x_shape = x.get_shape().as_list() 19 | w = tf.get_variable(name = name+'b',shape=[kernel_size, kernel_size, x_shape[3], out_planes]) 20 | return tf.nn.conv2d(input= x, filter=w, padding='SAME',strides=[1, stride, stride, 1],name= name ) 21 | 22 | def batch_norm(x,training): 23 | x_shape = x.get_shape() 24 | params_shape = x_shape[-1:] 25 | 26 | axis = list(range(len(x_shape) - 1)) 27 | beta = tf.get_variable(name='beta', 28 | shape=params_shape, 29 | initializer=tf.zeros_initializer) 30 | 31 | gamma = tf.get_variable(name='gamma', 32 | shape=params_shape, 33 | initializer=tf.ones_initializer) 34 | 35 | moving_mean = tf.get_variable(name='moving_mean', 36 | shape=params_shape, 37 | initializer=tf.zeros_initializer, 38 | trainable=False) 39 | 40 | moving_variance = tf.get_variable(name='moving_variance', 41 | shape=params_shape, 42 | initializer=tf.ones_initializer, 43 | trainable=False) 44 | 45 | tf.add_to_collection('BN_MEAN_VARIANCE', moving_mean) 46 | tf.add_to_collection('BN_MEAN_VARIANCE', moving_variance) 47 | 48 | # These ops will only be preformed when training. 49 | mean, variance = tf.nn.moments(x, axis) 50 | update_moving_mean = moving_averages.assign_moving_average(moving_mean, 51 | mean, 52 | 0.99, 53 | name='MovingAvgMean') 54 | update_moving_variance = moving_averages.assign_moving_average(moving_variance, 55 | variance, 56 | 0.99, 57 | name='MovingAvgVariance') 58 | 59 | tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_moving_mean) 60 | tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_moving_variance) 61 | 62 | mean, variance = tf.cond( 63 | pred= training, 64 | true_fn=lambda: (mean, variance), 65 | false_fn=lambda: (moving_mean, moving_variance) 66 | ) 67 | x = tf.nn.batch_normalization(x, mean, variance, beta, gamma, 1e-3) 68 | return x 69 | 70 | 71 | def Bottleneck(x, is_training,block_name, outplanes, stride=1, downsample=None): 72 | residual = x 73 | with tf.variable_scope(block_name + '11_1'): 74 | name = block_name + 'conv1' 75 | out = conv(x, 1, outplanes, name=name, stride=stride) 76 | out = batch_norm(out,is_training) 77 | out = tf.nn.relu(out) 78 | 79 | with tf.variable_scope(block_name + '33_2'): 80 | name = block_name + 'conv2' 81 | out = conv(x, 3, outplanes, name=name, stride=stride) 82 | out = batch_norm(out,is_training) 83 | out = tf.nn.relu(out) 84 | 85 | with tf.variable_scope(block_name + '11_3'): 86 | name = block_name + 'conv3' 87 | out = conv(out, 1, outplanes * 4, name=name, stride=stride) 88 | out = batch_norm(out,is_training) 89 | 90 | if downsample is not None: 91 | with tf.variable_scope(block_name + 'downsample'): 92 | residual = downsample(x,1,outplanes * 4, 'stage_dawnSample',stride = stride) 93 | residual = batch_norm(residual, is_training) 94 | 95 | out = out + residual 96 | out = tf.nn.relu(out) 97 | return out 98 | 99 | 100 | def BasicBlock(x, is_training,block_name, outplanes, stride=1, downsample=None): 101 | residual = x 102 | with tf.variable_scope(block_name + '33_1'): 103 | name = block_name + 'conv1' 104 | out = conv(x, 3, outplanes, name=name, stride=stride) 105 | out = batch_norm(out,is_training) 106 | print(out.shape) 107 | out = tf.nn.relu(out) 108 | with tf.variable_scope(block_name + '33_2'): 109 | name = block_name + 'conv2' 110 | out = conv(out, 3, outplanes, name=name, stride=stride) 111 | out = batch_norm(out,is_training) 112 | 113 | if downsample is not None: 114 | with tf.variable_scope(block_name + 'downsample'): 115 | residual = downsample(x) 116 | residual = batch_norm(residual, is_training) 117 | 118 | out = out + residual 119 | out = tf.nn.relu(out) 120 | return out 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | -------------------------------------------------------------------------------- /utils/Models_Config.py: -------------------------------------------------------------------------------- 1 | HIGH_RESOLUTION_NET = [ 2 | {'PRETRAINED_LAYERS': ['*'], 3 | 'STEM_INPLANES':64, 4 | 'FINAL_CONV_KERNEL':1, 5 | 'WITH_HEAD':True}, 6 | 7 | {'NUM_MODULES': 1, 8 | 'NUM_BRANCHES': 2, 9 | 'NUM_BLOCKS': [4, 4], 10 | 'NUM_CHANNELS': [32, 64], 11 | 'BLOCK': 'BASIC', 12 | 'FUSE_METHOD': 'SUM'}, 13 | 14 | {'NUM_MODULES': 1, 15 | 'NUM_BRANCHES': 3, 16 | 'NUM_BLOCKS': [4, 4, 4], 17 | 'NUM_CHANNELS': [32, 64, 128], 18 | 'BLOCK': 'BASIC', 19 | 'FUSE_METHOD': 'SUM'}, 20 | 21 | {'NUM_MODULES': 1, 22 | 'NUM_BRANCHES': 4, 23 | 'NUM_BLOCKS': [4, 4, 4, 4], 24 | 'NUM_CHANNELS': [32, 64, 128, 256], 25 | 'BLOCK': 'BASIC', 26 | 'FUSE_METHOD': 'SUM'}, 27 | ] 28 | 29 | -------------------------------------------------------------------------------- /utils/__pycache__/HighResolutionModule.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Chen/HRNet-V2/341e5f57714f04f8eedb4962a49adbd60c1bb47e/utils/__pycache__/HighResolutionModule.cpython-35.pyc -------------------------------------------------------------------------------- /utils/__pycache__/Models_Block.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Chen/HRNet-V2/341e5f57714f04f8eedb4962a49adbd60c1bb47e/utils/__pycache__/Models_Block.cpython-35.pyc -------------------------------------------------------------------------------- /utils/__pycache__/Models_Config.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Chen/HRNet-V2/341e5f57714f04f8eedb4962a49adbd60c1bb47e/utils/__pycache__/Models_Config.cpython-35.pyc -------------------------------------------------------------------------------- /utils/__pycache__/data_utils.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Chen/HRNet-V2/341e5f57714f04f8eedb4962a49adbd60c1bb47e/utils/__pycache__/data_utils.cpython-35.pyc -------------------------------------------------------------------------------- /utils/color_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | # 给标签图上色 5 | 6 | def color_annotation(label_path, output_path): 7 | 8 | ''' 9 | 10 | 给class图上色 11 | 12 | ''' 13 | 14 | img = cv2.imread(label_path,0)########### 15 | 16 | color = np.ones([img.shape[0], img.shape[1], 3]) 17 | 18 | 19 | color[img==0] = [255, 255, 255] #不透水地表 20 | 21 | color[img==1] = [255, 0, 0] #建筑 22 | 23 | color[img==2] = [0, 255, 0] # 树木 24 | 25 | color[img==3] = [255, 255, 0] # 低层植被 26 | 27 | color[img==4] = [0, 255, 255] # 汽车 28 | 29 | color[img == 5] = [0, 0, 255] #水面 30 | 31 | 32 | 33 | cv2.imwrite(output_path,color) -------------------------------------------------------------------------------- /utils/data_utils.py: -------------------------------------------------------------------------------- 1 | import cv2, os 2 | import numpy as np 3 | from sklearn.model_selection import train_test_split 4 | import tensorflow as tf 5 | 6 | 7 | class DataSet(): 8 | ''' 9 | 通过读入文件生成数据集 10 | ''' 11 | 12 | def __init__(self, image_path, label_path): 13 | 14 | self.image_path = np.array(image_path) 15 | self.label_path = np.array(label_path) 16 | self.batch_count = 0 17 | self.epoch_count = 0 18 | 19 | def num_examples(self): 20 | ''' 21 | 得到样本的数量 22 | :return: 23 | ''' 24 | 25 | return self.image_path.shape[0] 26 | 27 | 28 | def next_batch(self, batch_size): 29 | ''' 30 | next_batch函数 31 | :param batch_size: 32 | :return: 33 | ''' 34 | 35 | start = self.batch_count * batch_size 36 | end = start + batch_size 37 | self.batch_count += 1 38 | 39 | if end > self.image_path.shape[0]: 40 | self.batch_count = 0 41 | random_index = np.random.permutation(self.image_path.shape[0]) 42 | self.image_path = self.image_path[random_index] 43 | self.label_path = self.label_path[random_index] 44 | self.epoch_count += 1 45 | start = self.batch_count * batch_size 46 | end = start + batch_size 47 | self.batch_count += 1 48 | 49 | image_batch, label_batch = self.read_path(self.image_path[start:end], 50 | self.label_path[start:end]) 51 | return image_batch, label_batch 52 | 53 | def read_path(self, x_path, y_path): 54 | ''' 55 | 将路径读为图片 56 | :param x_path: 57 | :param y_path: 58 | :return: 59 | ''' 60 | x = [] 61 | y = [] 62 | for i in range(x_path.shape[0]): 63 | x.append(self.transform(cv2.imread(x_path[i], 1))) 64 | y.append(cv2.imread(y_path[i], 0)) 65 | 66 | return np.array(x), np.array(y) 67 | 68 | def transform(self, img): 69 | 70 | return img --------------------------------------------------------------------------------