一、背景：
由于项目需要，参考了多篇相关车道线检测论文与源码，设计了一套Tensorflow版车道线检测功能。
二、基本结构：
该模型主要由以下部分组成：
1、数据源：包括所有原始数据，分组后的数据；
2、数据预处理：包括数据的准备，数据的导入，数据的提取，数据的分组（训练与测试）；
3、配置文件：包括各种参数与超参数，如：训练周期，训练步长，批量数据，学习率，卷积核大小，全连接大小，训练模型存放路径（checkpoint），摘要存放路径(summary)等；
4、基础网络：包括基本的网络组件，基础网，
5、训练主文件：主入口，用于搭建生成图（graph）,会话（sess）,数据导入模型训练，ＧＰＵ配置，训练过程打印等
三、代码结构
以下为原始文件夹：
./data
-- ./InstanceSegmentationClass
-- ./JPEGImages
-- ./SegmentationClass
-- datasets_gen_culane.py 用于从上面三个图片目录生成list.txt,train.txt,test.txt
# coding=utf-8 #create date:12/5/2018 #modified date:2/12/2019 #author:jim.chen import os import glob import random import math import cv2 import numpy as np def gen_list_txt(rela_dir,img_dir,img_seg_dir,img_inst_dir): #cwd = os.getcwd() #print("gen_list_txt cwd:",cwd) list_txt = "list.txt" png_glob = img_seg_dir+'/*.png' png_list_path = glob.glob(png_glob) png_list=[] print("gen_list_txt png_list_path:",png_list_path) with open(list_txt,"w") as w_f: for png in png_list_path: path,name = os.path.splitext(os.path.basename(png)) print("path:",path) w_f.write(rela_dir+img_dir+'/'+path+'.jpg'+' '+rela_dir+png+' '+rela_dir+img_inst_dir+'/'+path+'.png'+'\n') w_f.close() with open(list_txt,"r") as r_f: for each_line in r_f: png_list.append(each_line) png_list.sort() print("gen_list_txt len(png_list):",len(png_list)) train=random.sample(png_list,int(math.floor(len(png_list)*9/10))) train.sort() print("gen_list_txt train:",train) val=list(set(png_list).difference(set(train))) print("gen_list_txt val:",val) enum_train_val=['train','val'] for item in enum_train_val: with open(item+'.txt','w') as w1_f: for num_item in eval(item): print("gen_list_txt num_item:",num_item) w1_f.write(num_item) w1_f.close() def sync_gt_2_img(img_dir,img_seg_dir,img_inst_dir): cwd = os.getcwd() print("sync_gt_2_img img_dir:",img_dir," img_seg_dir:",img_seg_dir," img_inst_dir:",img_inst_dir) img_full_dir = cwd + '/' +img_dir img_seg_full_dir = cwd + '/' +img_seg_dir img_inst_full_dir = cwd + '/' +img_inst_dir img_list = os.listdir(img_full_dir) for img in img_list: img_basename = os.path.splitext(img)[0] print("sync_gt_2_img img_basename:",img_basename) img_full_path = img_full_dir + '/'+ img img_seg_full_path = img_seg_full_dir + '/'+img_basename +'.png' #print("sync_gt_2_img img_seg_full_path:",img_seg_full_path) img_inst_full_path = img_inst_full_dir + '/'+img_basename +'.png' #print("sync_gt_2_img img_inst_full_path:",img_inst_full_path) if not os.path.exists(img_inst_full_path): print("sync_gt_2_img not os.path.exists(img_seg_full_path)") if os.path.exists(img_full_path): os.remove(img_full_path) if os.path.exists(img_full_path): os.remove(img_full_path) def sync_seg_2_inst(img_seg_dir,img_inst_dir): cwd = os.getcwd() print("sync_seg_2_inst img_seg_dir:",img_seg_dir," img_inst_dir:",img_inst_dir) img_seg_full_dir = cwd + '/' +img_seg_dir img_inst_full_dir = cwd + '/' +img_inst_dir img_list = os.listdir(img_seg_dir) for img in img_list: img_basename = os.path.splitext(img)[0] print("sync_seg_2_inst img_basename:",img_basename) img_seg_full_path = img_seg_full_dir + '/'+img_basename +'.jpg' img_inst_full_path = img_inst_full_dir + '/'+img_basename +'.png' if not os.path.exists(img_inst_full_path): if os.path.exists(img_seg_full_path): print("sync_seg_2_inst os.remove(img_seg_full_path):",img_seg_full_path) os.remove(img_seg_full_path) def gen_seg_color(img_inst_dir,img_seg_dir): cwd = os.getcwd() inPath = os.path.join(cwd,img_inst_dir) print(inPath) outPath=os.path.join(cwd,img_seg_dir) inPathDir = os.listdir(inPath) if not os.path.exists(outPath): os.makedirs(outPath) for l,file_name in enumerate(inPathDir): img_instance = cv2.imread(os.path.join(inPath,file_name)) h,w,c = img_instance.shape print("l:",l," img_instance.shape:",img_instance.shape) img_instance_new = np.zeros((h, w, c), dtype=np.uint8) for i in range(0,h): for j in range(0,w): #print(img_instance[i][j]) if img_instance[i][j][0] != 0: img_instance_new[i][j] = [255,255,255] img_instance_gray = cv2.cvtColor(img_instance_new, cv2.COLOR_BGR2GRAY) cv2.imwrite(os.path.join(outPath,file_name), img_instance_gray) print("generate segment finished!") def gen_inst_color(img_inst_dir): cwd = os.getcwd() inPath = os.path.join(cwd,img_inst_dir) print(inPath) outPath=os.path.join(cwd,"img_inst_new") inPathDir = os.listdir(inPath) if not os.path.exists(outPath): os.makedirs(outPath) for l,file_name in enumerate(inPathDir): img_instance = cv2.imread(os.path.join(inPath,file_name)) h,w,c = img_instance.shape print("l:",l," img_instance.shape:",img_instance.shape) img_instance_new = np.zeros((h, w, c), dtype=np.uint8) for i in range(0,h): for j in range(0,w): #print(img_instance[i][j]) if img_instance[i][j][0] == 2: img_instance_new[i][j] = [20,20,20] elif img_instance[i][j][0] == 3: img_instance_new[i][j] = [70,70,70] elif img_instance[i][j][0] == 4: img_instance_new[i][j] = [120,120,120] elif img_instance[i][j][0] == 5: img_instance_new[i][j] = [170,170,170] elif img_instance[i][j][0] == 6: img_instance_new[i][j] = [220,220,220] img_instance_gray = cv2.cvtColor(img_instance_new, cv2.COLOR_BGR2GRAY) cv2.imwrite(os.path.join(outPath,file_name), img_instance_gray) print("generate instance finished!") def detect_invalid_img(img_path): img_instance = cv2.imread(img_path) h,w,c = img_instance.shape print("detect_invalid_img img_instance.shape:",img_instance.shape) for i in range(0,h): for j in range(0,w): if img_instance[i][j][0] != 0: return False return True def filter_invalid_img(img_test_dir,img_seg_dir,img_inst_dir): cwd = os.getcwd() inPath = os.path.join(cwd,img_test_dir) inPathDir = os.listdir(inPath) print("filter_invalid_img inPathDir:",inPathDir) for l,file_name in enumerate(inPathDir): img_path = os.path.join(inPath,file_name) isdel = detect_invalid_img(img_path) if isdel: print("filter_invalid_img isdel:",isdel) os.remove(os.path.join(inPath,file_name)) def main(): print("main begin") rela_dir = "data/datasets_culane_all/" img_dir = "image" img_seg_dir = "gt_image_binary" img_inst_dir = "gt_image_instance" gen_list_txt(rela_dir,img_dir,img_seg_dir,img_inst_dir) #sync_gt_2_img(img_dir,img_seg_dir,img_inst_dir) #gen_seg_color(img_inst_dir,img_seg_dir) #gen_inst_color(img_inst_dir) #filter_invalid_img(img_dir,img_seg_dir,img_inst_dir) #sync_seg_2_inst(img_dir,img_inst_dir) print("main end") if __name__ == '__main__': main()
--list.txt
data/datasets_culane/JPEGImages/0000.jpg data/datasets_culane/SegmentationClass\0000.png data/datasets_culane/InstanceSegmentationClass/0000.png data/datasets_culane/JPEGImages/0001.jpg data/datasets_culane/SegmentationClass\0001.png data/datasets_culane/InstanceSegmentationClass/0001.png data/datasets_culane/JPEGImages/0002.jpg data/datasets_culane/SegmentationClass\0002.png data/datasets_culane/InstanceSegmentationClass/0002.png data/datasets_culane/JPEGImages/0003.jpg data/datasets_culane/SegmentationClass\0003.png data/datasets_culane/InstanceSegmentationClass/0003.png data/datasets_culane/JPEGImages/0004.jpg data/datasets_culane/SegmentationClass\0004.png data/datasets_culane/InstanceSegmentationClass/0004.png data/datasets_culane/JPEGImages/0005.jpg data/datasets_culane/SegmentationClass\0005.png data/datasets_culane/InstanceSegmentationClass/0005.png
--train.txt
data/datasets_culane/JPEGImages/0000.jpg data/datasets_culane/SegmentationClass\0000.png data/datasets_culane/InstanceSegmentationClass/0000.png data/datasets_culane/JPEGImages/0001.jpg data/datasets_culane/SegmentationClass\0001.png data/datasets_culane/InstanceSegmentationClass/0001.png data/datasets_culane/JPEGImages/0002.jpg data/datasets_culane/SegmentationClass\0002.png data/datasets_culane/InstanceSegmentationClass/0002.png data/datasets_culane/JPEGImages/0004.jpg data/datasets_culane/SegmentationClass\0004.png data/datasets_culane/InstanceSegmentationClass/0004.png
--val.txt
data/datasets_culane/JPEGImages/0005.jpg data/datasets_culane/SegmentationClass\0005.png data/datasets_culane/InstanceSegmentationClass/0005.png
data/datasets_culane/JPEGImages/0003.jpg data/datasets_culane/SegmentationClass\0003.png data/datasets_culane/InstanceSegmentationClass/0003.png
./data_provider
--data_processor.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import os.path as ops import cv2 import numpy as np try: from cv2 import cv2 except ImportError: pass class DataSet(object): def __init__(self, dataset_info_file): self._gt_img_list, self._gt_label_binary_list, \ self._gt_label_instance_list = self._init_dataset(dataset_info_file) self._random_dataset() self._next_batch_loop_count = 0 def _init_dataset(self, dataset_info_file): gt_img_list = [] gt_label_binary_list = [] gt_label_instance_list = [] assert ops.exists(dataset_info_file), '{:s} not exist'.format(dataset_info_file) with open(dataset_info_file, 'r') as file: for _info in file: info_tmp = _info.strip(' ').split() gt_img_list.append(info_tmp[0]) gt_label_binary_list.append(info_tmp[1]) gt_label_instance_list.append(info_tmp[2]) return gt_img_list, gt_label_binary_list, gt_label_instance_list def _random_dataset(self): assert len(self._gt_img_list) == len(self._gt_label_binary_list) == len(self._gt_label_instance_list) random_idx = np.random.permutation(len(self._gt_img_list)) new_gt_img_list = [] new_gt_label_binary_list = [] new_gt_label_instance_list = [] for index in random_idx: new_gt_img_list.append(self._gt_img_list[index]) new_gt_label_binary_list.append(self._gt_label_binary_list[index]) new_gt_label_instance_list.append(self._gt_label_instance_list[index]) self._gt_img_list = new_gt_img_list self._gt_label_binary_list = new_gt_label_binary_list self._gt_label_instance_list = new_gt_label_instance_list def next_batch(self, batch_size): """ :param batch_size: :return: """ assert len(self._gt_label_binary_list) == len(self._gt_label_instance_list) \ == len(self._gt_img_list) idx_start = batch_size * self._next_batch_loop_count idx_end = batch_size * self._next_batch_loop_count + batch_size if idx_start == 0 and idx_end > len(self._gt_label_binary_list): raise ValueError('Batch size cant be more than total numbers') if idx_end > len(self._gt_label_binary_list): self._random_dataset() self._next_batch_loop_count = 0 return self.next_batch(batch_size) else: gt_img_list = self._gt_img_list[idx_start:idx_end] gt_label_binary_list = self._gt_label_binary_list[idx_start:idx_end] gt_label_instance_list = self._gt_label_instance_list[idx_start:idx_end] gt_imgs = [] gt_labels_binary = [] gt_labels_instance = [] for gt_img_path in gt_img_list: gt_imgs.append(cv2.imread(gt_img_path, cv2.IMREAD_COLOR)) for gt_label_path in gt_label_binary_list: label_img = cv2.imread(gt_label_path, cv2.IMREAD_COLOR) label_binary = np.zeros([label_img.shape[0], label_img.shape[1]], dtype=np.uint8) idx = np.where((label_img[:, :, :] != [0, 0, 0]).all(axis=2)) label_binary[idx] = 1 gt_labels_binary.append(label_binary) for gt_label_path in gt_label_instance_list: label_img = cv2.imread(gt_label_path, cv2.IMREAD_UNCHANGED) gt_labels_instance.append(label_img) self._next_batch_loop_count += 1 return gt_imgs, gt_labels_binary, gt_labels_instance if __name__ == '__main__': val = DataSet('/media/baidu/Data/Semantic_Segmentation/TUSimple_Lane_Detection/training/val.txt') b1, b2, b3 = val.next_batch(50) c1, c2, c3 = val.next_batch(50) dd, d2, d3 = val.next_batch(50)
./config
--global_config.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- from easydict import EasyDict as edict __C = edict() # Consumers can get config by: from config import cfg cfg = __C # Train options __C.TRAIN = edict() # Set the shadownet training epochs __C.TRAIN.EPOCHS = 200010 # Set the display step __C.TRAIN.DISPLAY_STEP = 1 # Set the test display step during training process __C.TRAIN.TEST_DISPLAY_STEP = 1000 # Set the momentum parameter of the optimizer __C.TRAIN.MOMENTUM = 0.9 # Set the initial learning rate __C.TRAIN.LEARNING_RATE = 0.0005 # Set the GPU resource used during training process __C.TRAIN.GPU_MEMORY_FRACTION = 0.85 # Set the GPU allow growth parameter during tensorflow training process __C.TRAIN.TF_ALLOW_GROWTH = True # Set the shadownet training batch size __C.TRAIN.BATCH_SIZE = 1 # Set the shadownet validation batch size __C.TRAIN.VAL_BATCH_SIZE = 1 # Set the learning rate decay steps __C.TRAIN.LR_DECAY_STEPS = 410000 # Set the learning rate decay rate __C.TRAIN.LR_DECAY_RATE = 0.1 # Set the class numbers __C.TRAIN.CLASSES_NUMS = 2 # Set the image height __C.TRAIN.IMG_HEIGHT = 256 # Set the image width __C.TRAIN.IMG_WIDTH = 512 # Test options __C.TEST = edict() # Set the GPU resource used during testing process __C.TEST.GPU_MEMORY_FRACTION = 0.8 # Set the GPU allow growth parameter during tensorflow testing process __C.TEST.TF_ALLOW_GROWTH = True # Set the test batch size __C.TEST.BATCH_SIZE = 1
./encoder_decoder_model
--cnn_basenet.py
#!/usr/bin/env python # -*- coding: utf-8 -*- """ The base convolution neural networks mainly implement some useful cnn functions """ import tensorflow as tf import numpy as np class CNNBaseModel(object): """ Base model for other specific cnn ctpn_models """ def __init__(self): pass @staticmethod def conv2d(inputdata, out_channel, kernel_size, padding='SAME', stride=1, w_init=None, b_init=None, split=1, use_bias=True, data_format='NHWC', name=None): with tf.variable_scope(name): in_shape = inputdata.get_shape().as_list() channel_axis = 3 if data_format == 'NHWC' else 1 in_channel = in_shape[channel_axis] assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!" assert in_channel % split == 0 assert out_channel % split == 0 padding = padding.upper() if isinstance(kernel_size, list): filter_shape = [kernel_size[0], kernel_size[1]] + [in_channel / split, out_channel] else: filter_shape = [kernel_size, kernel_size] + [in_channel / split, out_channel] if isinstance(stride, list): strides = [1, stride[0], stride[1], 1] if data_format == 'NHWC' \ else [1, 1, stride[0], stride[1]] else: strides = [1, stride, stride, 1] if data_format == 'NHWC' \ else [1, 1, stride, stride] if w_init is None: w_init = tf.contrib.layers.variance_scaling_initializer() if b_init is None: b_init = tf.constant_initializer() w = tf.get_variable('W', filter_shape, initializer=w_init) b = None if use_bias: b = tf.get_variable('b', [out_channel], initializer=b_init) if split == 1: conv = tf.nn.conv2d(inputdata, w, strides, padding, data_format=data_format) else: inputs = tf.split(inputdata, split, channel_axis) kernels = tf.split(w, split, 3) outputs = [tf.nn.conv2d(i, k, strides, padding, data_format=data_format) for i, k in zip(inputs, kernels)] conv = tf.concat(outputs, channel_axis) ret = tf.identity(tf.nn.bias_add(conv, b, data_format=data_format) if use_bias else conv, name=name) return ret @staticmethod def relu(inputdata, name=None): return tf.nn.relu(features=inputdata, name=name) @staticmethod def sigmoid(inputdata, name=None): return tf.nn.sigmoid(x=inputdata, name=name) @staticmethod def maxpooling(inputdata, kernel_size, stride=None, padding='VALID', data_format='NHWC', name=None): padding = padding.upper() if stride is None: stride = kernel_size if isinstance(kernel_size, list): kernel = [1, kernel_size[0], kernel_size[1], 1] if data_format == 'NHWC' else \ [1, 1, kernel_size[0], kernel_size[1]] else: kernel = [1, kernel_size, kernel_size, 1] if data_format == 'NHWC' \ else [1, 1, kernel_size, kernel_size] if isinstance(stride, list): strides = [1, stride[0], stride[1], 1] if data_format == 'NHWC' \ else [1, 1, stride[0], stride[1]] else: strides = [1, stride, stride, 1] if data_format == 'NHWC' \ else [1, 1, stride, stride] return tf.nn.max_pool(value=inputdata, ksize=kernel, strides=strides, padding=padding, data_format=data_format, name=name) @staticmethod def avgpooling(inputdata, kernel_size, stride=None, padding='VALID', data_format='NHWC', name=None): if stride is None: stride = kernel_size kernel = [1, kernel_size, kernel_size, 1] if data_format == 'NHWC' \ else [1, 1, kernel_size, kernel_size] strides = [1, stride, stride, 1] if data_format == 'NHWC' else [1, 1, stride, stride] return tf.nn.avg_pool(value=inputdata, ksize=kernel, strides=strides, padding=padding, data_format=data_format, name=name) @staticmethod def globalavgpooling(inputdata, data_format='NHWC', name=None): assert inputdata.shape.ndims == 4 assert data_format in ['NHWC', 'NCHW'] axis = [1, 2] if data_format == 'NHWC' else [2, 3] return tf.reduce_mean(input_tensor=inputdata, axis=axis, name=name) @staticmethod def layernorm(inputdata, epsilon=1e-5, use_bias=True, use_scale=True, data_format='NHWC', name=None): shape = inputdata.get_shape().as_list() ndims = len(shape) assert ndims in [2, 4] mean, var = tf.nn.moments(inputdata, list(range(1, len(shape))), keep_dims=True) if data_format == 'NCHW': channnel = shape[1] new_shape = [1, channnel, 1, 1] else: channnel = shape[-1] new_shape = [1, 1, 1, channnel] if ndims == 2: new_shape = [1, channnel] if use_bias: beta = tf.get_variable('beta', [channnel], initializer=tf.constant_initializer()) beta = tf.reshape(beta, new_shape) else: beta = tf.zeros([1] * ndims, name='beta') if use_scale: gamma = tf.get_variable('gamma', [channnel], initializer=tf.constant_initializer(1.0)) gamma = tf.reshape(gamma, new_shape) else: gamma = tf.ones([1] * ndims, name='gamma') return tf.nn.batch_normalization(inputdata, mean, var, beta, gamma, epsilon, name=name) @staticmethod def instancenorm(inputdata, epsilon=1e-5, data_format='NHWC', use_affine=True, name=None): shape = inputdata.get_shape().as_list() if len(shape) != 4: raise ValueError("Input data of instancebn layer has to be 4D tensor") if data_format == 'NHWC': axis = [1, 2] ch = shape[3] new_shape = [1, 1, 1, ch] else: axis = [2, 3] ch = shape[1] new_shape = [1, ch, 1, 1] if ch is None: raise ValueError("Input of instancebn require known channel!") mean, var = tf.nn.moments(inputdata, axis, keep_dims=True) if not use_affine: return tf.divide(inputdata - mean, tf.sqrt(var + epsilon), name='output') beta = tf.get_variable('beta', [ch], initializer=tf.constant_initializer()) beta = tf.reshape(beta, new_shape) gamma = tf.get_variable('gamma', [ch], initializer=tf.constant_initializer(1.0)) gamma = tf.reshape(gamma, new_shape) return tf.nn.batch_normalization(inputdata, mean, var, beta, gamma, epsilon, name=name) @staticmethod def dropout(inputdata, keep_prob, noise_shape=None, name=None): return tf.nn.dropout(inputdata, keep_prob=keep_prob, noise_shape=noise_shape, name=name) @staticmethod def fullyconnect(inputdata, out_dim, w_init=None, b_init=None, use_bias=True, name=None): shape = inputdata.get_shape().as_list()[1:] if None not in shape: inputdata = tf.reshape(inputdata, [-1, int(np.prod(shape))]) else: inputdata = tf.reshape(inputdata, tf.stack([tf.shape(inputdata)[0], -1])) if w_init is None: w_init = tf.contrib.layers.variance_scaling_initializer() if b_init is None: b_init = tf.constant_initializer() ret = tf.layers.dense(inputs=inputdata, activation=lambda x: tf.identity(x, name='output'), use_bias=use_bias, name=name, kernel_initializer=w_init, bias_initializer=b_init, trainable=True, units=out_dim) return ret @staticmethod def layerbn(inputdata, is_training, name): return tf.layers.batch_normalization(inputs=inputdata, training=is_training, name=name) @staticmethod def squeeze(inputdata, axis=None, name=None): return tf.squeeze(input=inputdata, axis=axis, name=name) @staticmethod def deconv2d(inputdata, out_channel, kernel_size, padding='SAME', stride=1, w_init=None, b_init=None, use_bias=True, activation=None, data_format='channels_last', trainable=True, name=None): with tf.variable_scope(name): in_shape = inputdata.get_shape().as_list() channel_axis = 3 if data_format == 'channels_last' else 1 in_channel = in_shape[channel_axis] assert in_channel is not None, "[Deconv2D] Input cannot have unknown channel!" padding = padding.upper() if w_init is None: w_init = tf.contrib.layers.variance_scaling_initializer() if b_init is None: b_init = tf.constant_initializer() ret = tf.layers.conv2d_transpose(inputs=inputdata, filters=out_channel, kernel_size=kernel_size, strides=stride, padding=padding, data_format=data_format, activation=activation, use_bias=use_bias, kernel_initializer=w_init, bias_initializer=b_init, trainable=trainable, name=name) return ret @staticmethod def dilation_conv(input_tensor, k_size, out_dims, rate, padding='SAME', w_init=None, b_init=None, use_bias=False, name=None): with tf.variable_scope(name): in_shape = input_tensor.get_shape().as_list() in_channel = in_shape[3] assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!" padding = padding.upper() if isinstance(k_size, list): filter_shape = [k_size[0], k_size[1]] + [in_channel, out_dims] else: filter_shape = [k_size, k_size] + [in_channel, out_dims] if w_init is None: w_init = tf.contrib.layers.variance_scaling_initializer() if b_init is None: b_init = tf.constant_initializer() w = tf.get_variable('W', filter_shape, initializer=w_init) b = None if use_bias: b = tf.get_variable('b', [out_dims], initializer=b_init) conv = tf.nn.atrous_conv2d(value=input_tensor, filters=w, rate=rate, padding=padding, name='dilation_conv') if use_bias: ret = tf.add(conv, b) else: ret = conv return ret @staticmethod def spatial_dropout(input_tensor, keep_prob, is_training, name, seed=1234): tf.set_random_seed(seed=seed) def f1(): with tf.variable_scope(name): return input_tensor def f2(): with tf.variable_scope(name): num_feature_maps = [tf.shape(input_tensor)[0], tf.shape(input_tensor)[3]] random_tensor = keep_prob random_tensor += tf.random_uniform(num_feature_maps, seed=seed, dtype=input_tensor.dtype) binary_tensor = tf.floor(random_tensor) binary_tensor = tf.reshape(binary_tensor, [-1, 1, 1, tf.shape(input_tensor)[3]]) ret = input_tensor * binary_tensor return ret output = tf.cond(is_training, f2, f1) return output @staticmethod def lrelu(inputdata, name, alpha=0.2): with tf.variable_scope(name): return tf.nn.relu(inputdata) - alpha * tf.nn.relu(-inputdata)
--vgg_scnn_encoder.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- from collections import OrderedDict import tensorflow as tf import glog as log import math import sys sys.path.append('encoder_decoder_model') import cnn_basenet class VGG16Encoder(cnn_basenet.CNNBaseModel): def __init__(self, phase): super(VGG16Encoder, self).__init__() self._train_phase = tf.constant('train', dtype=tf.string) self._test_phase = tf.constant('test', dtype=tf.string) self._phase = phase self._is_training = self._init_phase() def _init_phase(self): return tf.equal(self._phase, self._train_phase) def _conv_stage(self, input_tensor, k_size, out_dims, name, stride=1, pad='SAME'): with tf.variable_scope(name): conv = self.conv2d(inputdata=input_tensor, out_channel=out_dims, kernel_size=k_size, stride=stride, use_bias=False, padding=pad, name='conv') bn = self.layerbn(inputdata=conv, is_training=self._is_training, name='bn') relu = self.relu(inputdata=bn, name='relu') return relu def _fc_stage(self, input_tensor, out_dims, name, use_bias=False): with tf.variable_scope(name): fc = self.fullyconnect(inputdata=input_tensor, out_dim=out_dims, use_bias=use_bias, name='fc') bn = self.layerbn(inputdata=fc, is_training=self._is_training, name='bn') relu = self.relu(inputdata=bn, name='relu') return relu def scnn_u2d_d2u(self,input_tensor): output_list_old = [] output_list_new = [] shape_list = input_tensor.get_shape().as_list() log.info("scnn_u2d_d2u shape_list:{:}".format(shape_list)) h_size = input_tensor.get_shape().as_list()[1] log.info("scnn_u2d_d2u h_size:{:}".format(h_size)) channel_size = input_tensor.get_shape().as_list()[3] #up2down conv for i in range(h_size): output_list_old.append(tf.expand_dims(input_tensor[:,i,:,:],axis=1)) output_list_new.append(tf.expand_dims(input_tensor[:,0,:,:],axis=1)) w_ud = tf.get_variable('w_ud',[1,9,channel_size,channel_size],initializer=tf.random_normal_initializer(0,math.sqrt(2.0/(9*channel_size*channel_size*2)))) with tf.variable_scope("scnn_u2d"): scnn_u2d = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_old[0],w_ud,[1,1,1,1],'SAME')),output_list_old[1]) output_list_new.append(scnn_u2d) for i in range(2,h_size): with tf.variable_scope("scnn_u2d",reuse=True): scnn_u2d = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_new[i-1],w_ud,[1,1,1,1],'SAME')),output_list_old[i]) output_list_new.append(scnn_u2d) #down2up conv output_list_old = output_list_new output_list_new = [] length = h_size-1 output_list_new.append(output_list_old[length]) w_du = tf.get_variable('w_du',[1,9,channel_size,channel_size],initializer=tf.random_normal_initializer(0,math.sqrt(2.0/(9*channel_size*channel_size*2)))) with tf.variable_scope('scnn_d2u'): scnn_d2u = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_old[length],w_du,[1,1,1,1],'SAME')),output_list_old[length-1]) output_list_new.append(scnn_d2u) for i in range(2,h_size): with tf.variable_scope("scnn_d2u",reuse=True): scnn_d2u = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_new[i-1],w_du,[1,1,1,1],'SAME')),output_list_old[length-i]) output_list_new.append(scnn_d2u) output_list_new.reverse() #log.info("scnn_u2d_d2u output_list_new:{:}".format(output_list_new)) out_tensor = tf.stack(output_list_new,axis = 1) out_tensor = tf.squeeze(out_tensor,axis=2) return out_tensor def scnn_l2r_r2l(self,input_tensor): output_list_old = [] output_list_new = [] shape_list = input_tensor.get_shape().as_list() log.info("scnn_l2r_r2l shape_list:{:}".format(shape_list)) w_size = input_tensor.get_shape().as_list()[2] log.info("scnn_l2r_r2l w_size:{:}".format(w_size)) channel_size = input_tensor.get_shape().as_list()[3] #left2right conv for i in range(w_size): output_list_old.append(tf.expand_dims(input_tensor[:,:,i,:],axis=2)) output_list_new.append(tf.expand_dims(input_tensor[:,:,0,:],axis=2)) w_lr = tf.get_variable('w_lr',[9,1,channel_size,channel_size],initializer=tf.random_normal_initializer(0,math.sqrt(2.0/(9*channel_size*channel_size*5)))) with tf.variable_scope("scnn_l2r"): scnn_l2r = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_old[0],w_lr,[1,1,1,1],'SAME')),output_list_old[1]) output_list_new.append(scnn_l2r) for i in range(2,w_size): with tf.variable_scope("scnn_l2r",reuse=True): scnn_l2r = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_new[i-1],w_lr,[1,1,1,1],'SAME')),output_list_old[i]) output_list_new.append(scnn_l2r) #log.info("output_list_new:{:}".format(output_list_new)) #right2left conv output_list_old = output_list_new output_list_new = [] length = w_size-1 output_list_new.append(output_list_old[length]) w_rl = tf.get_variable('w_rl',[9,1,channel_size,channel_size],initializer=tf.random_normal_initializer(0,math.sqrt(2.0/(9*channel_size*channel_size*5)))) with tf.variable_scope('scnn_r2l'): scnn_r2l = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_old[length],w_rl,[1,1,1,1],'SAME')),output_list_old[length-1]) output_list_new.append(scnn_r2l) for i in range(2,w_size): with tf.variable_scope("scnn_r2l",reuse=True): scnn_r2l = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_new[i-1],w_rl,[1,1,1,1],'SAME')),output_list_old[length-i]) output_list_new.append(scnn_r2l) output_list_new.reverse() out_tensor = tf.stack(output_list_new,axis = 2) out_tensor = tf.squeeze(out_tensor,axis=3) return out_tensor def encode(self, input_tensor, name): ret = OrderedDict() with tf.variable_scope(name): # conv stage 1_1 conv_1_1 = self._conv_stage(input_tensor=input_tensor, k_size=3, out_dims=64, name='conv1_1') log.info("encode conv_1_1:{:}".format(conv_1_1.get_shape().as_list())) # conv stage 1_2 conv_1_2 = self._conv_stage(input_tensor=conv_1_1, k_size=3, out_dims=64, name='conv1_2') log.info("encode conv_1_2:{:}".format(conv_1_2.get_shape().as_list())) # pool stage 1 pool1 = self.maxpooling(inputdata=conv_1_2, kernel_size=2, stride=2, name='pool1') log.info("encode pool1:{:}".format(pool1.get_shape().as_list())) # conv stage 2_1 conv_2_1 = self._conv_stage(input_tensor=pool1, k_size=3, out_dims=128, name='conv2_1') log.info("encode conv_2_1:{:}".format(conv_2_1.get_shape().as_list())) # conv stage 2_2 conv_2_2 = self._conv_stage(input_tensor=conv_2_1, k_size=3, out_dims=128, name='conv2_2') log.info("encode conv_2_2:{:}".format(conv_2_2.get_shape().as_list())) # pool stage 2 pool2 = self.maxpooling(inputdata=conv_2_2, kernel_size=2, stride=2, name='pool2') log.info("encode pool2:{:}".format(pool2.get_shape().as_list())) # conv stage 3_1 conv_3_1 = self._conv_stage(input_tensor=pool2, k_size=3, out_dims=256, name='conv3_1') log.info("encode conv_3_1:{:}".format(conv_3_1.get_shape().as_list())) # conv_stage 3_2 conv_3_2 = self._conv_stage(input_tensor=conv_3_1, k_size=3, out_dims=256, name='conv3_2') log.info("encode conv_3_2:{:}".format(conv_3_2.get_shape().as_list())) # conv stage 3_3 conv_3_3 = self._conv_stage(input_tensor=conv_3_2, k_size=3, out_dims=256, name='conv3_3') log.info("encode conv_3_3:{:}".format(conv_3_3.get_shape().as_list())) ret['conv_3_3'] = dict() ret['conv_3_3']['data'] = conv_3_3 ret['conv_3_3']['shape'] = conv_3_3.get_shape().as_list() # pool stage 3 pool3 = self.maxpooling(inputdata=conv_3_3, kernel_size=2, stride=2, name='pool3') log.info("encode pool3:{:}".format(pool3.get_shape().as_list())) ret['pool3'] = dict() ret['pool3']['data'] = pool3 ret['pool3']['shape'] = pool3.get_shape().as_list() # conv stage 4_1 conv_4_1 = self._conv_stage(input_tensor=pool3, k_size=3, out_dims=512, name='conv4_1') log.info("encode conv_4_1:{:}".format(conv_4_1.get_shape().as_list())) # conv stage 4_2 conv_4_2 = self._conv_stage(input_tensor=conv_4_1, k_size=3, out_dims=512, name='conv4_2') log.info("encode conv_4_2:{:}".format(conv_4_2.get_shape().as_list())) # conv stage 4_3 conv_4_3 = self._conv_stage(input_tensor=conv_4_2, k_size=3, out_dims=512, name='conv4_3') log.info("encode conv_4_3:{:}".format(conv_4_3.get_shape().as_list())) # pool stage 4 pool4 = self.maxpooling(inputdata=conv_4_3, kernel_size=2, stride=2, name='pool4') log.info("encode pool4:{:}".format(pool4.get_shape().as_list())) ret['pool4'] = dict() ret['pool4']['data'] = pool4 ret['pool4']['shape'] = pool4.get_shape().as_list() # conv stage 5_1 conv_5_1 = self._conv_stage(input_tensor=pool4, k_size=3, out_dims=512, name='conv5_1') log.info("encode conv_5_1:{:}".format(conv_5_1.get_shape().as_list())) # conv stage 5_2 conv_5_2 = self._conv_stage(input_tensor=conv_5_1, k_size=3, out_dims=512, name='conv5_2') log.info("encode conv_5_2:{:}".format(conv_5_2.get_shape().as_list())) # conv stage 5_3 conv_5_3 = self._conv_stage(input_tensor=conv_5_2, k_size=3, out_dims=512, name='conv5_3') log.info("encode conv_5_3:{:}".format(conv_5_3.get_shape().as_list())) # conv stage 6_1 conv_6_1 = self._conv_stage(input_tensor=conv_5_3, k_size=3, out_dims=128, name='conv6_1') log.info("encode conv_6_1:{:}".format(conv_6_1.get_shape().as_list())) scnn_ud = self.scnn_u2d_d2u(conv_6_1) log.info("encode scnn_ud:{:}".format(scnn_ud.get_shape().as_list())) scnn_lr = self.scnn_l2r_r2l(scnn_ud) log.info("encode scnn_lr:{:}".format(scnn_lr.get_shape().as_list())) # pool stage 5 pool5 = self.maxpooling(inputdata=scnn_lr, kernel_size=2, stride=2, name='pool5') log.info("encode pool5:{:}".format(pool5.get_shape().as_list())) ret['pool5'] = dict() ret['pool5']['data'] = pool5 ret['pool5']['shape'] = pool5.get_shape().as_list() # fc stage 1 # fc6 = self._fc_stage(input_tensor=pool5, out_dims=4096, name='fc6', # use_bias=False, flags=flags) # fc stage 2 # fc7 = self._fc_stage(input_tensor=fc6, out_dims=4096, name='fc7', # use_bias=False, flags=flags) return ret if __name__ == '__main__': a = tf.placeholder(dtype=tf.float32, shape=[1, 2048, 2048, 3], name='input') encoder = VGG16Encoder(phase=tf.constant('train', dtype=tf.string)) ret = encoder.encode(a, name='encode') for layer_name, layer_info in ret.items(): print('layer name: {:s} shape: {}'.format(layer_name, layer_info['shape']))
--dense_encoder.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import tensorflow as tf from collections import OrderedDict #from encoder_decoder_model import cnn_basenet import cnn_basenet class DenseEncoder(cnn_basenet.CNNBaseModel): """ 基于DenseNet的编码器 """ def __init__(self, l, n, growthrate, phase, with_bc=False, bc_theta=0.5): super(DenseEncoder, self).__init__() self._L = l self._block_depth = int((l - n - 1) / n) self._N = n self._growthrate = growthrate self._with_bc = with_bc self._phase = phase self._train_phase = tf.constant('train', dtype=tf.string) self._test_phase = tf.constant('test', dtype=tf.string) self._is_training = self._init_phase() self._bc_theta = bc_theta return def _init_phase(self): return tf.equal(self._phase, self._train_phase) def __str__(self): encoder_info = 'A densenet with net depth: {:d} block nums: ' \ '{:d} growth rate: {:d} block depth: {:d}'.\ format(self._L, self._N, self._growthrate, self._block_depth) return encoder_info def _composite_conv(self, inputdata, out_channel, name): with tf.variable_scope(name): bn_1 = self.layerbn(inputdata=inputdata, is_training=self._is_training, name='bn_1') relu_1 = self.relu(bn_1, name='relu_1') if self._with_bc: conv_1 = self.conv2d(inputdata=relu_1, out_channel=out_channel, kernel_size=1, padding='SAME', stride=1, use_bias=False, name='conv_1') bn_2 = self.layerbn(inputdata=conv_1, is_training=self._is_training, name='bn_2') relu_2 = self.relu(inputdata=bn_2, name='relu_2') conv_2 = self.conv2d(inputdata=relu_2, out_channel=out_channel, kernel_size=3, stride=1, padding='SAME', use_bias=False, name='conv_2') return conv_2 else: conv_2 = self.conv2d(inputdata=relu_1, out_channel=out_channel, kernel_size=3, stride=1, padding='SAME', use_bias=False, name='conv_2') return conv_2 def _denseconnect_layers(self, inputdata, name): with tf.variable_scope(name): conv_out = self._composite_conv(inputdata=inputdata, name='composite_conv', out_channel=self._growthrate) concate_cout = tf.concat(values=[conv_out, inputdata], axis=3, name='concatenate') return concate_cout def _transition_layers(self, inputdata, name): """ Mainly implement the Pooling layer mentioned in DenseNet paper :param inputdata: :param name: :return: """ input_channels = inputdata.get_shape().as_list()[3] with tf.variable_scope(name): # First batch norm bn = self.layerbn(inputdata=inputdata, is_training=self._is_training, name='bn') # Second 1*1 conv if self._with_bc: out_channels = int(input_channels * self._bc_theta) conv = self.conv2d(inputdata=bn, out_channel=out_channels, kernel_size=1, stride=1, use_bias=False, name='conv') # Third average pooling avgpool_out = self.avgpooling(inputdata=conv, kernel_size=2, stride=2, name='avgpool') return avgpool_out else: conv = self.conv2d(inputdata=bn, out_channel=input_channels, kernel_size=1, stride=1, use_bias=False, name='conv') # Third average pooling avgpool_out = self.avgpooling(inputdata=conv, kernel_size=2, stride=2, name='avgpool') return avgpool_out def _dense_block(self, inputdata, name): """ Mainly implement the dense block mentioned in DenseNet figure 1 :param inputdata: :param name: :return: """ block_input = inputdata with tf.variable_scope(name): for i in range(self._block_depth): block_layer_name = '{:s}_layer_{:d}'.format(name, i + 1) block_input = self._denseconnect_layers(inputdata=block_input, name=block_layer_name) return block_input def encode(self, input_tensor, name): """ DenseNet编码 :param input_tensor: :param name: :return: """ encode_ret = OrderedDict() # First apply a 3*3 16 out channels conv layer # mentioned in DenseNet paper Implementation Details part with tf.variable_scope(name): conv1 = self.conv2d(inputdata=input_tensor, out_channel=16, kernel_size=3, use_bias=False, name='conv1') dense_block_input = conv1 # Second apply dense block stage for dense_block_nums in range(self._N): dense_block_name = 'Dense_Block_{:d}'.format(dense_block_nums + 1) # dense connectivity dense_block_out = self._dense_block(inputdata=dense_block_input, name=dense_block_name) # apply the trainsition part dense_block_out = self._transition_layers(inputdata=dense_block_out, name=dense_block_name) dense_block_input = dense_block_out encode_ret[dense_block_name] = dict() encode_ret[dense_block_name]['data'] = dense_block_out encode_ret[dense_block_name]['shape'] = dense_block_out.get_shape().as_list() return encode_ret if __name__ == '__main__': input_tensor = tf.placeholder(dtype=tf.float32, shape=[None, 384, 1248, 3], name='input_tensor') encoder = DenseEncoder(l=100, growthrate=16, with_bc=True, phase=tf.constant('train'), n=5) ret = encoder.encode(input_tensor=input_tensor, name='Dense_Encode') for layer_name, layer_info in ret.items(): print('layer_name: {:s} shape: {}'.format(layer_name, layer_info['shape']))
--fcn_decoder.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import tensorflow as tf #from encoder_decoder_model import cnn_basenet #from encoder_decoder_model import vgg_encoder #from encoder_decoder_model import dense_encoder import cnn_basenet import vgg_encoder import dense_encoder class FCNDecoder(cnn_basenet.CNNBaseModel): def __init__(self, phase): """ """ super(FCNDecoder, self).__init__() self._train_phase = tf.constant('train', dtype=tf.string) self._phase = phase self._is_training = self._init_phase() def _init_phase(self): """ :return: """ return tf.equal(self._phase, self._train_phase) def decode(self, input_tensor_dict, decode_layer_list, name): """ 解码特征信息反卷积还原 :param input_tensor_dict: :param decode_layer_list: 需要解码的层名称需要由深到浅顺序写 eg. ['pool5', 'pool4', 'pool3'] :param name: :return: """ ret = dict() with tf.variable_scope(name): # score stage 1 input_tensor = input_tensor_dict[decode_layer_list[0]]['data'] score = self.conv2d(inputdata=input_tensor, out_channel=64, kernel_size=1, use_bias=False, name='score_origin') ret['score'] = dict() ret['score']['data'] = score ret['score']['shape'] = score.get_shape().as_list() decode_layer_list = decode_layer_list[1:] print("len(decode_layer_list):",len(decode_layer_list)) for i in range(len(decode_layer_list)): deconv = self.deconv2d(inputdata=score, out_channel=64, kernel_size=4, stride=2, use_bias=False, name='deconv_{:d}'.format(i + 1)) input_tensor = input_tensor_dict[decode_layer_list[i]]['data'] score = self.conv2d(inputdata=input_tensor, out_channel=64, kernel_size=1, use_bias=False, name='score_{:d}'.format(i + 1)) fused = tf.add(deconv, score, name='fuse_{:d}'.format(i + 1)) score = fused ret['fuse_{:d}'.format(i + 1)] = dict() ret['fuse_{:d}'.format(i + 1)]['data'] = fused ret['fuse_{:d}'.format(i + 1)]['shape'] = fused.get_shape().as_list() deconv_final = self.deconv2d(inputdata=score, out_channel=64, kernel_size=16, stride=8, use_bias=False, name='deconv_final') score_final = self.conv2d(inputdata=deconv_final, out_channel=2, kernel_size=1, use_bias=False, name='score_final') ret['logits'] = score_final ret['deconv'] = deconv_final ret['logits'] = dict() ret['logits']['data'] = score_final ret['logits']['shape'] = score_final.get_shape().as_list() ret['deconv'] = dict() ret['deconv']['data'] = deconv_final ret['deconv']['shape'] = deconv_final.get_shape().as_list() return ret if __name__ == '__main__': vgg_encoder = vgg_encoder.VGG16Encoder(phase=tf.constant('train', tf.string)) dense_encoder = dense_encoder.DenseEncoder(l=40, growthrate=12, with_bc=True, phase='train', n=5) decoder = FCNDecoder(phase='train') in_tensor = tf.placeholder(dtype=tf.float32, shape=[None, 256, 512, 3], name='input') vgg_encode_ret = vgg_encoder.encode(in_tensor, name='vgg_encoder') dense_encode_ret = dense_encoder.encode(in_tensor, name='dense_encoder') decode_ret = decoder.decode(vgg_encode_ret, name='decoder', decode_layer_list=['pool5', 'pool4', 'pool3']) for layer_name, layer_info in decode_ret.items(): print('layer name: {:s} shape: {}'.format(layer_name, layer_info['shape']))
./merge_model
--merge_model.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import tensorflow as tf from encoder_decoder_model import vgg_encoder from encoder_decoder_model import fcn_decoder from encoder_decoder_model import dense_encoder from encoder_decoder_model import cnn_basenet from lanenet_model import lanenet_discriminative_loss from encoder_decoder_model import vgg_scnn_encoder import glog class LaneNet(cnn_basenet.CNNBaseModel): """ 实现语义分割模型 """ def __init__(self, phase, net_flag='vgg'): """ """ super(LaneNet, self).__init__() self._net_flag = net_flag self._phase = phase if self._net_flag == 'vgg': self._encoder = vgg_encoder.VGG16Encoder(phase=phase) elif self._net_flag == 'vgg_scnn': self._encoder = vgg_scnn_encoder.VGG16Encoder(phase=phase) elif self._net_flag == 'dense': self._encoder = dense_encoder.DenseEncoder(l=20, growthrate=8, with_bc=True, phase=phase, n=5) self._decoder = fcn_decoder.FCNDecoder(phase=phase) return def __str__(self): """ :return: """ info = 'Semantic Segmentation use {:s} as basenet to encode'.format(self._net_flag) return info def _build_model(self, input_tensor, name): """ 前向传播过程 :param input_tensor: :param name: :return: """ with tf.variable_scope(name): # first encode encode_ret = self._encoder.encode(input_tensor=input_tensor, name='encode') # second decode if self._net_flag.lower() == 'vgg': decode_ret = self._decoder.decode(input_tensor_dict=encode_ret, name='decode', decode_layer_list=['pool5', 'pool4', 'pool3']) return decode_ret if self._net_flag.lower() == 'vgg_scnn': decode_ret = self._decoder.decode(input_tensor_dict=encode_ret, name='decode', decode_layer_list=['pool5', 'pool4', 'pool3']) return decode_ret elif self._net_flag.lower() == 'dense': decode_ret = self._decoder.decode(input_tensor_dict=encode_ret, name='decode', decode_layer_list=['Dense_Block_5', 'Dense_Block_4', 'Dense_Block_3']) return decode_ret def compute_loss(self, input_tensor, binary_label, instance_label, name): """ 计算LaneNet模型损失函数 :param input_tensor: :param binary_label: :param instance_label: :param name: :return: """ with tf.variable_scope(name): # 前向传播获取logits inference_ret = self._build_model(input_tensor=input_tensor, name='inference') glog.info('compute_loss inference_ret:{:}'.format(inference_ret)) # 计算二值分割损失函数 decode_logits = inference_ret['logits'] binary_label_plain = tf.reshape( binary_label, shape=[binary_label.get_shape().as_list()[0] * binary_label.get_shape().as_list()[1] * binary_label.get_shape().as_list()[2]]) glog.info('compute_loss binary_label_plain:{:}'.format(binary_label_plain)) # 加入class weights unique_labels, unique_id, counts = tf.unique_with_counts(binary_label_plain) counts = tf.cast(counts, tf.float32) glog.info('compute_loss counts:{:}'.format(counts)) inverse_weights = tf.divide(1.0, tf.log(tf.add(tf.divide(tf.constant(1.0), counts), tf.constant(1.02)))) glog.info('compute_loss inverse_weights:{:}'.format(inverse_weights)) inverse_weights = tf.gather(inverse_weights, binary_label) glog.info('compute_loss gather inverse_weights:{:}'.format(inverse_weights)) binary_segmenatation_loss = tf.losses.sparse_softmax_cross_entropy( labels=binary_label, logits=decode_logits, weights=inverse_weights) glog.info('compute_loss binary_segmenatation_loss:{:}'.format(binary_segmenatation_loss)) binary_segmenatation_loss = tf.reduce_mean(binary_segmenatation_loss) glog.info('compute_loss reduce_mean binary_segmenatation_loss:{:}'.format(binary_segmenatation_loss)) # 计算discriminative loss损失函数 decode_deconv = inference_ret['deconv'] # 像素嵌入 pix_embedding = self.conv2d(inputdata=decode_deconv, out_channel=4, kernel_size=1, use_bias=False, name='pix_embedding_conv') pix_embedding = self.relu(inputdata=pix_embedding, name='pix_embedding_relu') # 计算discriminative loss image_shape = (pix_embedding.get_shape().as_list()[1], pix_embedding.get_shape().as_list()[2]) glog.info('compute_loss image_shape:{:}'.format(image_shape)) disc_loss, l_var, l_dist, l_reg = \ lanenet_discriminative_loss.discriminative_loss( pix_embedding, instance_label, 4, image_shape, 0.5, 3.0, 1.0, 1.0, 0.001) glog.info('compute_loss disc_loss:{:}'.format(disc_loss)) # 合并损失 l2_reg_loss = tf.constant(0.0, tf.float32) for vv in tf.trainable_variables(): if 'bn' in vv.name: continue else: l2_reg_loss = tf.add(l2_reg_loss, tf.nn.l2_loss(vv)) l2_reg_loss *= 0.001 total_loss = 0.5 * binary_segmenatation_loss + 0.5 * disc_loss + l2_reg_loss ret = { 'total_loss': total_loss, 'binary_seg_logits': decode_logits, 'instance_seg_logits': pix_embedding, 'binary_seg_loss': binary_segmenatation_loss, 'discriminative_loss': disc_loss } return ret def inference(self, input_tensor, name): """ :param input_tensor: :param name: :return: """ with tf.variable_scope(name): # 前向传播获取logits inference_ret = self._build_model(input_tensor=input_tensor, name='inference') # 计算二值分割损失函数 decode_logits = inference_ret['logits'] binary_seg_ret = tf.nn.softmax(logits=decode_logits) binary_seg_ret = tf.argmax(binary_seg_ret, axis=-1) # 计算像素嵌入 decode_deconv = inference_ret['deconv'] # 像素嵌入 pix_embedding = self.conv2d(inputdata=decode_deconv, out_channel=4, kernel_size=1, use_bias=False, name='pix_embedding_conv') pix_embedding = self.relu(inputdata=pix_embedding, name='pix_embedding_relu') return binary_seg_ret, pix_embedding if __name__ == '__main__': model = LaneNet(tf.constant('train', dtype=tf.string)) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input') binary_label = tf.placeholder(dtype=tf.int64, shape=[1, 256, 512, 1], name='label') instance_label = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 1], name='label') ret = model.compute_loss(input_tensor=input_tensor, binary_label=binary_label, instance_label=instance_label, name='loss') for vv in tf.trainable_variables(): if 'bn' in vv.name: continue print(vv.name)
--dirscriminative_loss.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import tensorflow as tf import glog def discriminative_loss_single( prediction, correct_label, feature_dim, label_shape, delta_v, delta_d, param_var, param_dist, param_reg): """ 论文equ(1)提到的实例分割损失函数 :param prediction: inference of network :param correct_label: instance label :param feature_dim: feature dimension of prediction :param label_shape: shape of label :param delta_v: cut off variance distance :param delta_d: cut off cluster distance :param param_var: weight for intra cluster variance :param param_dist: weight for inter cluster distances :param param_reg: weight regularization """ # 像素对齐为一行 correct_label = tf.reshape( correct_label, [ label_shape[1] * label_shape[0]]) reshaped_pred = tf.reshape( prediction, [ label_shape[1] * label_shape[0], feature_dim]) # 统计实例个数 unique_labels, unique_id, counts = tf.unique_with_counts(correct_label) counts = tf.cast(counts, tf.float32) num_instances = tf.size(unique_labels) glog.info('discriminative_loss_single counts:{:} num_instances:{:}'.format(counts,num_instances)) # 计算pixel embedding均值向量 segmented_sum = tf.unsorted_segment_sum( reshaped_pred, unique_id, num_instances) mu = tf.div(segmented_sum, tf.reshape(counts, (-1, 1))) mu_expand = tf.gather(mu, unique_id) # 计算公式的loss(var) distance = tf.norm(tf.subtract(mu_expand, reshaped_pred), axis=1) distance = tf.subtract(distance, delta_v) distance = tf.clip_by_value(distance, 0., distance) distance = tf.square(distance) l_var = tf.unsorted_segment_sum(distance, unique_id, num_instances) l_var = tf.div(l_var, counts) l_var = tf.reduce_sum(l_var) l_var = tf.divide(l_var, tf.cast(num_instances, tf.float32)) # 计算公式的loss(dist) mu_interleaved_rep = tf.tile(mu, [num_instances, 1]) mu_band_rep = tf.tile(mu, [1, num_instances]) mu_band_rep = tf.reshape( mu_band_rep, (num_instances * num_instances, feature_dim)) mu_diff = tf.subtract(mu_band_rep, mu_interleaved_rep) # 去除掩模上的零点 intermediate_tensor = tf.reduce_sum(tf.abs(mu_diff), axis=1) zero_vector = tf.zeros(1, dtype=tf.float32) bool_mask = tf.not_equal(intermediate_tensor, zero_vector) mu_diff_bool = tf.boolean_mask(mu_diff, bool_mask) mu_norm = tf.norm(mu_diff_bool, axis=1) mu_norm = tf.subtract(2. * delta_d, mu_norm) mu_norm = tf.clip_by_value(mu_norm, 0., mu_norm) mu_norm = tf.square(mu_norm) l_dist = tf.reduce_mean(mu_norm) # 计算原始Discriminative Loss论文中提到的正则项损失 l_reg = tf.reduce_mean(tf.norm(mu, axis=1)) # 合并损失按照原始Discriminative Loss论文中提到的参数合并 param_scale = 1. l_var = param_var * l_var l_dist = param_dist * l_dist l_reg = param_reg * l_reg loss = param_scale * (l_var + l_dist + l_reg) return loss, l_var, l_dist, l_reg def discriminative_loss(prediction, correct_label, feature_dim, image_shape, delta_v, delta_d, param_var, param_dist, param_reg): """ 按照论文的思想迭代计算loss损失 :return: discriminative loss and its three components """ def cond(label, batch, out_loss, out_var, out_dist, out_reg, i): return tf.less(i, tf.shape(batch)[0]) def body(label, batch, out_loss, out_var, out_dist, out_reg, i): disc_loss, l_var, l_dist, l_reg = discriminative_loss_single( prediction[i], correct_label[i], feature_dim, image_shape, delta_v, delta_d, param_var, param_dist, param_reg) out_loss = out_loss.write(i, disc_loss) out_var = out_var.write(i, l_var) out_dist = out_dist.write(i, l_dist) out_reg = out_reg.write(i, l_reg) return label, batch, out_loss, out_var, out_dist, out_reg, i + 1 # TensorArray is a data structure that support dynamic writing output_ta_loss = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True) output_ta_var = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True) output_ta_dist = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True) output_ta_reg = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True) _, _, out_loss_op, out_var_op, out_dist_op, out_reg_op, _ = tf.while_loop( cond, body, [ correct_label, prediction, output_ta_loss, output_ta_var, output_ta_dist, output_ta_reg, 0]) out_loss_op = out_loss_op.stack() out_var_op = out_var_op.stack() out_dist_op = out_dist_op.stack() out_reg_op = out_reg_op.stack() disc_loss = tf.reduce_mean(out_loss_op) l_var = tf.reduce_mean(out_var_op) l_dist = tf.reduce_mean(out_dist_op) l_reg = tf.reduce_mean(out_reg_op) return disc_loss, l_var, l_dist, l_reg
--postpostprecess.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ LaneNet模型后处理 """ import numpy as np import matplotlib.pyplot as plt import cv2 import glog try: from cv2 import cv2 except ImportError: pass class LaneNetPoseProcessor(object): """ """ def __init__(self): """ """ pass @staticmethod def _morphological_process(image, kernel_size=5): """ :param image: :param kernel_size: :return: """ if image.dtype is not np.uint8: image = np.array(image, np.uint8) glog.info("_morphological_process image shape len:{:d}".format(len(image.shape))) if len(image.shape) == 3: image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) glog.info("_morphological_process image shape len:{:d}".format(len(image.shape))) kernel = cv2.getStructuringElement(shape=cv2.MORPH_ELLIPSE, ksize=(kernel_size, kernel_size)) # close operation fille hole closing = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel, iterations=1) return closing @staticmethod def _connect_components_analysis(image): """ :param image: :return: """ glog.info("_connect_components_analysis image shape len:{:d}".format(len(image.shape))) if len(image.shape) == 3: gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: gray_image = image return cv2.connectedComponentsWithStats(gray_image, connectivity=8, ltype=cv2.CV_32S) def postprocess(self, image, minarea_threshold=15): """ :param image: :param minarea_threshold: 连通域分析阈值 :return: """ # 首先进行图像形态学运算 morphological_ret = self._morphological_process(image, kernel_size=5) glog.info("postprocess image shape len:{:d}".format(len(image.shape))) # 进行连通域分析 connect_components_analysis_ret = self._connect_components_analysis(image=morphological_ret) glog.info("postprocess connect_components_analysis_ret:{:}".format(connect_components_analysis_ret)) # 排序连通域并删除过小的连通域 labels = connect_components_analysis_ret[1] stats = connect_components_analysis_ret[2] glog.info("postprocess labels:{:}".format(labels)) glog.info("postprocess stats:{:}".format(stats)) for index, stat in enumerate(stats): if stat[4] <= minarea_threshold: idx = np.where(labels == index) morphological_ret[idx] = 0 return morphological_ret if __name__ == '__main__': processor = LaneNetPoseProcessor() image = cv2.imread('D:/Code/github/tf_lanenet/data/training_data_example/gt_image_binary/0000.png', cv2.IMREAD_UNCHANGED) #IMREAD_GRAYSCALE postprocess_ret = processor.postprocess(image) plt.figure('src') plt.imshow(image) plt.figure('post') plt.imshow(postprocess_ret) plt.show()
--cluster.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 实现LaneNet中实例分割的聚类部分 """ import numpy as np import glog as log import matplotlib.pyplot as plt from sklearn.cluster import MeanShift from sklearn.cluster import DBSCAN import time import warnings import cv2 import glog try: from cv2 import cv2 except ImportError: pass class LaneNetCluster(object): """ 实例分割聚类器 """ def __init__(self): """ """ self._color_map = [np.array([255, 0, 0]), np.array([0, 255, 0]), np.array([0, 0, 255]), np.array([125, 125, 0]), np.array([0, 125, 125]), np.array([125, 0, 125]), np.array([50, 100, 50]), np.array([100, 50, 100])] pass @staticmethod def _cluster(prediction, bandwidth): """ 实现论文SectionⅡ的cluster部分 :param prediction: :param bandwidth: :return: """ ms = MeanShift(bandwidth, bin_seeding=True) # log.info('开始Mean shift聚类 ...') tic = time.time() try: ms.fit(prediction) except ValueError as err: log.error(err) return 0, [], [] # log.info('Mean Shift耗时: {:.5f}s'.format(time.time() - tic)) labels = ms.labels_ cluster_centers = ms.cluster_centers_ num_clusters = cluster_centers.shape[0] # log.info('聚类簇个数为: {:d}'.format(num_clusters)) return num_clusters, labels, cluster_centers @staticmethod def _cluster_v2(prediction): """ dbscan cluster :param prediction: :return: """ db = DBSCAN(eps=0.7, min_samples=200).fit(prediction) db_labels = db.labels_ unique_labels = np.unique(db_labels) unique_labels = [tmp for tmp in unique_labels if tmp != -1] log.info('聚类簇个数为: {:d}'.format(len(unique_labels))) num_clusters = len(unique_labels) cluster_centers = db.components_ return num_clusters, db_labels, cluster_centers @staticmethod def _get_lane_area(binary_seg_ret, instance_seg_ret): """ 通过二值分割掩码图在实例分割图上获取所有车道线的特征向量 :param binary_seg_ret: :param instance_seg_ret: :return: """ idx = np.where(binary_seg_ret == 1) print("_get_lane_area idx:",idx) print("_get_lane_area idx len:",len(idx)) print("_get_lane_area idx len[0]:",len(idx[0])) print("_get_lane_area idx len[1]:",len(idx[1])) lane_embedding_feats = [] lane_coordinate = [] for i in range(len(idx[0])): lane_embedding_feats.append(instance_seg_ret[idx[0][i], idx[1][i]]) #print("_get_lane_area instance_seg_ret[idx[0][i], idx[1][i]]:",instance_seg_ret[idx[0][i], idx[1][i]]) lane_coordinate.append([idx[0][i], idx[1][i]]) #print("_get_lane_area idx[0][i]:",idx[0][i]," , idx[1][i]:", idx[1][i]) return np.array(lane_embedding_feats, np.float32), np.array(lane_coordinate, np.int64) @staticmethod def _thresh_coord(coord): """ 过滤实例车道线位置坐标点,假设车道线是连续的, 因此车道线点的坐标变换应该是平滑变化的不应该出现跳变 :param coord: [(x, y)] :return: """ pts_x = coord[:, 0] mean_x = np.mean(pts_x) idx = np.where(np.abs(pts_x - mean_x) < mean_x) return coord[idx[0]] @staticmethod def _lane_fit(lane_pts): """ 车道线多项式拟合 :param lane_pts: :return: """ if not isinstance(lane_pts, np.ndarray): lane_pts = np.array(lane_pts, np.float32) x = lane_pts[:, 0] y = lane_pts[:, 1] x_fit = [] y_fit = [] with warnings.catch_warnings(): warnings.filterwarnings('error') try: f1 = np.polyfit(x, y, 3) p1 = np.poly1d(f1) x_min = int(np.min(x)) x_max = int(np.max(x)) x_fit = [] for i in range(x_min, x_max + 1): x_fit.append(i) y_fit = p1(x_fit) except Warning as e: x_fit = x y_fit = y finally: return zip(x_fit, y_fit) def get_lane_mask(self, binary_seg_ret, instance_seg_ret): """ :param binary_seg_ret: :param instance_seg_ret: :return: """ lane_embedding_feats, lane_coordinate = self._get_lane_area(binary_seg_ret, instance_seg_ret) num_clusters, labels, cluster_centers = self._cluster(lane_embedding_feats, bandwidth=1.5) # 聚类簇超过八个则选择其中类内样本最多的八个聚类簇保留下来 if num_clusters > 8: cluster_sample_nums = [] for i in range(num_clusters): cluster_sample_nums.append(len(np.where(labels == i)[0])) sort_idx = np.argsort(-np.array(cluster_sample_nums, np.int64)) cluster_index = np.array(range(num_clusters))[sort_idx[0:4]] else: cluster_index = range(num_clusters) mask_image = np.zeros(shape=[binary_seg_ret.shape[0], binary_seg_ret.shape[1], 3], dtype=np.uint8) for index, i in enumerate(cluster_index): idx = np.where(labels == i) coord = lane_coordinate[idx] # coord = self._thresh_coord(coord) coord = np.flip(coord, axis=1) # coord = (coord[:, 0], coord[:, 1]) color = (int(self._color_map[index][0]), int(self._color_map[index][1]), int(self._color_map[index][2])) coord = np.array([coord]) cv2.polylines(img=mask_image, pts=coord, isClosed=False, color=color, thickness=2) # mask_image[coord] = color return mask_image if __name__ == '__main__': binary_seg_image = cv2.imread('D:/Code/github/tf_lanenet/data/training_data_example/gt_image_binary/0000.png', cv2.IMREAD_GRAYSCALE) print("binary_seg_image shape:",binary_seg_image.shape) binary_seg_image[np.where(binary_seg_image == 255)] = 1 print("binary_seg_image np.where(binary_seg_image == 255):",np.where(binary_seg_image == 255)) instance_seg_image = cv2.imread('D:/Code/github/tf_lanenet/data/training_data_example/gt_image_instance/0000.png', cv2.IMREAD_UNCHANGED) glog.info("__name__ instance_seg_image shape len:{:d}".format(len(instance_seg_image.shape))) instance_seg_image = cv2.cvtColor(instance_seg_image, cv2.COLOR_GRAY2BGR) glog.info("__name__ instance_seg_image shape len:{:d}".format(len(instance_seg_image.shape))) #print("instance_seg_image shape:",instance_seg_image.shape) ele_mex = np.max(instance_seg_image, axis=(0,1)) print("ele_mex:",ele_mex) for i in range(3): if ele_mex[i] == 0: scale = 1 else: scale = 255 / ele_mex[i] instance_seg_image[:, :, i] *= int(scale) embedding_image = np.array(instance_seg_image, np.uint8) cluster = LaneNetCluster() mask_image = cluster.get_lane_mask(binary_seg_ret=binary_seg_image,instance_seg_ret=instance_seg_image) det_img = embedding_image+mask_image plt.figure('det_img') plt.imshow(det_img[:, :, (2, 1, 0)]) #plt.figure('embedding') #plt.imshow(embedding_image[:, :, (2, 1, 0)]) #plt.figure('mask_image') #plt.imshow(mask_image[:, :, (2, 1, 0)]) plt.show()
--train_lane_scnn.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import argparse import math import os import os.path as ops import time import cv2 import glog as log import numpy as np import tensorflow as tf from config import global_config from lanenet_model import lanenet_merge_model from data_provider import lanenet_data_processor CFG = global_config.cfg VGG_MEAN = [103.939, 116.779, 123.68] def init_args(): """ :return: """ parser = argparse.ArgumentParser() parser.add_argument('--dataset_dir', type=str,default='data/datasets_culane', help='The training dataset dir path') parser.add_argument('--net', type=str, default='vgg', help='Which base net work to use') parser.add_argument('--weights_path', type=str,default='model/lanenet_culane_vgg_2019-02-02-14-05-16.ckpt-200000',help='The pretrained weights path') return parser.parse_args() def minmax_scale(input_arr): """ :param input_arr: :return: """ min_val = np.min(input_arr) max_val = np.max(input_arr) output_arr = (input_arr - min_val) * 255.0 / (max_val - min_val) return output_arr def train_net(dataset_dir, weights_path=None, net_flag='vgg'): """ :param dataset_dir: :param net_flag: choose which base network to use :param weights_path: :return: """ train_dataset_file = ops.join(dataset_dir, 'train.txt') val_dataset_file = ops.join(dataset_dir, 'val.txt') print('train_dataset_file:',train_dataset_file) print('val_dataset_file:',val_dataset_file) assert ops.exists(train_dataset_file) train_dataset = lanenet_data_processor.DataSet(train_dataset_file) val_dataset = lanenet_data_processor.DataSet(val_dataset_file) with tf.device('/gpu:1'): input_tensor = tf.placeholder(dtype=tf.float32, shape=[CFG.TRAIN.BATCH_SIZE, CFG.TRAIN.IMG_HEIGHT, CFG.TRAIN.IMG_WIDTH, 3], name='input_tensor') binary_label_tensor = tf.placeholder(dtype=tf.int64, shape=[CFG.TRAIN.BATCH_SIZE, CFG.TRAIN.IMG_HEIGHT, CFG.TRAIN.IMG_WIDTH, 1], name='binary_input_label') instance_label_tensor = tf.placeholder(dtype=tf.float32, shape=[CFG.TRAIN.BATCH_SIZE, CFG.TRAIN.IMG_HEIGHT, CFG.TRAIN.IMG_WIDTH], name='instance_input_label') phase = tf.placeholder(dtype=tf.string, shape=None, name='net_phase') net = lanenet_merge_model.LaneNet(net_flag=net_flag, phase=phase) # calculate the loss compute_ret = net.compute_loss(input_tensor=input_tensor, binary_label=binary_label_tensor, instance_label=instance_label_tensor, name='lanenet_model') total_loss = compute_ret['total_loss'] binary_seg_loss = compute_ret['binary_seg_loss'] disc_loss = compute_ret['discriminative_loss'] pix_embedding = compute_ret['instance_seg_logits'] # calculate the accuracy out_logits = compute_ret['binary_seg_logits'] out_logits = tf.nn.softmax(logits=out_logits) out_logits_out = tf.argmax(out_logits, axis=-1) #out = tf.argmax(out_logits, axis=-1) #out = tf.expand_dims(out, axis=-1) out = tf.expand_dims(out_logits_out,axis=-1) idx = tf.where(tf.equal(binary_label_tensor, 1)) pix_cls_ret = tf.gather_nd(out, idx) accuracy = tf.count_nonzero(pix_cls_ret) accuracy = tf.divide(accuracy, tf.cast(tf.shape(pix_cls_ret)[0], tf.int64)) global_step = tf.Variable(0, trainable=False) learning_rate = tf.train.exponential_decay(CFG.TRAIN.LEARNING_RATE, global_step, 100000, 0.1, staircase=True) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): optimizer = tf.train.MomentumOptimizer( learning_rate=learning_rate, momentum=0.9).minimize(loss=total_loss, var_list=tf.trainable_variables(), global_step=global_step) # Set tf saver saver = tf.train.Saver() model_save_dir = 'model/lanenet_culane' if not ops.exists(model_save_dir): os.makedirs(model_save_dir) train_start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) model_name = 'lanenet_culane_{:s}_{:s}.ckpt'.format(net_flag, str(train_start_time)) model_save_path = ops.join(model_save_dir, model_name) # Set tf summary tboard_save_path = 'tboard/lanenet_culane/{:s}'.format(net_flag) if not ops.exists(tboard_save_path): os.makedirs(tboard_save_path) train_cost_scalar = tf.summary.scalar(name='train_cost', tensor=total_loss) val_cost_scalar = tf.summary.scalar(name='val_cost', tensor=total_loss) train_accuracy_scalar = tf.summary.scalar(name='train_accuracy', tensor=accuracy) val_accuracy_scalar = tf.summary.scalar(name='val_accuracy', tensor=accuracy) train_binary_seg_loss_scalar = tf.summary.scalar(name='train_binary_seg_loss', tensor=binary_seg_loss) val_binary_seg_loss_scalar = tf.summary.scalar(name='val_binary_seg_loss', tensor=binary_seg_loss) train_instance_seg_loss_scalar = tf.summary.scalar(name='train_instance_seg_loss', tensor=disc_loss) val_instance_seg_loss_scalar = tf.summary.scalar(name='val_instance_seg_loss', tensor=disc_loss) learning_rate_scalar = tf.summary.scalar(name='learning_rate', tensor=learning_rate) train_merge_summary_op = tf.summary.merge([train_accuracy_scalar, train_cost_scalar, learning_rate_scalar, train_binary_seg_loss_scalar, train_instance_seg_loss_scalar]) val_merge_summary_op = tf.summary.merge([val_accuracy_scalar, val_cost_scalar, val_binary_seg_loss_scalar, val_instance_seg_loss_scalar]) # Set sess configuration sess_config = tf.ConfigProto(allow_soft_placement=True) sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TRAIN.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.TRAIN.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) summary_writer = tf.summary.FileWriter(tboard_save_path) summary_writer.add_graph(sess.graph) # Set the training parameters train_epochs = CFG.TRAIN.EPOCHS log.info('Global configuration is as follows:') log.info(CFG) with sess.as_default(): tf.train.write_graph(graph_or_graph_def=sess.graph, logdir='', name='{:s}/lanenet_model.pb'.format(model_save_dir)) if weights_path is None: log.info('Training from scratch') init = tf.global_variables_initializer() sess.run(init) else: log.info('Restore model from last model checkpoint {:s}'.format(weights_path)) saver.restore(sess=sess, save_path=weights_path) # 加载预训练参数 log.info('jim.chen train_net net_flag:',net_flag) if net_flag == 'vgg' and weights_path is None: pretrained_weights = np.load( './data/vgg16.npy', encoding='latin1').item() log.info('jim.chen train_net net_flag is 1vgg') for vv in tf.trainable_variables(): weights_key = vv.name.split('/')[-3] try: weights = pretrained_weights[weights_key][0] _op = tf.assign(vv, weights) sess.run(_op) except Exception as e: continue train_cost_time_mean = [] val_cost_time_mean = [] for epoch in range(train_epochs): # training part t_start = time.time() with tf.device('/cpu:0'): gt_imgs, binary_gt_labels,instance_gt_labels = train_dataset.next_batch(CFG.TRAIN.BATCH_SIZE) gt_imgs = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_LINEAR) for tmp in gt_imgs] gt_imgs = [tmp - VGG_MEAN for tmp in gt_imgs] binary_gt_labels = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_NEAREST) for tmp in binary_gt_labels] binary_gt_labels = [np.expand_dims(tmp, axis=-1) for tmp in binary_gt_labels] instance_gt_labels = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_NEAREST) for tmp in instance_gt_labels] phase_train = 'train' _, c, train_accuracy, train_summary, binary_loss, instance_loss, embedding, binary_seg_img = \ sess.run([optimizer, total_loss, accuracy, train_merge_summary_op, binary_seg_loss, disc_loss, pix_embedding, out_logits_out], feed_dict={input_tensor: gt_imgs, binary_label_tensor: binary_gt_labels, instance_label_tensor: instance_gt_labels, phase: phase_train}) if math.isnan(c) or math.isnan(instance_loss) or math.isnan(binary_loss): log.error('cost is: {:.5f}'.format(c)) log.error('binary cost is: {:.5f}'.format(binary_loss)) log.error('instance cost is: {:.5f}'.format(instance_loss)) cv2.imwrite('nan_image.png', gt_imgs[0] + VGG_MEAN) cv2.imwrite('nan_instance_label.png', instance_gt_labels[0]) cv2.imwrite('nan_binary_label.png', binary_gt_labels[0] * 255) return if epoch % 100 == 0: cv2.imwrite('image.png', gt_imgs[0] + VGG_MEAN) cv2.imwrite('binary_label.png', binary_gt_labels[0] * 255) cv2.imwrite('instance_label.png', instance_gt_labels[0]) cv2.imwrite('binary_seg_img.png', binary_seg_img[0] * 255) for i in range(4): embedding[0][:, :, i] = minmax_scale(embedding[0][:, :, i]) embedding_image = np.array(embedding[0], np.uint8) cv2.imwrite('embedding.png', embedding_image) cost_time = time.time() - t_start train_cost_time_mean.append(cost_time) summary_writer.add_summary(summary=train_summary, global_step=epoch) # validation part with tf.device('/cpu:0'): gt_imgs_val, binary_gt_labels_val, instance_gt_labels_val \ = val_dataset.next_batch(CFG.TRAIN.VAL_BATCH_SIZE) gt_imgs_val = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_LINEAR) for tmp in gt_imgs_val] gt_imgs_val = [tmp - VGG_MEAN for tmp in gt_imgs_val] binary_gt_labels_val = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp) for tmp in binary_gt_labels_val] binary_gt_labels_val = [np.expand_dims(tmp, axis=-1) for tmp in binary_gt_labels_val] instance_gt_labels_val = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_NEAREST) for tmp in instance_gt_labels_val] phase_val = 'test' t_start_val = time.time() c_val, val_summary, val_accuracy, val_binary_seg_loss, val_instance_seg_loss = \ sess.run([total_loss, val_merge_summary_op, accuracy, binary_seg_loss, disc_loss], feed_dict={input_tensor: gt_imgs_val, binary_label_tensor: binary_gt_labels_val, instance_label_tensor: instance_gt_labels_val, phase: phase_val}) if epoch % 100 == 0: cv2.imwrite('test_image.png', gt_imgs_val[0] + VGG_MEAN) summary_writer.add_summary(val_summary, global_step=epoch) cost_time_val = time.time() - t_start_val val_cost_time_mean.append(cost_time_val) if epoch % CFG.TRAIN.DISPLAY_STEP == 0: log.info('Epoch: {:d} total_loss= {:6f} binary_seg_loss= {:6f} instance_seg_loss= {:6f} accuracy= {:6f}' ' mean_cost_time= {:5f}s '. format(epoch + 1, c, binary_loss, instance_loss, train_accuracy, np.mean(train_cost_time_mean))) train_cost_time_mean.clear() if epoch % CFG.TRAIN.TEST_DISPLAY_STEP == 0: log.info('Epoch_Val: {:d} total_loss= {:6f} binary_seg_loss= {:6f} ' 'instance_seg_loss= {:6f} accuracy= {:6f} ' 'mean_cost_time= {:5f}s '. format(epoch + 1, c_val, val_binary_seg_loss, val_instance_seg_loss, val_accuracy, np.mean(val_cost_time_mean))) val_cost_time_mean.clear() if epoch % 2000 == 0: saver.save(sess=sess, save_path=model_save_path, global_step=epoch) sess.close() return if __name__ == '__main__': # init args args = init_args() # train lanenet train_net(args.dataset_dir, args.weights_path, net_flag=args.net)
以下是模型训练过程中生成的文件夹：
./summary
./figure
./checkpoint
在主目录下，执行python train_lanenet_scnn.py,没有问题的话，可以开始训练了．．．
下一篇：
1、数据源：包括所有原始数据，分组后的数据；
2、数据预处理：包括数据的准备，数据的导入，数据的提取，数据的分组（训练与测试）；
3、配置文件：包括各种参数与超参数，如：训练周期，训练步长，批量数据，学习率，卷积核大小，全连接大小，训练模型存放路径（checkpoint），摘要存放路径(summary)等；
4、基础网络：包括基本的网络组件，基础网，
5、训练主文件：主入口，用于搭建生成图（graph）,会话（sess）,数据导入模型训练，ＧＰＵ配置，训练过程打印等
三、代码结构
以下为原始文件夹：
./data
-- ./InstanceSegmentationClass
-- ./JPEGImages
-- ./SegmentationClass
-- datasets_gen_culane.py 用于从上面三个图片目录生成list.txt,train.txt,test.txt
# coding=utf-8 #create date:12/5/2018 #modified date:2/12/2019 #author:jim.chen import os import glob import random import math import cv2 import numpy as np def gen_list_txt(rela_dir,img_dir,img_seg_dir,img_inst_dir): #cwd = os.getcwd() #print("gen_list_txt cwd:",cwd) list_txt = "list.txt" png_glob = img_seg_dir+'/*.png' png_list_path = glob.glob(png_glob) png_list=[] print("gen_list_txt png_list_path:",png_list_path) with open(list_txt,"w") as w_f: for png in png_list_path: path,name = os.path.splitext(os.path.basename(png)) print("path:",path) w_f.write(rela_dir+img_dir+'/'+path+'.jpg'+' '+rela_dir+png+' '+rela_dir+img_inst_dir+'/'+path+'.png'+'\n') w_f.close() with open(list_txt,"r") as r_f: for each_line in r_f: png_list.append(each_line) png_list.sort() print("gen_list_txt len(png_list):",len(png_list)) train=random.sample(png_list,int(math.floor(len(png_list)*9/10))) train.sort() print("gen_list_txt train:",train) val=list(set(png_list).difference(set(train))) print("gen_list_txt val:",val) enum_train_val=['train','val'] for item in enum_train_val: with open(item+'.txt','w') as w1_f: for num_item in eval(item): print("gen_list_txt num_item:",num_item) w1_f.write(num_item) w1_f.close() def sync_gt_2_img(img_dir,img_seg_dir,img_inst_dir): cwd = os.getcwd() print("sync_gt_2_img img_dir:",img_dir," img_seg_dir:",img_seg_dir," img_inst_dir:",img_inst_dir) img_full_dir = cwd + '/' +img_dir img_seg_full_dir = cwd + '/' +img_seg_dir img_inst_full_dir = cwd + '/' +img_inst_dir img_list = os.listdir(img_full_dir) for img in img_list: img_basename = os.path.splitext(img)[0] print("sync_gt_2_img img_basename:",img_basename) img_full_path = img_full_dir + '/'+ img img_seg_full_path = img_seg_full_dir + '/'+img_basename +'.png' #print("sync_gt_2_img img_seg_full_path:",img_seg_full_path) img_inst_full_path = img_inst_full_dir + '/'+img_basename +'.png' #print("sync_gt_2_img img_inst_full_path:",img_inst_full_path) if not os.path.exists(img_inst_full_path): print("sync_gt_2_img not os.path.exists(img_seg_full_path)") if os.path.exists(img_full_path): os.remove(img_full_path) if os.path.exists(img_full_path): os.remove(img_full_path) def sync_seg_2_inst(img_seg_dir,img_inst_dir): cwd = os.getcwd() print("sync_seg_2_inst img_seg_dir:",img_seg_dir," img_inst_dir:",img_inst_dir) img_seg_full_dir = cwd + '/' +img_seg_dir img_inst_full_dir = cwd + '/' +img_inst_dir img_list = os.listdir(img_seg_dir) for img in img_list: img_basename = os.path.splitext(img)[0] print("sync_seg_2_inst img_basename:",img_basename) img_seg_full_path = img_seg_full_dir + '/'+img_basename +'.jpg' img_inst_full_path = img_inst_full_dir + '/'+img_basename +'.png' if not os.path.exists(img_inst_full_path): if os.path.exists(img_seg_full_path): print("sync_seg_2_inst os.remove(img_seg_full_path):",img_seg_full_path) os.remove(img_seg_full_path) def gen_seg_color(img_inst_dir,img_seg_dir): cwd = os.getcwd() inPath = os.path.join(cwd,img_inst_dir) print(inPath) outPath=os.path.join(cwd,img_seg_dir) inPathDir = os.listdir(inPath) if not os.path.exists(outPath): os.makedirs(outPath) for l,file_name in enumerate(inPathDir): img_instance = cv2.imread(os.path.join(inPath,file_name)) h,w,c = img_instance.shape print("l:",l," img_instance.shape:",img_instance.shape) img_instance_new = np.zeros((h, w, c), dtype=np.uint8) for i in range(0,h): for j in range(0,w): #print(img_instance[i][j]) if img_instance[i][j][0] != 0: img_instance_new[i][j] = [255,255,255] img_instance_gray = cv2.cvtColor(img_instance_new, cv2.COLOR_BGR2GRAY) cv2.imwrite(os.path.join(outPath,file_name), img_instance_gray) print("generate segment finished!") def gen_inst_color(img_inst_dir): cwd = os.getcwd() inPath = os.path.join(cwd,img_inst_dir) print(inPath) outPath=os.path.join(cwd,"img_inst_new") inPathDir = os.listdir(inPath) if not os.path.exists(outPath): os.makedirs(outPath) for l,file_name in enumerate(inPathDir): img_instance = cv2.imread(os.path.join(inPath,file_name)) h,w,c = img_instance.shape print("l:",l," img_instance.shape:",img_instance.shape) img_instance_new = np.zeros((h, w, c), dtype=np.uint8) for i in range(0,h): for j in range(0,w): #print(img_instance[i][j]) if img_instance[i][j][0] == 2: img_instance_new[i][j] = [20,20,20] elif img_instance[i][j][0] == 3: img_instance_new[i][j] = [70,70,70] elif img_instance[i][j][0] == 4: img_instance_new[i][j] = [120,120,120] elif img_instance[i][j][0] == 5: img_instance_new[i][j] = [170,170,170] elif img_instance[i][j][0] == 6: img_instance_new[i][j] = [220,220,220] img_instance_gray = cv2.cvtColor(img_instance_new, cv2.COLOR_BGR2GRAY) cv2.imwrite(os.path.join(outPath,file_name), img_instance_gray) print("generate instance finished!") def detect_invalid_img(img_path): img_instance = cv2.imread(img_path) h,w,c = img_instance.shape print("detect_invalid_img img_instance.shape:",img_instance.shape) for i in range(0,h): for j in range(0,w): if img_instance[i][j][0] != 0: return False return True def filter_invalid_img(img_test_dir,img_seg_dir,img_inst_dir): cwd = os.getcwd() inPath = os.path.join(cwd,img_test_dir) inPathDir = os.listdir(inPath) print("filter_invalid_img inPathDir:",inPathDir) for l,file_name in enumerate(inPathDir): img_path = os.path.join(inPath,file_name) isdel = detect_invalid_img(img_path) if isdel: print("filter_invalid_img isdel:",isdel) os.remove(os.path.join(inPath,file_name)) def main(): print("main begin") rela_dir = "data/datasets_culane_all/" img_dir = "image" img_seg_dir = "gt_image_binary" img_inst_dir = "gt_image_instance" gen_list_txt(rela_dir,img_dir,img_seg_dir,img_inst_dir) #sync_gt_2_img(img_dir,img_seg_dir,img_inst_dir) #gen_seg_color(img_inst_dir,img_seg_dir) #gen_inst_color(img_inst_dir) #filter_invalid_img(img_dir,img_seg_dir,img_inst_dir) #sync_seg_2_inst(img_dir,img_inst_dir) print("main end") if __name__ == '__main__': main()
--list.txt
data/datasets_culane/JPEGImages/0000.jpg data/datasets_culane/SegmentationClass\0000.png data/datasets_culane/InstanceSegmentationClass/0000.png data/datasets_culane/JPEGImages/0001.jpg data/datasets_culane/SegmentationClass\0001.png data/datasets_culane/InstanceSegmentationClass/0001.png data/datasets_culane/JPEGImages/0002.jpg data/datasets_culane/SegmentationClass\0002.png data/datasets_culane/InstanceSegmentationClass/0002.png data/datasets_culane/JPEGImages/0003.jpg data/datasets_culane/SegmentationClass\0003.png data/datasets_culane/InstanceSegmentationClass/0003.png data/datasets_culane/JPEGImages/0004.jpg data/datasets_culane/SegmentationClass\0004.png data/datasets_culane/InstanceSegmentationClass/0004.png data/datasets_culane/JPEGImages/0005.jpg data/datasets_culane/SegmentationClass\0005.png data/datasets_culane/InstanceSegmentationClass/0005.png
--train.txt
data/datasets_culane/JPEGImages/0000.jpg data/datasets_culane/SegmentationClass\0000.png data/datasets_culane/InstanceSegmentationClass/0000.png data/datasets_culane/JPEGImages/0001.jpg data/datasets_culane/SegmentationClass\0001.png data/datasets_culane/InstanceSegmentationClass/0001.png data/datasets_culane/JPEGImages/0002.jpg data/datasets_culane/SegmentationClass\0002.png data/datasets_culane/InstanceSegmentationClass/0002.png data/datasets_culane/JPEGImages/0004.jpg data/datasets_culane/SegmentationClass\0004.png data/datasets_culane/InstanceSegmentationClass/0004.png
--val.txt
data/datasets_culane/JPEGImages/0005.jpg data/datasets_culane/SegmentationClass\0005.png data/datasets_culane/InstanceSegmentationClass/0005.png
data/datasets_culane/JPEGImages/0003.jpg data/datasets_culane/SegmentationClass\0003.png data/datasets_culane/InstanceSegmentationClass/0003.png
./data_provider
--data_processor.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import os.path as ops import cv2 import numpy as np try: from cv2 import cv2 except ImportError: pass class DataSet(object): def __init__(self, dataset_info_file): self._gt_img_list, self._gt_label_binary_list, \ self._gt_label_instance_list = self._init_dataset(dataset_info_file) self._random_dataset() self._next_batch_loop_count = 0 def _init_dataset(self, dataset_info_file): gt_img_list = [] gt_label_binary_list = [] gt_label_instance_list = [] assert ops.exists(dataset_info_file), '{:s} not exist'.format(dataset_info_file) with open(dataset_info_file, 'r') as file: for _info in file: info_tmp = _info.strip(' ').split() gt_img_list.append(info_tmp[0]) gt_label_binary_list.append(info_tmp[1]) gt_label_instance_list.append(info_tmp[2]) return gt_img_list, gt_label_binary_list, gt_label_instance_list def _random_dataset(self): assert len(self._gt_img_list) == len(self._gt_label_binary_list) == len(self._gt_label_instance_list) random_idx = np.random.permutation(len(self._gt_img_list)) new_gt_img_list = [] new_gt_label_binary_list = [] new_gt_label_instance_list = [] for index in random_idx: new_gt_img_list.append(self._gt_img_list[index]) new_gt_label_binary_list.append(self._gt_label_binary_list[index]) new_gt_label_instance_list.append(self._gt_label_instance_list[index]) self._gt_img_list = new_gt_img_list self._gt_label_binary_list = new_gt_label_binary_list self._gt_label_instance_list = new_gt_label_instance_list def next_batch(self, batch_size): """ :param batch_size: :return: """ assert len(self._gt_label_binary_list) == len(self._gt_label_instance_list) \ == len(self._gt_img_list) idx_start = batch_size * self._next_batch_loop_count idx_end = batch_size * self._next_batch_loop_count + batch_size if idx_start == 0 and idx_end > len(self._gt_label_binary_list): raise ValueError('Batch size cant be more than total numbers') if idx_end > len(self._gt_label_binary_list): self._random_dataset() self._next_batch_loop_count = 0 return self.next_batch(batch_size) else: gt_img_list = self._gt_img_list[idx_start:idx_end] gt_label_binary_list = self._gt_label_binary_list[idx_start:idx_end] gt_label_instance_list = self._gt_label_instance_list[idx_start:idx_end] gt_imgs = [] gt_labels_binary = [] gt_labels_instance = [] for gt_img_path in gt_img_list: gt_imgs.append(cv2.imread(gt_img_path, cv2.IMREAD_COLOR)) for gt_label_path in gt_label_binary_list: label_img = cv2.imread(gt_label_path, cv2.IMREAD_COLOR) label_binary = np.zeros([label_img.shape[0], label_img.shape[1]], dtype=np.uint8) idx = np.where((label_img[:, :, :] != [0, 0, 0]).all(axis=2)) label_binary[idx] = 1 gt_labels_binary.append(label_binary) for gt_label_path in gt_label_instance_list: label_img = cv2.imread(gt_label_path, cv2.IMREAD_UNCHANGED) gt_labels_instance.append(label_img) self._next_batch_loop_count += 1 return gt_imgs, gt_labels_binary, gt_labels_instance if __name__ == '__main__': val = DataSet('/media/baidu/Data/Semantic_Segmentation/TUSimple_Lane_Detection/training/val.txt') b1, b2, b3 = val.next_batch(50) c1, c2, c3 = val.next_batch(50) dd, d2, d3 = val.next_batch(50)
./config
--global_config.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- from easydict import EasyDict as edict __C = edict() # Consumers can get config by: from config import cfg cfg = __C # Train options __C.TRAIN = edict() # Set the shadownet training epochs __C.TRAIN.EPOCHS = 200010 # Set the display step __C.TRAIN.DISPLAY_STEP = 1 # Set the test display step during training process __C.TRAIN.TEST_DISPLAY_STEP = 1000 # Set the momentum parameter of the optimizer __C.TRAIN.MOMENTUM = 0.9 # Set the initial learning rate __C.TRAIN.LEARNING_RATE = 0.0005 # Set the GPU resource used during training process __C.TRAIN.GPU_MEMORY_FRACTION = 0.85 # Set the GPU allow growth parameter during tensorflow training process __C.TRAIN.TF_ALLOW_GROWTH = True # Set the shadownet training batch size __C.TRAIN.BATCH_SIZE = 1 # Set the shadownet validation batch size __C.TRAIN.VAL_BATCH_SIZE = 1 # Set the learning rate decay steps __C.TRAIN.LR_DECAY_STEPS = 410000 # Set the learning rate decay rate __C.TRAIN.LR_DECAY_RATE = 0.1 # Set the class numbers __C.TRAIN.CLASSES_NUMS = 2 # Set the image height __C.TRAIN.IMG_HEIGHT = 256 # Set the image width __C.TRAIN.IMG_WIDTH = 512 # Test options __C.TEST = edict() # Set the GPU resource used during testing process __C.TEST.GPU_MEMORY_FRACTION = 0.8 # Set the GPU allow growth parameter during tensorflow testing process __C.TEST.TF_ALLOW_GROWTH = True # Set the test batch size __C.TEST.BATCH_SIZE = 1
./encoder_decoder_model
--cnn_basenet.py
#!/usr/bin/env python # -*- coding: utf-8 -*- """ The base convolution neural networks mainly implement some useful cnn functions """ import tensorflow as tf import numpy as np class CNNBaseModel(object): """ Base model for other specific cnn ctpn_models """ def __init__(self): pass @staticmethod def conv2d(inputdata, out_channel, kernel_size, padding='SAME', stride=1, w_init=None, b_init=None, split=1, use_bias=True, data_format='NHWC', name=None): with tf.variable_scope(name): in_shape = inputdata.get_shape().as_list() channel_axis = 3 if data_format == 'NHWC' else 1 in_channel = in_shape[channel_axis] assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!" assert in_channel % split == 0 assert out_channel % split == 0 padding = padding.upper() if isinstance(kernel_size, list): filter_shape = [kernel_size[0], kernel_size[1]] + [in_channel / split, out_channel] else: filter_shape = [kernel_size, kernel_size] + [in_channel / split, out_channel] if isinstance(stride, list): strides = [1, stride[0], stride[1], 1] if data_format == 'NHWC' \ else [1, 1, stride[0], stride[1]] else: strides = [1, stride, stride, 1] if data_format == 'NHWC' \ else [1, 1, stride, stride] if w_init is None: w_init = tf.contrib.layers.variance_scaling_initializer() if b_init is None: b_init = tf.constant_initializer() w = tf.get_variable('W', filter_shape, initializer=w_init) b = None if use_bias: b = tf.get_variable('b', [out_channel], initializer=b_init) if split == 1: conv = tf.nn.conv2d(inputdata, w, strides, padding, data_format=data_format) else: inputs = tf.split(inputdata, split, channel_axis) kernels = tf.split(w, split, 3) outputs = [tf.nn.conv2d(i, k, strides, padding, data_format=data_format) for i, k in zip(inputs, kernels)] conv = tf.concat(outputs, channel_axis) ret = tf.identity(tf.nn.bias_add(conv, b, data_format=data_format) if use_bias else conv, name=name) return ret @staticmethod def relu(inputdata, name=None): return tf.nn.relu(features=inputdata, name=name) @staticmethod def sigmoid(inputdata, name=None): return tf.nn.sigmoid(x=inputdata, name=name) @staticmethod def maxpooling(inputdata, kernel_size, stride=None, padding='VALID', data_format='NHWC', name=None): padding = padding.upper() if stride is None: stride = kernel_size if isinstance(kernel_size, list): kernel = [1, kernel_size[0], kernel_size[1], 1] if data_format == 'NHWC' else \ [1, 1, kernel_size[0], kernel_size[1]] else: kernel = [1, kernel_size, kernel_size, 1] if data_format == 'NHWC' \ else [1, 1, kernel_size, kernel_size] if isinstance(stride, list): strides = [1, stride[0], stride[1], 1] if data_format == 'NHWC' \ else [1, 1, stride[0], stride[1]] else: strides = [1, stride, stride, 1] if data_format == 'NHWC' \ else [1, 1, stride, stride] return tf.nn.max_pool(value=inputdata, ksize=kernel, strides=strides, padding=padding, data_format=data_format, name=name) @staticmethod def avgpooling(inputdata, kernel_size, stride=None, padding='VALID', data_format='NHWC', name=None): if stride is None: stride = kernel_size kernel = [1, kernel_size, kernel_size, 1] if data_format == 'NHWC' \ else [1, 1, kernel_size, kernel_size] strides = [1, stride, stride, 1] if data_format == 'NHWC' else [1, 1, stride, stride] return tf.nn.avg_pool(value=inputdata, ksize=kernel, strides=strides, padding=padding, data_format=data_format, name=name) @staticmethod def globalavgpooling(inputdata, data_format='NHWC', name=None): assert inputdata.shape.ndims == 4 assert data_format in ['NHWC', 'NCHW'] axis = [1, 2] if data_format == 'NHWC' else [2, 3] return tf.reduce_mean(input_tensor=inputdata, axis=axis, name=name) @staticmethod def layernorm(inputdata, epsilon=1e-5, use_bias=True, use_scale=True, data_format='NHWC', name=None): shape = inputdata.get_shape().as_list() ndims = len(shape) assert ndims in [2, 4] mean, var = tf.nn.moments(inputdata, list(range(1, len(shape))), keep_dims=True) if data_format == 'NCHW': channnel = shape[1] new_shape = [1, channnel, 1, 1] else: channnel = shape[-1] new_shape = [1, 1, 1, channnel] if ndims == 2: new_shape = [1, channnel] if use_bias: beta = tf.get_variable('beta', [channnel], initializer=tf.constant_initializer()) beta = tf.reshape(beta, new_shape) else: beta = tf.zeros([1] * ndims, name='beta') if use_scale: gamma = tf.get_variable('gamma', [channnel], initializer=tf.constant_initializer(1.0)) gamma = tf.reshape(gamma, new_shape) else: gamma = tf.ones([1] * ndims, name='gamma') return tf.nn.batch_normalization(inputdata, mean, var, beta, gamma, epsilon, name=name) @staticmethod def instancenorm(inputdata, epsilon=1e-5, data_format='NHWC', use_affine=True, name=None): shape = inputdata.get_shape().as_list() if len(shape) != 4: raise ValueError("Input data of instancebn layer has to be 4D tensor") if data_format == 'NHWC': axis = [1, 2] ch = shape[3] new_shape = [1, 1, 1, ch] else: axis = [2, 3] ch = shape[1] new_shape = [1, ch, 1, 1] if ch is None: raise ValueError("Input of instancebn require known channel!") mean, var = tf.nn.moments(inputdata, axis, keep_dims=True) if not use_affine: return tf.divide(inputdata - mean, tf.sqrt(var + epsilon), name='output') beta = tf.get_variable('beta', [ch], initializer=tf.constant_initializer()) beta = tf.reshape(beta, new_shape) gamma = tf.get_variable('gamma', [ch], initializer=tf.constant_initializer(1.0)) gamma = tf.reshape(gamma, new_shape) return tf.nn.batch_normalization(inputdata, mean, var, beta, gamma, epsilon, name=name) @staticmethod def dropout(inputdata, keep_prob, noise_shape=None, name=None): return tf.nn.dropout(inputdata, keep_prob=keep_prob, noise_shape=noise_shape, name=name) @staticmethod def fullyconnect(inputdata, out_dim, w_init=None, b_init=None, use_bias=True, name=None): shape = inputdata.get_shape().as_list()[1:] if None not in shape: inputdata = tf.reshape(inputdata, [-1, int(np.prod(shape))]) else: inputdata = tf.reshape(inputdata, tf.stack([tf.shape(inputdata)[0], -1])) if w_init is None: w_init = tf.contrib.layers.variance_scaling_initializer() if b_init is None: b_init = tf.constant_initializer() ret = tf.layers.dense(inputs=inputdata, activation=lambda x: tf.identity(x, name='output'), use_bias=use_bias, name=name, kernel_initializer=w_init, bias_initializer=b_init, trainable=True, units=out_dim) return ret @staticmethod def layerbn(inputdata, is_training, name): return tf.layers.batch_normalization(inputs=inputdata, training=is_training, name=name) @staticmethod def squeeze(inputdata, axis=None, name=None): return tf.squeeze(input=inputdata, axis=axis, name=name) @staticmethod def deconv2d(inputdata, out_channel, kernel_size, padding='SAME', stride=1, w_init=None, b_init=None, use_bias=True, activation=None, data_format='channels_last', trainable=True, name=None): with tf.variable_scope(name): in_shape = inputdata.get_shape().as_list() channel_axis = 3 if data_format == 'channels_last' else 1 in_channel = in_shape[channel_axis] assert in_channel is not None, "[Deconv2D] Input cannot have unknown channel!" padding = padding.upper() if w_init is None: w_init = tf.contrib.layers.variance_scaling_initializer() if b_init is None: b_init = tf.constant_initializer() ret = tf.layers.conv2d_transpose(inputs=inputdata, filters=out_channel, kernel_size=kernel_size, strides=stride, padding=padding, data_format=data_format, activation=activation, use_bias=use_bias, kernel_initializer=w_init, bias_initializer=b_init, trainable=trainable, name=name) return ret @staticmethod def dilation_conv(input_tensor, k_size, out_dims, rate, padding='SAME', w_init=None, b_init=None, use_bias=False, name=None): with tf.variable_scope(name): in_shape = input_tensor.get_shape().as_list() in_channel = in_shape[3] assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!" padding = padding.upper() if isinstance(k_size, list): filter_shape = [k_size[0], k_size[1]] + [in_channel, out_dims] else: filter_shape = [k_size, k_size] + [in_channel, out_dims] if w_init is None: w_init = tf.contrib.layers.variance_scaling_initializer() if b_init is None: b_init = tf.constant_initializer() w = tf.get_variable('W', filter_shape, initializer=w_init) b = None if use_bias: b = tf.get_variable('b', [out_dims], initializer=b_init) conv = tf.nn.atrous_conv2d(value=input_tensor, filters=w, rate=rate, padding=padding, name='dilation_conv') if use_bias: ret = tf.add(conv, b) else: ret = conv return ret @staticmethod def spatial_dropout(input_tensor, keep_prob, is_training, name, seed=1234): tf.set_random_seed(seed=seed) def f1(): with tf.variable_scope(name): return input_tensor def f2(): with tf.variable_scope(name): num_feature_maps = [tf.shape(input_tensor)[0], tf.shape(input_tensor)[3]] random_tensor = keep_prob random_tensor += tf.random_uniform(num_feature_maps, seed=seed, dtype=input_tensor.dtype) binary_tensor = tf.floor(random_tensor) binary_tensor = tf.reshape(binary_tensor, [-1, 1, 1, tf.shape(input_tensor)[3]]) ret = input_tensor * binary_tensor return ret output = tf.cond(is_training, f2, f1) return output @staticmethod def lrelu(inputdata, name, alpha=0.2): with tf.variable_scope(name): return tf.nn.relu(inputdata) - alpha * tf.nn.relu(-inputdata)
--vgg_scnn_encoder.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- from collections import OrderedDict import tensorflow as tf import glog as log import math import sys sys.path.append('encoder_decoder_model') import cnn_basenet class VGG16Encoder(cnn_basenet.CNNBaseModel): def __init__(self, phase): super(VGG16Encoder, self).__init__() self._train_phase = tf.constant('train', dtype=tf.string) self._test_phase = tf.constant('test', dtype=tf.string) self._phase = phase self._is_training = self._init_phase() def _init_phase(self): return tf.equal(self._phase, self._train_phase) def _conv_stage(self, input_tensor, k_size, out_dims, name, stride=1, pad='SAME'): with tf.variable_scope(name): conv = self.conv2d(inputdata=input_tensor, out_channel=out_dims, kernel_size=k_size, stride=stride, use_bias=False, padding=pad, name='conv') bn = self.layerbn(inputdata=conv, is_training=self._is_training, name='bn') relu = self.relu(inputdata=bn, name='relu') return relu def _fc_stage(self, input_tensor, out_dims, name, use_bias=False): with tf.variable_scope(name): fc = self.fullyconnect(inputdata=input_tensor, out_dim=out_dims, use_bias=use_bias, name='fc') bn = self.layerbn(inputdata=fc, is_training=self._is_training, name='bn') relu = self.relu(inputdata=bn, name='relu') return relu def scnn_u2d_d2u(self,input_tensor): output_list_old = [] output_list_new = [] shape_list = input_tensor.get_shape().as_list() log.info("scnn_u2d_d2u shape_list:{:}".format(shape_list)) h_size = input_tensor.get_shape().as_list()[1] log.info("scnn_u2d_d2u h_size:{:}".format(h_size)) channel_size = input_tensor.get_shape().as_list()[3] #up2down conv for i in range(h_size): output_list_old.append(tf.expand_dims(input_tensor[:,i,:,:],axis=1)) output_list_new.append(tf.expand_dims(input_tensor[:,0,:,:],axis=1)) w_ud = tf.get_variable('w_ud',[1,9,channel_size,channel_size],initializer=tf.random_normal_initializer(0,math.sqrt(2.0/(9*channel_size*channel_size*2)))) with tf.variable_scope("scnn_u2d"): scnn_u2d = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_old[0],w_ud,[1,1,1,1],'SAME')),output_list_old[1]) output_list_new.append(scnn_u2d) for i in range(2,h_size): with tf.variable_scope("scnn_u2d",reuse=True): scnn_u2d = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_new[i-1],w_ud,[1,1,1,1],'SAME')),output_list_old[i]) output_list_new.append(scnn_u2d) #down2up conv output_list_old = output_list_new output_list_new = [] length = h_size-1 output_list_new.append(output_list_old[length]) w_du = tf.get_variable('w_du',[1,9,channel_size,channel_size],initializer=tf.random_normal_initializer(0,math.sqrt(2.0/(9*channel_size*channel_size*2)))) with tf.variable_scope('scnn_d2u'): scnn_d2u = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_old[length],w_du,[1,1,1,1],'SAME')),output_list_old[length-1]) output_list_new.append(scnn_d2u) for i in range(2,h_size): with tf.variable_scope("scnn_d2u",reuse=True): scnn_d2u = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_new[i-1],w_du,[1,1,1,1],'SAME')),output_list_old[length-i]) output_list_new.append(scnn_d2u) output_list_new.reverse() #log.info("scnn_u2d_d2u output_list_new:{:}".format(output_list_new)) out_tensor = tf.stack(output_list_new,axis = 1) out_tensor = tf.squeeze(out_tensor,axis=2) return out_tensor def scnn_l2r_r2l(self,input_tensor): output_list_old = [] output_list_new = [] shape_list = input_tensor.get_shape().as_list() log.info("scnn_l2r_r2l shape_list:{:}".format(shape_list)) w_size = input_tensor.get_shape().as_list()[2] log.info("scnn_l2r_r2l w_size:{:}".format(w_size)) channel_size = input_tensor.get_shape().as_list()[3] #left2right conv for i in range(w_size): output_list_old.append(tf.expand_dims(input_tensor[:,:,i,:],axis=2)) output_list_new.append(tf.expand_dims(input_tensor[:,:,0,:],axis=2)) w_lr = tf.get_variable('w_lr',[9,1,channel_size,channel_size],initializer=tf.random_normal_initializer(0,math.sqrt(2.0/(9*channel_size*channel_size*5)))) with tf.variable_scope("scnn_l2r"): scnn_l2r = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_old[0],w_lr,[1,1,1,1],'SAME')),output_list_old[1]) output_list_new.append(scnn_l2r) for i in range(2,w_size): with tf.variable_scope("scnn_l2r",reuse=True): scnn_l2r = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_new[i-1],w_lr,[1,1,1,1],'SAME')),output_list_old[i]) output_list_new.append(scnn_l2r) #log.info("output_list_new:{:}".format(output_list_new)) #right2left conv output_list_old = output_list_new output_list_new = [] length = w_size-1 output_list_new.append(output_list_old[length]) w_rl = tf.get_variable('w_rl',[9,1,channel_size,channel_size],initializer=tf.random_normal_initializer(0,math.sqrt(2.0/(9*channel_size*channel_size*5)))) with tf.variable_scope('scnn_r2l'): scnn_r2l = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_old[length],w_rl,[1,1,1,1],'SAME')),output_list_old[length-1]) output_list_new.append(scnn_r2l) for i in range(2,w_size): with tf.variable_scope("scnn_r2l",reuse=True): scnn_r2l = tf.add(tf.nn.relu(tf.nn.conv2d(output_list_new[i-1],w_rl,[1,1,1,1],'SAME')),output_list_old[length-i]) output_list_new.append(scnn_r2l) output_list_new.reverse() out_tensor = tf.stack(output_list_new,axis = 2) out_tensor = tf.squeeze(out_tensor,axis=3) return out_tensor def encode(self, input_tensor, name): ret = OrderedDict() with tf.variable_scope(name): # conv stage 1_1 conv_1_1 = self._conv_stage(input_tensor=input_tensor, k_size=3, out_dims=64, name='conv1_1') log.info("encode conv_1_1:{:}".format(conv_1_1.get_shape().as_list())) # conv stage 1_2 conv_1_2 = self._conv_stage(input_tensor=conv_1_1, k_size=3, out_dims=64, name='conv1_2') log.info("encode conv_1_2:{:}".format(conv_1_2.get_shape().as_list())) # pool stage 1 pool1 = self.maxpooling(inputdata=conv_1_2, kernel_size=2, stride=2, name='pool1') log.info("encode pool1:{:}".format(pool1.get_shape().as_list())) # conv stage 2_1 conv_2_1 = self._conv_stage(input_tensor=pool1, k_size=3, out_dims=128, name='conv2_1') log.info("encode conv_2_1:{:}".format(conv_2_1.get_shape().as_list())) # conv stage 2_2 conv_2_2 = self._conv_stage(input_tensor=conv_2_1, k_size=3, out_dims=128, name='conv2_2') log.info("encode conv_2_2:{:}".format(conv_2_2.get_shape().as_list())) # pool stage 2 pool2 = self.maxpooling(inputdata=conv_2_2, kernel_size=2, stride=2, name='pool2') log.info("encode pool2:{:}".format(pool2.get_shape().as_list())) # conv stage 3_1 conv_3_1 = self._conv_stage(input_tensor=pool2, k_size=3, out_dims=256, name='conv3_1') log.info("encode conv_3_1:{:}".format(conv_3_1.get_shape().as_list())) # conv_stage 3_2 conv_3_2 = self._conv_stage(input_tensor=conv_3_1, k_size=3, out_dims=256, name='conv3_2') log.info("encode conv_3_2:{:}".format(conv_3_2.get_shape().as_list())) # conv stage 3_3 conv_3_3 = self._conv_stage(input_tensor=conv_3_2, k_size=3, out_dims=256, name='conv3_3') log.info("encode conv_3_3:{:}".format(conv_3_3.get_shape().as_list())) ret['conv_3_3'] = dict() ret['conv_3_3']['data'] = conv_3_3 ret['conv_3_3']['shape'] = conv_3_3.get_shape().as_list() # pool stage 3 pool3 = self.maxpooling(inputdata=conv_3_3, kernel_size=2, stride=2, name='pool3') log.info("encode pool3:{:}".format(pool3.get_shape().as_list())) ret['pool3'] = dict() ret['pool3']['data'] = pool3 ret['pool3']['shape'] = pool3.get_shape().as_list() # conv stage 4_1 conv_4_1 = self._conv_stage(input_tensor=pool3, k_size=3, out_dims=512, name='conv4_1') log.info("encode conv_4_1:{:}".format(conv_4_1.get_shape().as_list())) # conv stage 4_2 conv_4_2 = self._conv_stage(input_tensor=conv_4_1, k_size=3, out_dims=512, name='conv4_2') log.info("encode conv_4_2:{:}".format(conv_4_2.get_shape().as_list())) # conv stage 4_3 conv_4_3 = self._conv_stage(input_tensor=conv_4_2, k_size=3, out_dims=512, name='conv4_3') log.info("encode conv_4_3:{:}".format(conv_4_3.get_shape().as_list())) # pool stage 4 pool4 = self.maxpooling(inputdata=conv_4_3, kernel_size=2, stride=2, name='pool4') log.info("encode pool4:{:}".format(pool4.get_shape().as_list())) ret['pool4'] = dict() ret['pool4']['data'] = pool4 ret['pool4']['shape'] = pool4.get_shape().as_list() # conv stage 5_1 conv_5_1 = self._conv_stage(input_tensor=pool4, k_size=3, out_dims=512, name='conv5_1') log.info("encode conv_5_1:{:}".format(conv_5_1.get_shape().as_list())) # conv stage 5_2 conv_5_2 = self._conv_stage(input_tensor=conv_5_1, k_size=3, out_dims=512, name='conv5_2') log.info("encode conv_5_2:{:}".format(conv_5_2.get_shape().as_list())) # conv stage 5_3 conv_5_3 = self._conv_stage(input_tensor=conv_5_2, k_size=3, out_dims=512, name='conv5_3') log.info("encode conv_5_3:{:}".format(conv_5_3.get_shape().as_list())) # conv stage 6_1 conv_6_1 = self._conv_stage(input_tensor=conv_5_3, k_size=3, out_dims=128, name='conv6_1') log.info("encode conv_6_1:{:}".format(conv_6_1.get_shape().as_list())) scnn_ud = self.scnn_u2d_d2u(conv_6_1) log.info("encode scnn_ud:{:}".format(scnn_ud.get_shape().as_list())) scnn_lr = self.scnn_l2r_r2l(scnn_ud) log.info("encode scnn_lr:{:}".format(scnn_lr.get_shape().as_list())) # pool stage 5 pool5 = self.maxpooling(inputdata=scnn_lr, kernel_size=2, stride=2, name='pool5') log.info("encode pool5:{:}".format(pool5.get_shape().as_list())) ret['pool5'] = dict() ret['pool5']['data'] = pool5 ret['pool5']['shape'] = pool5.get_shape().as_list() # fc stage 1 # fc6 = self._fc_stage(input_tensor=pool5, out_dims=4096, name='fc6', # use_bias=False, flags=flags) # fc stage 2 # fc7 = self._fc_stage(input_tensor=fc6, out_dims=4096, name='fc7', # use_bias=False, flags=flags) return ret if __name__ == '__main__': a = tf.placeholder(dtype=tf.float32, shape=[1, 2048, 2048, 3], name='input') encoder = VGG16Encoder(phase=tf.constant('train', dtype=tf.string)) ret = encoder.encode(a, name='encode') for layer_name, layer_info in ret.items(): print('layer name: {:s} shape: {}'.format(layer_name, layer_info['shape']))
--dense_encoder.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import tensorflow as tf from collections import OrderedDict #from encoder_decoder_model import cnn_basenet import cnn_basenet class DenseEncoder(cnn_basenet.CNNBaseModel): """ 基于DenseNet的编码器 """ def __init__(self, l, n, growthrate, phase, with_bc=False, bc_theta=0.5): super(DenseEncoder, self).__init__() self._L = l self._block_depth = int((l - n - 1) / n) self._N = n self._growthrate = growthrate self._with_bc = with_bc self._phase = phase self._train_phase = tf.constant('train', dtype=tf.string) self._test_phase = tf.constant('test', dtype=tf.string) self._is_training = self._init_phase() self._bc_theta = bc_theta return def _init_phase(self): return tf.equal(self._phase, self._train_phase) def __str__(self): encoder_info = 'A densenet with net depth: {:d} block nums: ' \ '{:d} growth rate: {:d} block depth: {:d}'.\ format(self._L, self._N, self._growthrate, self._block_depth) return encoder_info def _composite_conv(self, inputdata, out_channel, name): with tf.variable_scope(name): bn_1 = self.layerbn(inputdata=inputdata, is_training=self._is_training, name='bn_1') relu_1 = self.relu(bn_1, name='relu_1') if self._with_bc: conv_1 = self.conv2d(inputdata=relu_1, out_channel=out_channel, kernel_size=1, padding='SAME', stride=1, use_bias=False, name='conv_1') bn_2 = self.layerbn(inputdata=conv_1, is_training=self._is_training, name='bn_2') relu_2 = self.relu(inputdata=bn_2, name='relu_2') conv_2 = self.conv2d(inputdata=relu_2, out_channel=out_channel, kernel_size=3, stride=1, padding='SAME', use_bias=False, name='conv_2') return conv_2 else: conv_2 = self.conv2d(inputdata=relu_1, out_channel=out_channel, kernel_size=3, stride=1, padding='SAME', use_bias=False, name='conv_2') return conv_2 def _denseconnect_layers(self, inputdata, name): with tf.variable_scope(name): conv_out = self._composite_conv(inputdata=inputdata, name='composite_conv', out_channel=self._growthrate) concate_cout = tf.concat(values=[conv_out, inputdata], axis=3, name='concatenate') return concate_cout def _transition_layers(self, inputdata, name): """ Mainly implement the Pooling layer mentioned in DenseNet paper :param inputdata: :param name: :return: """ input_channels = inputdata.get_shape().as_list()[3] with tf.variable_scope(name): # First batch norm bn = self.layerbn(inputdata=inputdata, is_training=self._is_training, name='bn') # Second 1*1 conv if self._with_bc: out_channels = int(input_channels * self._bc_theta) conv = self.conv2d(inputdata=bn, out_channel=out_channels, kernel_size=1, stride=1, use_bias=False, name='conv') # Third average pooling avgpool_out = self.avgpooling(inputdata=conv, kernel_size=2, stride=2, name='avgpool') return avgpool_out else: conv = self.conv2d(inputdata=bn, out_channel=input_channels, kernel_size=1, stride=1, use_bias=False, name='conv') # Third average pooling avgpool_out = self.avgpooling(inputdata=conv, kernel_size=2, stride=2, name='avgpool') return avgpool_out def _dense_block(self, inputdata, name): """ Mainly implement the dense block mentioned in DenseNet figure 1 :param inputdata: :param name: :return: """ block_input = inputdata with tf.variable_scope(name): for i in range(self._block_depth): block_layer_name = '{:s}_layer_{:d}'.format(name, i + 1) block_input = self._denseconnect_layers(inputdata=block_input, name=block_layer_name) return block_input def encode(self, input_tensor, name): """ DenseNet编码 :param input_tensor: :param name: :return: """ encode_ret = OrderedDict() # First apply a 3*3 16 out channels conv layer # mentioned in DenseNet paper Implementation Details part with tf.variable_scope(name): conv1 = self.conv2d(inputdata=input_tensor, out_channel=16, kernel_size=3, use_bias=False, name='conv1') dense_block_input = conv1 # Second apply dense block stage for dense_block_nums in range(self._N): dense_block_name = 'Dense_Block_{:d}'.format(dense_block_nums + 1) # dense connectivity dense_block_out = self._dense_block(inputdata=dense_block_input, name=dense_block_name) # apply the trainsition part dense_block_out = self._transition_layers(inputdata=dense_block_out, name=dense_block_name) dense_block_input = dense_block_out encode_ret[dense_block_name] = dict() encode_ret[dense_block_name]['data'] = dense_block_out encode_ret[dense_block_name]['shape'] = dense_block_out.get_shape().as_list() return encode_ret if __name__ == '__main__': input_tensor = tf.placeholder(dtype=tf.float32, shape=[None, 384, 1248, 3], name='input_tensor') encoder = DenseEncoder(l=100, growthrate=16, with_bc=True, phase=tf.constant('train'), n=5) ret = encoder.encode(input_tensor=input_tensor, name='Dense_Encode') for layer_name, layer_info in ret.items(): print('layer_name: {:s} shape: {}'.format(layer_name, layer_info['shape']))
--fcn_decoder.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import tensorflow as tf #from encoder_decoder_model import cnn_basenet #from encoder_decoder_model import vgg_encoder #from encoder_decoder_model import dense_encoder import cnn_basenet import vgg_encoder import dense_encoder class FCNDecoder(cnn_basenet.CNNBaseModel): def __init__(self, phase): """ """ super(FCNDecoder, self).__init__() self._train_phase = tf.constant('train', dtype=tf.string) self._phase = phase self._is_training = self._init_phase() def _init_phase(self): """ :return: """ return tf.equal(self._phase, self._train_phase) def decode(self, input_tensor_dict, decode_layer_list, name): """ 解码特征信息反卷积还原 :param input_tensor_dict: :param decode_layer_list: 需要解码的层名称需要由深到浅顺序写 eg. ['pool5', 'pool4', 'pool3'] :param name: :return: """ ret = dict() with tf.variable_scope(name): # score stage 1 input_tensor = input_tensor_dict[decode_layer_list[0]]['data'] score = self.conv2d(inputdata=input_tensor, out_channel=64, kernel_size=1, use_bias=False, name='score_origin') ret['score'] = dict() ret['score']['data'] = score ret['score']['shape'] = score.get_shape().as_list() decode_layer_list = decode_layer_list[1:] print("len(decode_layer_list):",len(decode_layer_list)) for i in range(len(decode_layer_list)): deconv = self.deconv2d(inputdata=score, out_channel=64, kernel_size=4, stride=2, use_bias=False, name='deconv_{:d}'.format(i + 1)) input_tensor = input_tensor_dict[decode_layer_list[i]]['data'] score = self.conv2d(inputdata=input_tensor, out_channel=64, kernel_size=1, use_bias=False, name='score_{:d}'.format(i + 1)) fused = tf.add(deconv, score, name='fuse_{:d}'.format(i + 1)) score = fused ret['fuse_{:d}'.format(i + 1)] = dict() ret['fuse_{:d}'.format(i + 1)]['data'] = fused ret['fuse_{:d}'.format(i + 1)]['shape'] = fused.get_shape().as_list() deconv_final = self.deconv2d(inputdata=score, out_channel=64, kernel_size=16, stride=8, use_bias=False, name='deconv_final') score_final = self.conv2d(inputdata=deconv_final, out_channel=2, kernel_size=1, use_bias=False, name='score_final') ret['logits'] = score_final ret['deconv'] = deconv_final ret['logits'] = dict() ret['logits']['data'] = score_final ret['logits']['shape'] = score_final.get_shape().as_list() ret['deconv'] = dict() ret['deconv']['data'] = deconv_final ret['deconv']['shape'] = deconv_final.get_shape().as_list() return ret if __name__ == '__main__': vgg_encoder = vgg_encoder.VGG16Encoder(phase=tf.constant('train', tf.string)) dense_encoder = dense_encoder.DenseEncoder(l=40, growthrate=12, with_bc=True, phase='train', n=5) decoder = FCNDecoder(phase='train') in_tensor = tf.placeholder(dtype=tf.float32, shape=[None, 256, 512, 3], name='input') vgg_encode_ret = vgg_encoder.encode(in_tensor, name='vgg_encoder') dense_encode_ret = dense_encoder.encode(in_tensor, name='dense_encoder') decode_ret = decoder.decode(vgg_encode_ret, name='decoder', decode_layer_list=['pool5', 'pool4', 'pool3']) for layer_name, layer_info in decode_ret.items(): print('layer name: {:s} shape: {}'.format(layer_name, layer_info['shape']))
./merge_model
--merge_model.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import tensorflow as tf from encoder_decoder_model import vgg_encoder from encoder_decoder_model import fcn_decoder from encoder_decoder_model import dense_encoder from encoder_decoder_model import cnn_basenet from lanenet_model import lanenet_discriminative_loss from encoder_decoder_model import vgg_scnn_encoder import glog class LaneNet(cnn_basenet.CNNBaseModel): """ 实现语义分割模型 """ def __init__(self, phase, net_flag='vgg'): """ """ super(LaneNet, self).__init__() self._net_flag = net_flag self._phase = phase if self._net_flag == 'vgg': self._encoder = vgg_encoder.VGG16Encoder(phase=phase) elif self._net_flag == 'vgg_scnn': self._encoder = vgg_scnn_encoder.VGG16Encoder(phase=phase) elif self._net_flag == 'dense': self._encoder = dense_encoder.DenseEncoder(l=20, growthrate=8, with_bc=True, phase=phase, n=5) self._decoder = fcn_decoder.FCNDecoder(phase=phase) return def __str__(self): """ :return: """ info = 'Semantic Segmentation use {:s} as basenet to encode'.format(self._net_flag) return info def _build_model(self, input_tensor, name): """ 前向传播过程 :param input_tensor: :param name: :return: """ with tf.variable_scope(name): # first encode encode_ret = self._encoder.encode(input_tensor=input_tensor, name='encode') # second decode if self._net_flag.lower() == 'vgg': decode_ret = self._decoder.decode(input_tensor_dict=encode_ret, name='decode', decode_layer_list=['pool5', 'pool4', 'pool3']) return decode_ret if self._net_flag.lower() == 'vgg_scnn': decode_ret = self._decoder.decode(input_tensor_dict=encode_ret, name='decode', decode_layer_list=['pool5', 'pool4', 'pool3']) return decode_ret elif self._net_flag.lower() == 'dense': decode_ret = self._decoder.decode(input_tensor_dict=encode_ret, name='decode', decode_layer_list=['Dense_Block_5', 'Dense_Block_4', 'Dense_Block_3']) return decode_ret def compute_loss(self, input_tensor, binary_label, instance_label, name): """ 计算LaneNet模型损失函数 :param input_tensor: :param binary_label: :param instance_label: :param name: :return: """ with tf.variable_scope(name): # 前向传播获取logits inference_ret = self._build_model(input_tensor=input_tensor, name='inference') glog.info('compute_loss inference_ret:{:}'.format(inference_ret)) # 计算二值分割损失函数 decode_logits = inference_ret['logits'] binary_label_plain = tf.reshape( binary_label, shape=[binary_label.get_shape().as_list()[0] * binary_label.get_shape().as_list()[1] * binary_label.get_shape().as_list()[2]]) glog.info('compute_loss binary_label_plain:{:}'.format(binary_label_plain)) # 加入class weights unique_labels, unique_id, counts = tf.unique_with_counts(binary_label_plain) counts = tf.cast(counts, tf.float32) glog.info('compute_loss counts:{:}'.format(counts)) inverse_weights = tf.divide(1.0, tf.log(tf.add(tf.divide(tf.constant(1.0), counts), tf.constant(1.02)))) glog.info('compute_loss inverse_weights:{:}'.format(inverse_weights)) inverse_weights = tf.gather(inverse_weights, binary_label) glog.info('compute_loss gather inverse_weights:{:}'.format(inverse_weights)) binary_segmenatation_loss = tf.losses.sparse_softmax_cross_entropy( labels=binary_label, logits=decode_logits, weights=inverse_weights) glog.info('compute_loss binary_segmenatation_loss:{:}'.format(binary_segmenatation_loss)) binary_segmenatation_loss = tf.reduce_mean(binary_segmenatation_loss) glog.info('compute_loss reduce_mean binary_segmenatation_loss:{:}'.format(binary_segmenatation_loss)) # 计算discriminative loss损失函数 decode_deconv = inference_ret['deconv'] # 像素嵌入 pix_embedding = self.conv2d(inputdata=decode_deconv, out_channel=4, kernel_size=1, use_bias=False, name='pix_embedding_conv') pix_embedding = self.relu(inputdata=pix_embedding, name='pix_embedding_relu') # 计算discriminative loss image_shape = (pix_embedding.get_shape().as_list()[1], pix_embedding.get_shape().as_list()[2]) glog.info('compute_loss image_shape:{:}'.format(image_shape)) disc_loss, l_var, l_dist, l_reg = \ lanenet_discriminative_loss.discriminative_loss( pix_embedding, instance_label, 4, image_shape, 0.5, 3.0, 1.0, 1.0, 0.001) glog.info('compute_loss disc_loss:{:}'.format(disc_loss)) # 合并损失 l2_reg_loss = tf.constant(0.0, tf.float32) for vv in tf.trainable_variables(): if 'bn' in vv.name: continue else: l2_reg_loss = tf.add(l2_reg_loss, tf.nn.l2_loss(vv)) l2_reg_loss *= 0.001 total_loss = 0.5 * binary_segmenatation_loss + 0.5 * disc_loss + l2_reg_loss ret = { 'total_loss': total_loss, 'binary_seg_logits': decode_logits, 'instance_seg_logits': pix_embedding, 'binary_seg_loss': binary_segmenatation_loss, 'discriminative_loss': disc_loss } return ret def inference(self, input_tensor, name): """ :param input_tensor: :param name: :return: """ with tf.variable_scope(name): # 前向传播获取logits inference_ret = self._build_model(input_tensor=input_tensor, name='inference') # 计算二值分割损失函数 decode_logits = inference_ret['logits'] binary_seg_ret = tf.nn.softmax(logits=decode_logits) binary_seg_ret = tf.argmax(binary_seg_ret, axis=-1) # 计算像素嵌入 decode_deconv = inference_ret['deconv'] # 像素嵌入 pix_embedding = self.conv2d(inputdata=decode_deconv, out_channel=4, kernel_size=1, use_bias=False, name='pix_embedding_conv') pix_embedding = self.relu(inputdata=pix_embedding, name='pix_embedding_relu') return binary_seg_ret, pix_embedding if __name__ == '__main__': model = LaneNet(tf.constant('train', dtype=tf.string)) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input') binary_label = tf.placeholder(dtype=tf.int64, shape=[1, 256, 512, 1], name='label') instance_label = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 1], name='label') ret = model.compute_loss(input_tensor=input_tensor, binary_label=binary_label, instance_label=instance_label, name='loss') for vv in tf.trainable_variables(): if 'bn' in vv.name: continue print(vv.name)
--dirscriminative_loss.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import tensorflow as tf import glog def discriminative_loss_single( prediction, correct_label, feature_dim, label_shape, delta_v, delta_d, param_var, param_dist, param_reg): """ 论文equ(1)提到的实例分割损失函数 :param prediction: inference of network :param correct_label: instance label :param feature_dim: feature dimension of prediction :param label_shape: shape of label :param delta_v: cut off variance distance :param delta_d: cut off cluster distance :param param_var: weight for intra cluster variance :param param_dist: weight for inter cluster distances :param param_reg: weight regularization """ # 像素对齐为一行 correct_label = tf.reshape( correct_label, [ label_shape[1] * label_shape[0]]) reshaped_pred = tf.reshape( prediction, [ label_shape[1] * label_shape[0], feature_dim]) # 统计实例个数 unique_labels, unique_id, counts = tf.unique_with_counts(correct_label) counts = tf.cast(counts, tf.float32) num_instances = tf.size(unique_labels) glog.info('discriminative_loss_single counts:{:} num_instances:{:}'.format(counts,num_instances)) # 计算pixel embedding均值向量 segmented_sum = tf.unsorted_segment_sum( reshaped_pred, unique_id, num_instances) mu = tf.div(segmented_sum, tf.reshape(counts, (-1, 1))) mu_expand = tf.gather(mu, unique_id) # 计算公式的loss(var) distance = tf.norm(tf.subtract(mu_expand, reshaped_pred), axis=1) distance = tf.subtract(distance, delta_v) distance = tf.clip_by_value(distance, 0., distance) distance = tf.square(distance) l_var = tf.unsorted_segment_sum(distance, unique_id, num_instances) l_var = tf.div(l_var, counts) l_var = tf.reduce_sum(l_var) l_var = tf.divide(l_var, tf.cast(num_instances, tf.float32)) # 计算公式的loss(dist) mu_interleaved_rep = tf.tile(mu, [num_instances, 1]) mu_band_rep = tf.tile(mu, [1, num_instances]) mu_band_rep = tf.reshape( mu_band_rep, (num_instances * num_instances, feature_dim)) mu_diff = tf.subtract(mu_band_rep, mu_interleaved_rep) # 去除掩模上的零点 intermediate_tensor = tf.reduce_sum(tf.abs(mu_diff), axis=1) zero_vector = tf.zeros(1, dtype=tf.float32) bool_mask = tf.not_equal(intermediate_tensor, zero_vector) mu_diff_bool = tf.boolean_mask(mu_diff, bool_mask) mu_norm = tf.norm(mu_diff_bool, axis=1) mu_norm = tf.subtract(2. * delta_d, mu_norm) mu_norm = tf.clip_by_value(mu_norm, 0., mu_norm) mu_norm = tf.square(mu_norm) l_dist = tf.reduce_mean(mu_norm) # 计算原始Discriminative Loss论文中提到的正则项损失 l_reg = tf.reduce_mean(tf.norm(mu, axis=1)) # 合并损失按照原始Discriminative Loss论文中提到的参数合并 param_scale = 1. l_var = param_var * l_var l_dist = param_dist * l_dist l_reg = param_reg * l_reg loss = param_scale * (l_var + l_dist + l_reg) return loss, l_var, l_dist, l_reg def discriminative_loss(prediction, correct_label, feature_dim, image_shape, delta_v, delta_d, param_var, param_dist, param_reg): """ 按照论文的思想迭代计算loss损失 :return: discriminative loss and its three components """ def cond(label, batch, out_loss, out_var, out_dist, out_reg, i): return tf.less(i, tf.shape(batch)[0]) def body(label, batch, out_loss, out_var, out_dist, out_reg, i): disc_loss, l_var, l_dist, l_reg = discriminative_loss_single( prediction[i], correct_label[i], feature_dim, image_shape, delta_v, delta_d, param_var, param_dist, param_reg) out_loss = out_loss.write(i, disc_loss) out_var = out_var.write(i, l_var) out_dist = out_dist.write(i, l_dist) out_reg = out_reg.write(i, l_reg) return label, batch, out_loss, out_var, out_dist, out_reg, i + 1 # TensorArray is a data structure that support dynamic writing output_ta_loss = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True) output_ta_var = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True) output_ta_dist = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True) output_ta_reg = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True) _, _, out_loss_op, out_var_op, out_dist_op, out_reg_op, _ = tf.while_loop( cond, body, [ correct_label, prediction, output_ta_loss, output_ta_var, output_ta_dist, output_ta_reg, 0]) out_loss_op = out_loss_op.stack() out_var_op = out_var_op.stack() out_dist_op = out_dist_op.stack() out_reg_op = out_reg_op.stack() disc_loss = tf.reduce_mean(out_loss_op) l_var = tf.reduce_mean(out_var_op) l_dist = tf.reduce_mean(out_dist_op) l_reg = tf.reduce_mean(out_reg_op) return disc_loss, l_var, l_dist, l_reg
--postpostprecess.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ LaneNet模型后处理 """ import numpy as np import matplotlib.pyplot as plt import cv2 import glog try: from cv2 import cv2 except ImportError: pass class LaneNetPoseProcessor(object): """ """ def __init__(self): """ """ pass @staticmethod def _morphological_process(image, kernel_size=5): """ :param image: :param kernel_size: :return: """ if image.dtype is not np.uint8: image = np.array(image, np.uint8) glog.info("_morphological_process image shape len:{:d}".format(len(image.shape))) if len(image.shape) == 3: image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) glog.info("_morphological_process image shape len:{:d}".format(len(image.shape))) kernel = cv2.getStructuringElement(shape=cv2.MORPH_ELLIPSE, ksize=(kernel_size, kernel_size)) # close operation fille hole closing = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel, iterations=1) return closing @staticmethod def _connect_components_analysis(image): """ :param image: :return: """ glog.info("_connect_components_analysis image shape len:{:d}".format(len(image.shape))) if len(image.shape) == 3: gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: gray_image = image return cv2.connectedComponentsWithStats(gray_image, connectivity=8, ltype=cv2.CV_32S) def postprocess(self, image, minarea_threshold=15): """ :param image: :param minarea_threshold: 连通域分析阈值 :return: """ # 首先进行图像形态学运算 morphological_ret = self._morphological_process(image, kernel_size=5) glog.info("postprocess image shape len:{:d}".format(len(image.shape))) # 进行连通域分析 connect_components_analysis_ret = self._connect_components_analysis(image=morphological_ret) glog.info("postprocess connect_components_analysis_ret:{:}".format(connect_components_analysis_ret)) # 排序连通域并删除过小的连通域 labels = connect_components_analysis_ret[1] stats = connect_components_analysis_ret[2] glog.info("postprocess labels:{:}".format(labels)) glog.info("postprocess stats:{:}".format(stats)) for index, stat in enumerate(stats): if stat[4] <= minarea_threshold: idx = np.where(labels == index) morphological_ret[idx] = 0 return morphological_ret if __name__ == '__main__': processor = LaneNetPoseProcessor() image = cv2.imread('D:/Code/github/tf_lanenet/data/training_data_example/gt_image_binary/0000.png', cv2.IMREAD_UNCHANGED) #IMREAD_GRAYSCALE postprocess_ret = processor.postprocess(image) plt.figure('src') plt.imshow(image) plt.figure('post') plt.imshow(postprocess_ret) plt.show()
--cluster.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 实现LaneNet中实例分割的聚类部分 """ import numpy as np import glog as log import matplotlib.pyplot as plt from sklearn.cluster import MeanShift from sklearn.cluster import DBSCAN import time import warnings import cv2 import glog try: from cv2 import cv2 except ImportError: pass class LaneNetCluster(object): """ 实例分割聚类器 """ def __init__(self): """ """ self._color_map = [np.array([255, 0, 0]), np.array([0, 255, 0]), np.array([0, 0, 255]), np.array([125, 125, 0]), np.array([0, 125, 125]), np.array([125, 0, 125]), np.array([50, 100, 50]), np.array([100, 50, 100])] pass @staticmethod def _cluster(prediction, bandwidth): """ 实现论文SectionⅡ的cluster部分 :param prediction: :param bandwidth: :return: """ ms = MeanShift(bandwidth, bin_seeding=True) # log.info('开始Mean shift聚类 ...') tic = time.time() try: ms.fit(prediction) except ValueError as err: log.error(err) return 0, [], [] # log.info('Mean Shift耗时: {:.5f}s'.format(time.time() - tic)) labels = ms.labels_ cluster_centers = ms.cluster_centers_ num_clusters = cluster_centers.shape[0] # log.info('聚类簇个数为: {:d}'.format(num_clusters)) return num_clusters, labels, cluster_centers @staticmethod def _cluster_v2(prediction): """ dbscan cluster :param prediction: :return: """ db = DBSCAN(eps=0.7, min_samples=200).fit(prediction) db_labels = db.labels_ unique_labels = np.unique(db_labels) unique_labels = [tmp for tmp in unique_labels if tmp != -1] log.info('聚类簇个数为: {:d}'.format(len(unique_labels))) num_clusters = len(unique_labels) cluster_centers = db.components_ return num_clusters, db_labels, cluster_centers @staticmethod def _get_lane_area(binary_seg_ret, instance_seg_ret): """ 通过二值分割掩码图在实例分割图上获取所有车道线的特征向量 :param binary_seg_ret: :param instance_seg_ret: :return: """ idx = np.where(binary_seg_ret == 1) print("_get_lane_area idx:",idx) print("_get_lane_area idx len:",len(idx)) print("_get_lane_area idx len[0]:",len(idx[0])) print("_get_lane_area idx len[1]:",len(idx[1])) lane_embedding_feats = [] lane_coordinate = [] for i in range(len(idx[0])): lane_embedding_feats.append(instance_seg_ret[idx[0][i], idx[1][i]]) #print("_get_lane_area instance_seg_ret[idx[0][i], idx[1][i]]:",instance_seg_ret[idx[0][i], idx[1][i]]) lane_coordinate.append([idx[0][i], idx[1][i]]) #print("_get_lane_area idx[0][i]:",idx[0][i]," , idx[1][i]:", idx[1][i]) return np.array(lane_embedding_feats, np.float32), np.array(lane_coordinate, np.int64) @staticmethod def _thresh_coord(coord): """ 过滤实例车道线位置坐标点,假设车道线是连续的, 因此车道线点的坐标变换应该是平滑变化的不应该出现跳变 :param coord: [(x, y)] :return: """ pts_x = coord[:, 0] mean_x = np.mean(pts_x) idx = np.where(np.abs(pts_x - mean_x) < mean_x) return coord[idx[0]] @staticmethod def _lane_fit(lane_pts): """ 车道线多项式拟合 :param lane_pts: :return: """ if not isinstance(lane_pts, np.ndarray): lane_pts = np.array(lane_pts, np.float32) x = lane_pts[:, 0] y = lane_pts[:, 1] x_fit = [] y_fit = [] with warnings.catch_warnings(): warnings.filterwarnings('error') try: f1 = np.polyfit(x, y, 3) p1 = np.poly1d(f1) x_min = int(np.min(x)) x_max = int(np.max(x)) x_fit = [] for i in range(x_min, x_max + 1): x_fit.append(i) y_fit = p1(x_fit) except Warning as e: x_fit = x y_fit = y finally: return zip(x_fit, y_fit) def get_lane_mask(self, binary_seg_ret, instance_seg_ret): """ :param binary_seg_ret: :param instance_seg_ret: :return: """ lane_embedding_feats, lane_coordinate = self._get_lane_area(binary_seg_ret, instance_seg_ret) num_clusters, labels, cluster_centers = self._cluster(lane_embedding_feats, bandwidth=1.5) # 聚类簇超过八个则选择其中类内样本最多的八个聚类簇保留下来 if num_clusters > 8: cluster_sample_nums = [] for i in range(num_clusters): cluster_sample_nums.append(len(np.where(labels == i)[0])) sort_idx = np.argsort(-np.array(cluster_sample_nums, np.int64)) cluster_index = np.array(range(num_clusters))[sort_idx[0:4]] else: cluster_index = range(num_clusters) mask_image = np.zeros(shape=[binary_seg_ret.shape[0], binary_seg_ret.shape[1], 3], dtype=np.uint8) for index, i in enumerate(cluster_index): idx = np.where(labels == i) coord = lane_coordinate[idx] # coord = self._thresh_coord(coord) coord = np.flip(coord, axis=1) # coord = (coord[:, 0], coord[:, 1]) color = (int(self._color_map[index][0]), int(self._color_map[index][1]), int(self._color_map[index][2])) coord = np.array([coord]) cv2.polylines(img=mask_image, pts=coord, isClosed=False, color=color, thickness=2) # mask_image[coord] = color return mask_image if __name__ == '__main__': binary_seg_image = cv2.imread('D:/Code/github/tf_lanenet/data/training_data_example/gt_image_binary/0000.png', cv2.IMREAD_GRAYSCALE) print("binary_seg_image shape:",binary_seg_image.shape) binary_seg_image[np.where(binary_seg_image == 255)] = 1 print("binary_seg_image np.where(binary_seg_image == 255):",np.where(binary_seg_image == 255)) instance_seg_image = cv2.imread('D:/Code/github/tf_lanenet/data/training_data_example/gt_image_instance/0000.png', cv2.IMREAD_UNCHANGED) glog.info("__name__ instance_seg_image shape len:{:d}".format(len(instance_seg_image.shape))) instance_seg_image = cv2.cvtColor(instance_seg_image, cv2.COLOR_GRAY2BGR) glog.info("__name__ instance_seg_image shape len:{:d}".format(len(instance_seg_image.shape))) #print("instance_seg_image shape:",instance_seg_image.shape) ele_mex = np.max(instance_seg_image, axis=(0,1)) print("ele_mex:",ele_mex) for i in range(3): if ele_mex[i] == 0: scale = 1 else: scale = 255 / ele_mex[i] instance_seg_image[:, :, i] *= int(scale) embedding_image = np.array(instance_seg_image, np.uint8) cluster = LaneNetCluster() mask_image = cluster.get_lane_mask(binary_seg_ret=binary_seg_image,instance_seg_ret=instance_seg_image) det_img = embedding_image+mask_image plt.figure('det_img') plt.imshow(det_img[:, :, (2, 1, 0)]) #plt.figure('embedding') #plt.imshow(embedding_image[:, :, (2, 1, 0)]) #plt.figure('mask_image') #plt.imshow(mask_image[:, :, (2, 1, 0)]) plt.show()
--train_lane_scnn.py
#!/usr/bin/env python3 # -*- coding: utf-8 -*- import argparse import math import os import os.path as ops import time import cv2 import glog as log import numpy as np import tensorflow as tf from config import global_config from lanenet_model import lanenet_merge_model from data_provider import lanenet_data_processor CFG = global_config.cfg VGG_MEAN = [103.939, 116.779, 123.68] def init_args(): """ :return: """ parser = argparse.ArgumentParser() parser.add_argument('--dataset_dir', type=str,default='data/datasets_culane', help='The training dataset dir path') parser.add_argument('--net', type=str, default='vgg', help='Which base net work to use') parser.add_argument('--weights_path', type=str,default='model/lanenet_culane_vgg_2019-02-02-14-05-16.ckpt-200000',help='The pretrained weights path') return parser.parse_args() def minmax_scale(input_arr): """ :param input_arr: :return: """ min_val = np.min(input_arr) max_val = np.max(input_arr) output_arr = (input_arr - min_val) * 255.0 / (max_val - min_val) return output_arr def train_net(dataset_dir, weights_path=None, net_flag='vgg'): """ :param dataset_dir: :param net_flag: choose which base network to use :param weights_path: :return: """ train_dataset_file = ops.join(dataset_dir, 'train.txt') val_dataset_file = ops.join(dataset_dir, 'val.txt') print('train_dataset_file:',train_dataset_file) print('val_dataset_file:',val_dataset_file) assert ops.exists(train_dataset_file) train_dataset = lanenet_data_processor.DataSet(train_dataset_file) val_dataset = lanenet_data_processor.DataSet(val_dataset_file) with tf.device('/gpu:1'): input_tensor = tf.placeholder(dtype=tf.float32, shape=[CFG.TRAIN.BATCH_SIZE, CFG.TRAIN.IMG_HEIGHT, CFG.TRAIN.IMG_WIDTH, 3], name='input_tensor') binary_label_tensor = tf.placeholder(dtype=tf.int64, shape=[CFG.TRAIN.BATCH_SIZE, CFG.TRAIN.IMG_HEIGHT, CFG.TRAIN.IMG_WIDTH, 1], name='binary_input_label') instance_label_tensor = tf.placeholder(dtype=tf.float32, shape=[CFG.TRAIN.BATCH_SIZE, CFG.TRAIN.IMG_HEIGHT, CFG.TRAIN.IMG_WIDTH], name='instance_input_label') phase = tf.placeholder(dtype=tf.string, shape=None, name='net_phase') net = lanenet_merge_model.LaneNet(net_flag=net_flag, phase=phase) # calculate the loss compute_ret = net.compute_loss(input_tensor=input_tensor, binary_label=binary_label_tensor, instance_label=instance_label_tensor, name='lanenet_model') total_loss = compute_ret['total_loss'] binary_seg_loss = compute_ret['binary_seg_loss'] disc_loss = compute_ret['discriminative_loss'] pix_embedding = compute_ret['instance_seg_logits'] # calculate the accuracy out_logits = compute_ret['binary_seg_logits'] out_logits = tf.nn.softmax(logits=out_logits) out_logits_out = tf.argmax(out_logits, axis=-1) #out = tf.argmax(out_logits, axis=-1) #out = tf.expand_dims(out, axis=-1) out = tf.expand_dims(out_logits_out,axis=-1) idx = tf.where(tf.equal(binary_label_tensor, 1)) pix_cls_ret = tf.gather_nd(out, idx) accuracy = tf.count_nonzero(pix_cls_ret) accuracy = tf.divide(accuracy, tf.cast(tf.shape(pix_cls_ret)[0], tf.int64)) global_step = tf.Variable(0, trainable=False) learning_rate = tf.train.exponential_decay(CFG.TRAIN.LEARNING_RATE, global_step, 100000, 0.1, staircase=True) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): optimizer = tf.train.MomentumOptimizer( learning_rate=learning_rate, momentum=0.9).minimize(loss=total_loss, var_list=tf.trainable_variables(), global_step=global_step) # Set tf saver saver = tf.train.Saver() model_save_dir = 'model/lanenet_culane' if not ops.exists(model_save_dir): os.makedirs(model_save_dir) train_start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) model_name = 'lanenet_culane_{:s}_{:s}.ckpt'.format(net_flag, str(train_start_time)) model_save_path = ops.join(model_save_dir, model_name) # Set tf summary tboard_save_path = 'tboard/lanenet_culane/{:s}'.format(net_flag) if not ops.exists(tboard_save_path): os.makedirs(tboard_save_path) train_cost_scalar = tf.summary.scalar(name='train_cost', tensor=total_loss) val_cost_scalar = tf.summary.scalar(name='val_cost', tensor=total_loss) train_accuracy_scalar = tf.summary.scalar(name='train_accuracy', tensor=accuracy) val_accuracy_scalar = tf.summary.scalar(name='val_accuracy', tensor=accuracy) train_binary_seg_loss_scalar = tf.summary.scalar(name='train_binary_seg_loss', tensor=binary_seg_loss) val_binary_seg_loss_scalar = tf.summary.scalar(name='val_binary_seg_loss', tensor=binary_seg_loss) train_instance_seg_loss_scalar = tf.summary.scalar(name='train_instance_seg_loss', tensor=disc_loss) val_instance_seg_loss_scalar = tf.summary.scalar(name='val_instance_seg_loss', tensor=disc_loss) learning_rate_scalar = tf.summary.scalar(name='learning_rate', tensor=learning_rate) train_merge_summary_op = tf.summary.merge([train_accuracy_scalar, train_cost_scalar, learning_rate_scalar, train_binary_seg_loss_scalar, train_instance_seg_loss_scalar]) val_merge_summary_op = tf.summary.merge([val_accuracy_scalar, val_cost_scalar, val_binary_seg_loss_scalar, val_instance_seg_loss_scalar]) # Set sess configuration sess_config = tf.ConfigProto(allow_soft_placement=True) sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TRAIN.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.TRAIN.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) summary_writer = tf.summary.FileWriter(tboard_save_path) summary_writer.add_graph(sess.graph) # Set the training parameters train_epochs = CFG.TRAIN.EPOCHS log.info('Global configuration is as follows:') log.info(CFG) with sess.as_default(): tf.train.write_graph(graph_or_graph_def=sess.graph, logdir='', name='{:s}/lanenet_model.pb'.format(model_save_dir)) if weights_path is None: log.info('Training from scratch') init = tf.global_variables_initializer() sess.run(init) else: log.info('Restore model from last model checkpoint {:s}'.format(weights_path)) saver.restore(sess=sess, save_path=weights_path) # 加载预训练参数 log.info('jim.chen train_net net_flag:',net_flag) if net_flag == 'vgg' and weights_path is None: pretrained_weights = np.load( './data/vgg16.npy', encoding='latin1').item() log.info('jim.chen train_net net_flag is 1vgg') for vv in tf.trainable_variables(): weights_key = vv.name.split('/')[-3] try: weights = pretrained_weights[weights_key][0] _op = tf.assign(vv, weights) sess.run(_op) except Exception as e: continue train_cost_time_mean = [] val_cost_time_mean = [] for epoch in range(train_epochs): # training part t_start = time.time() with tf.device('/cpu:0'): gt_imgs, binary_gt_labels,instance_gt_labels = train_dataset.next_batch(CFG.TRAIN.BATCH_SIZE) gt_imgs = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_LINEAR) for tmp in gt_imgs] gt_imgs = [tmp - VGG_MEAN for tmp in gt_imgs] binary_gt_labels = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_NEAREST) for tmp in binary_gt_labels] binary_gt_labels = [np.expand_dims(tmp, axis=-1) for tmp in binary_gt_labels] instance_gt_labels = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_NEAREST) for tmp in instance_gt_labels] phase_train = 'train' _, c, train_accuracy, train_summary, binary_loss, instance_loss, embedding, binary_seg_img = \ sess.run([optimizer, total_loss, accuracy, train_merge_summary_op, binary_seg_loss, disc_loss, pix_embedding, out_logits_out], feed_dict={input_tensor: gt_imgs, binary_label_tensor: binary_gt_labels, instance_label_tensor: instance_gt_labels, phase: phase_train}) if math.isnan(c) or math.isnan(instance_loss) or math.isnan(binary_loss): log.error('cost is: {:.5f}'.format(c)) log.error('binary cost is: {:.5f}'.format(binary_loss)) log.error('instance cost is: {:.5f}'.format(instance_loss)) cv2.imwrite('nan_image.png', gt_imgs[0] + VGG_MEAN) cv2.imwrite('nan_instance_label.png', instance_gt_labels[0]) cv2.imwrite('nan_binary_label.png', binary_gt_labels[0] * 255) return if epoch % 100 == 0: cv2.imwrite('image.png', gt_imgs[0] + VGG_MEAN) cv2.imwrite('binary_label.png', binary_gt_labels[0] * 255) cv2.imwrite('instance_label.png', instance_gt_labels[0]) cv2.imwrite('binary_seg_img.png', binary_seg_img[0] * 255) for i in range(4): embedding[0][:, :, i] = minmax_scale(embedding[0][:, :, i]) embedding_image = np.array(embedding[0], np.uint8) cv2.imwrite('embedding.png', embedding_image) cost_time = time.time() - t_start train_cost_time_mean.append(cost_time) summary_writer.add_summary(summary=train_summary, global_step=epoch) # validation part with tf.device('/cpu:0'): gt_imgs_val, binary_gt_labels_val, instance_gt_labels_val \ = val_dataset.next_batch(CFG.TRAIN.VAL_BATCH_SIZE) gt_imgs_val = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_LINEAR) for tmp in gt_imgs_val] gt_imgs_val = [tmp - VGG_MEAN for tmp in gt_imgs_val] binary_gt_labels_val = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp) for tmp in binary_gt_labels_val] binary_gt_labels_val = [np.expand_dims(tmp, axis=-1) for tmp in binary_gt_labels_val] instance_gt_labels_val = [cv2.resize(tmp, dsize=(CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), dst=tmp, interpolation=cv2.INTER_NEAREST) for tmp in instance_gt_labels_val] phase_val = 'test' t_start_val = time.time() c_val, val_summary, val_accuracy, val_binary_seg_loss, val_instance_seg_loss = \ sess.run([total_loss, val_merge_summary_op, accuracy, binary_seg_loss, disc_loss], feed_dict={input_tensor: gt_imgs_val, binary_label_tensor: binary_gt_labels_val, instance_label_tensor: instance_gt_labels_val, phase: phase_val}) if epoch % 100 == 0: cv2.imwrite('test_image.png', gt_imgs_val[0] + VGG_MEAN) summary_writer.add_summary(val_summary, global_step=epoch) cost_time_val = time.time() - t_start_val val_cost_time_mean.append(cost_time_val) if epoch % CFG.TRAIN.DISPLAY_STEP == 0: log.info('Epoch: {:d} total_loss= {:6f} binary_seg_loss= {:6f} instance_seg_loss= {:6f} accuracy= {:6f}' ' mean_cost_time= {:5f}s '. format(epoch + 1, c, binary_loss, instance_loss, train_accuracy, np.mean(train_cost_time_mean))) train_cost_time_mean.clear() if epoch % CFG.TRAIN.TEST_DISPLAY_STEP == 0: log.info('Epoch_Val: {:d} total_loss= {:6f} binary_seg_loss= {:6f} ' 'instance_seg_loss= {:6f} accuracy= {:6f} ' 'mean_cost_time= {:5f}s '. format(epoch + 1, c_val, val_binary_seg_loss, val_instance_seg_loss, val_accuracy, np.mean(val_cost_time_mean))) val_cost_time_mean.clear() if epoch % 2000 == 0: saver.save(sess=sess, save_path=model_save_path, global_step=epoch) sess.close() return if __name__ == '__main__': # init args args = init_args() # train lanenet train_net(args.dataset_dir, args.weights_path, net_flag=args.net)
以下是模型训练过程中生成的文件夹：
./summary
./figure
./checkpoint
在主目录下，执行python train_lanenet_scnn.py,没有问题的话，可以开始训练了．．．
下一篇：