def init_lanenet(self): ''' initlize the tensorflow model ''' self.input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') phase_tensor = tf.constant('test', tf.string) net = lanenet.LaneNet(phase=phase_tensor, net_flag='vgg') self.binary_seg_ret, self.instance_seg_ret = net.inference( input_tensor=self.input_tensor, name='lanenet_model') # self.cluster = lanenet_cluster.LaneNetCluster() self.postprocessor = lanenet_postprocess.LaneNetPostProcessor() saver = tf.train.Saver() # Set sess configuration if self.use_gpu: sess_config = tf.ConfigProto(device_count={'GPU': 1}) else: sess_config = tf.ConfigProto(device_count={'CPU': 0}) sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.TRAIN.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' self.sess = tf.Session(config=sess_config) saver.restore(sess=self.sess, save_path=self.weight_path) rospy.loginfo("Done Initializing")
def __init__(self, weights, config): """ Initializes a LanePredictor which is used to register a callback for the RGB images and predict lanes. Args: weights: The path of the weights to be used in the prediction. config: The config to be used for tensorflow. """ self.input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') self.net = lanenet.LaneNet(phase='test', net_flag='vgg') self.binary_seg_ret, self.instance_seg_ret = self.net.inference( input_tensor=self.input_tensor, name='lanenet_model') self.postprocessor = lanenet_postprocess.LaneNetPostProcessor( ipm_remap_file_path= './dependencies/lanenet-lane-detection/data/tusimple_ipm_remap.yml' ) sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = \ config.TEST.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = config.TRAIN.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' self.sess = tf.Session(config=sess_config).__enter__() saver = tf.train.Saver() saver.restore(sess=self.sess, save_path=weights)
def __init__(self, camera_stream, detected_lanes_stream, flags): camera_stream.add_callback(self.on_camera_frame, [detected_lanes_stream]) self._flags = flags self._logger = erdos.utils.setup_logging(self.config.name, self.config.log_file_name) pylot.utils.set_tf_loglevel(logging.ERROR) self._input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test') self._binary_seg_ret, self._instance_seg_ret = net.inference( input_tensor=self._input_tensor, name='LaneNet') self._gpu_options = tf.GPUOptions( allow_growth=True, visible_device_list=str(self._flags.lane_detection_gpu_index), per_process_gpu_memory_fraction=flags. lane_detection_gpu_memory_fraction, allocator_type='BFC') self._tf_session = tf.Session(config=tf.ConfigProto( gpu_options=self._gpu_options, allow_soft_placement=True)) with tf.variable_scope(name_or_scope='moving_avg'): variable_averages = tf.train.ExponentialMovingAverage(0.9995) variables_to_restore = variable_averages.variables_to_restore() self._postprocessor = lanenet_postprocess.LaneNetPostProcessor() saver = tf.train.Saver(variables_to_restore) with self._tf_session.as_default(): saver.restore(sess=self._tf_session, save_path=flags.lanenet_detection_model_path)
def _test_lanenet(image_path, weights_path): """ :param image_path: :param weights_path: :return: """ assert ops.exists(image_path), '{:s} not exist'.format(image_path) log.info('Start reading image and preprocessing') t_start = time.time() image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = cv2.resize(image, (CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 log.info('Image load complete, cost time: {:.5f}s'.format(time.time() - t_start)) input_tensor = tf.placeholder( dtype=tf.float32, shape=[1, CFG.TRAIN.IMG_HEIGHT, CFG.TRAIN.IMG_WIDTH, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model', return_score=True) saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.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) with sess.as_default(): model_file = tf.train.latest_checkpoint(weights_path) saver.restore(sess=sess, save_path=model_file) t_start = time.time() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) t_cost = time.time() - t_start log.info('Single imgae inference cost time: {:.5f}s'.format(t_cost)) plt.figure('src_image') plt.imshow(image_vis[:, :, (2, 1, 0)]) plt.axis('off') plt.figure('binary_image') plt.imshow(binary_seg_image[0, :, :, 1] * 255, cmap='gray') plt.axis('off') plt.show() sess.close() return
def generate_on_jiqing(src_path, pred_path, weights_path): image_list = get_all_files(src_path) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') saver = tf.train.Saver() sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.TRAIN.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' # best fit with coalescing 内存管理算法 sess = tf.Session(config=sess_config) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) for index, image_path in tqdm.tqdm(enumerate(image_list), total=len(image_list)): image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = image[540:840] image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 # 归一化 (只归一未改变维数) postprocessor = lanenet_postprocess.LaneNetPostProcessor_noremap() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) postprocess_result = postprocessor.postprocess_noremap( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis, data_source='jiqing') # binary = lanenet_postprocess._morphological_process(binary_seg_image[0]) mask = postprocess_result['mask_image'] # binary = lanenet_postprocess._morphological_process(mask[0]) binary = cv2.resize(mask, (1920, 300), interpolation=cv2.INTER_LINEAR) back = np.zeros(shape=(1080, 1920, 3), dtype=np.uint8) for i in range(1920): for j in range(300): # back[j+540, i] = binary[j, i] back[j + 540, i, 0] = binary[j, i, 0] back[j + 540, i, 1] = binary[j, i, 1] back[j + 540, i, 2] = binary[j, i, 2] dst_image_path = ops.join(pred_path, image_path.split('/')[-2]) if not ops.exists(dst_image_path): os.makedirs(dst_image_path) dst_image_path = ops.join(dst_image_path, image_path.split('/')[-1]) cv2.imwrite(dst_image_path, back) return
def graph_setting(): input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') postprocessor = lanenet_postprocess.LaneNetPostProcessor()
def plot_auc(txt_file, weights_path): input_tensor = tf.placeholder( dtype=tf.float32, shape=[1, CFG.TRAIN.IMG_HEIGHT, CFG.TRAIN.IMG_WIDTH, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model', return_score=True) saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.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) with sess.as_default(): model_file = tf.train.latest_checkpoint(weights_path) saver.restore(sess=sess, save_path=model_file) gts, preds = [], [] with open(txt_file, 'r') as f: lines = f.readlines() for line in lines: image_path, gt_path, _ = line.strip().split() log.info('Processing {}.'.format(image_path)) image = cv2.imread(image_path, cv2.IMREAD_COLOR) image = cv2.resize(image, (CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 seg_image_gt = cv2.imread(gt_path, cv2.IMREAD_GRAYSCALE) seg_image_gt = cv2.resize( seg_image_gt, (CFG.TRAIN.IMG_WIDTH, CFG.TRAIN.IMG_HEIGHT), interpolation=cv2.INTER_LINEAR) seg_image_score = sess.run([binary_seg_ret], feed_dict={input_tensor: [image]}) seg_image_score = seg_image_score[0][0, :, :, 1] gts.append((seg_image_gt > 0)) preds.append(seg_image_score) sess.close() gt = np.array(gts) pred = np.array(preds) fpr, tpr, _ = roc_curve(gt.flatten(), pred.flatten()) plt.title('ROC') plt.plot(fpr, tpr) plt.xlabel('fp') plt.ylabel('tp') plt.savefig('./tboard/roc_{}.png'.format(txt_file.split('/')[-1][:-4])) return
def convert_ckpt_into_pb_file(ckpt_file_path, pb_file_path): """ :param ckpt_file_path: :param pb_file_path: :return: """ # construct compute graph with tf.compat.v1.variable_scope('lanenet'): input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', cfg=CFG) binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='LaneNet') with tf.compat.v1.variable_scope('lanenet/'): binary_seg_ret = tf.cast(binary_seg_ret, dtype=tf.float32) binary_seg_ret = tf.squeeze(binary_seg_ret, axis=0, name='final_binary_output') instance_seg_ret = tf.squeeze(instance_seg_ret, axis=0, name='final_pixel_embedding_output') # define moving average version of the learned variables for eval with tf.compat.v1.variable_scope(name_or_scope='moving_avg'): variable_averages = tf.train.ExponentialMovingAverage( CFG.SOLVER.MOVING_AVE_DECAY) variables_to_restore = variable_averages.variables_to_restore() # create a session saver = tf.train.Saver(variables_to_restore) sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = 0.85 sess_config.gpu_options.allow_growth = False sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) with sess.as_default(): saver.restore(sess, ckpt_file_path) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=[ 'lanenet/input_tensor', 'lanenet/final_binary_output', 'lanenet/final_pixel_embedding_output' ]) with tf.gfile.GFile(pb_file_path, "wb") as f: f.write(converted_graph_def.SerializeToString())
def generate(src_path, pred_path, weights_path): image_list = get_all_files(src_path) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') postprocessor = lanenet_postprocess.LaneNetPostProcessor_noremap() saver = tf.train.Saver() sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.TRAIN.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' # best fit with coalescing 内存管理算法 sess = tf.Session(config=sess_config) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) for image_path in image_list: image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image # image = image[:][300:850] # image = cv2.resize(image, (1920, 1280), interpolation=cv2.INTER_LINEAR) image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 # 归一化 (只归一未改变维数) t_strat = time.time() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) t_inference = time.time() - t_strat postprocess_result = postprocessor.postprocess_noremap( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis, data_source='caltech') t_cost = time.time() - t_strat print(t_inference, t_cost) dst_image_path = ops.join(pred_path, image_path.split('/')[-2]) if not ops.exists(dst_image_path): os.makedirs(dst_image_path) dst_image_path = ops.join(dst_image_path, image_path.split('/')[-1]) cv2.imwrite(dst_image_path, image_vis) # postprocess_result['mask_image']) return
def generate_tusimple_sample(scr_path, pred_path, weights_path): listOfFiles = list() for (dirpath, dirnames, filenames) in os.walk(scr_path): listOfFiles += [os.path.join(dirpath, file) for file in filenames] print(listOfFiles) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') postprocessor = lanenet_postprocess.LaneNetPostProcessor_noremap() saver = tf.train.Saver() sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.TRAIN.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' # best fit with coalescing 内存管理算法 sess = tf.Session(config=sess_config) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) for index, image_path in tqdm.tqdm(enumerate(listOfFiles), total=len(listOfFiles)): image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 # 归一化 (只归一未改变维数) binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) postprocess_result = postprocessor.postprocess_noremap( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis, ) dst_image_path = ops.join(pred_path, image_path.split('/')[-2]) if not ops.exists(dst_image_path): os.makedirs(dst_image_path) dst_image_path = ops.join(dst_image_path, image_path.split('/')[-1]) cv2.imwrite(dst_image_path, image_vis) return
def convert_ckpt_into_pb_file(ckpt_file_path, pb_file_path): """ :param ckpt_file_path: :param pb_file_path: :return: """ # construct compute graph with tf.variable_scope('lanenet'): input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') with tf.variable_scope('lanenet/'): binary_seg_ret = tf.cast(binary_seg_ret, dtype=tf.float32) binary_seg_ret = tf.squeeze(binary_seg_ret, axis=0, name='final_binary_output') instance_seg_ret = tf.squeeze(instance_seg_ret, axis=0, name='final_pixel_embedding_output') # create a session saver = tf.train.Saver() sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = 0.85 sess_config.gpu_options.allow_growth = False sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) with sess.as_default(): saver.restore(sess, ckpt_file_path) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=[ 'lanenet/input_tensor', 'lanenet/final_binary_output', 'lanenet/final_pixel_embedding_output' ]) with tf.gfile.GFile(pb_file_path, "wb") as f: f.write(converted_graph_def.SerializeToString())
def optimize_model(model_path): inputGraph = tf.GraphDef() with tf.gfile.Open(model_path, 'rb') as model: data2read = model.read() inputGraph.ParseFromString(data2read) with tf.variable_scope('lanenet'): input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') with tf.variable_scope('lanenet/'): binary_seg_ret = tf.cast(binary_seg_ret, dtype=tf.float32) binary_seg_ret = tf.identity(binary_seg_ret, name= 'final_binary_output') instance_seg_ret = tf.identity(instance_seg_ret, name='final_pixel_embedding_output') # binary_seg_ret = tf.squeeze(binary_seg_ret, axis=0, name='final_binary_output') # 删除所有大小是 1 的维度 # instance_seg_ret = tf.squeeze(instance_seg_ret, axis=0, name='final_pixel_embedding_output') """outputGraph = optimize_for_inference_lib.optimize_for_inference( inputGraph, input_node_names=['lanenet/input_tensor'], output_node_names=[ 'lanenet/final_binary_output', 'lanenet/final_pixel_embedding_output'], placeholder_type_enum=tf.int32.as_datatype_enum )""" outputGraph = TransformGraph( inputGraph, ['lanenet/input_tensor'], ['lanenet/final_binary_output', 'lanenet/final_pixel_embedding_output'], ['remove_nodes(op=Identity, op=CheckNumerics)', 'merge_duplicate_nodes', 'strip_unused_nodes', 'fold_constants(ignore_errors=true)', 'fold_batch_norms', 'fold_old_batch_norms', 'quantize_weights', 'quantize_nodes', 'sort_by_execution_order'] ) new_name = model_path.split('/')[-2] + '/OptimizedGraph.pb' model = tf.gfile.FastGFile(new_name, 'w') model.write(outputGraph.SerializeToString())
def test_lanenet(image_path, weights_path): """ :param image_path: :param weights_path: :return: """ assert ops.exists(image_path), '{:s} not exist'.format(image_path) log.info('Start reading image and preprocessing') t_start = time.time() image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 log.info('Image load complete, cost time: {:.5f}s'.format(time.time() - t_start)) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') n_images = 0 total_runtime = 0 for i in range(1000): start_time = time.time() binary_seg_ret, instance_seg_ret = net.inference( input_tensor=input_tensor, name='lanenet_model') # mask = binary_seg_ret[0] # mask *= 255 # mask = mask.astype(np.uint8) run_time = time.time() - start_time total_runtime += run_time n_images += 1 avg_time = total_runtime / n_images print("Avg time: " + str(avg_time) + " - FPS: " + str(1 / avg_time)) return
def __init__(self, weights_path): with tf.device('/cpu:0'): self.input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') self.binary_seg_ret, self.instance_seg_ret = net.inference( input_tensor=self.input_tensor, name='lanenet_model') self.postprocessor = lanenet_postprocess.LaneNetPostProcessor() saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.TRAIN.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' self.sess = tf.Session(config=sess_config) with self.sess.as_default(): saver.restore(sess=self.sess, save_path=weights_path)
pb_file = "model/tusimple_lanenet_mobilenet_v2_1005/culane_lanenet_mobilenet_v2_1213.pb" #ckpt_file = cfg.YOLO.DEMO_WEIGHT ckpt_file = 'model/tusimple_lanenet_mobilenet_v2_1005/culane_lanenet_mobilenet_v2_1005_reduce_train.ckpt' # output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"] # #output_node_names = ["input_tensor", "lanenet_loss/inference/decode/score_final/Conv2D", "lanenet_loss/pix_embedding_conv/Conv2D"] output_node_names = [ "input_tensor", "lanenet_model/mobilenet_v2_backend/binary_seg/ArgMax", "lanenet_model/mobilenet_v2_backend/instance_seg/pix_embedding_conv/Conv2D" ] input_data = tf.placeholder(dtype=tf.float32, name='input_tensor', shape=[None, None, None, 3]) phase_tensor = tf.constant('false', tf.string) net = lanenet.LaneNet(phase=phase_tensor, net_flag='mobilenet_v2') model = net.inference(input_data, name='lanenet_model') #print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox) sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) saver = tf.train.Saver() saver.restore(sess, ckpt_file) converted_graph_def = tf.graph_util.convert_variables_to_constants( sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) with tf.gfile.GFile(pb_file, "wb") as f: f.write(converted_graph_def.SerializeToString())
def test_lanenet(image_path, weights_path): """ :param image_path: :param weights_path: :return: """ assert ops.exists(image_path), '{:s} not exist'.format(image_path) log.info('Start reading image and preprocessing') t_start = time.time() image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 log.info('Image load complete, cost time: {:.5f}s'.format(time.time() - t_start)) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') postprocessor = lanenet_postprocess.LaneNetPostProcessor() saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.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) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) writer = tf.summary.FileWriter('logs', sess.graph) t_start = time.time() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]} ) writer.close() t_cost = time.time() - t_start log.info('Single imgae inference cost time: {:.5f}s'.format(t_cost)) postprocess_result = postprocessor.postprocess( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis ) mask_image = postprocess_result['mask_image'] for i in range(CFG.TRAIN.EMBEDDING_FEATS_DIMS): instance_seg_image[0][:, :, i] = minmax_scale(instance_seg_image[0][:, :, i]) embedding_image = np.array(instance_seg_image[0], np.uint8) print([n.name for n in sess.graph.as_graph_def().node]) #output_node_names = "lanenet_model/vgg_backend/instance_seg/pix_embedding_conv/Conv2D" output_node_names = ["lanenet_model/vgg_backend/binary_seg/ArgMax"] output_node_names = ["lanenet_model/vgg_backend/instance_seg/pix_embedding_conv/Conv2D", "lanenet_model/vgg_backend/binary_seg/ArgMax"] input_graph_def = sess.graph.as_graph_def() # output_node_names="train_IteratorGetNext" output_graph_def = graph_util.convert_variables_to_constants( sess, # The session input_graph_def, # input_graph_def is useful for retrieving the nodes output_node_names) import tensorflow.contrib.tensorrt as trt trt_graph = trt.create_inference_graph( input_graph_def=output_graph_def, outputs=output_node_names, max_batch_size=1, max_workspace_size_bytes=1 << 25, precision_mode='INT8', minimum_segment_size=50 ) with open('lanenet_trt_2outputs.pb', 'wb') as f: f.write(trt_graph.SerializeToString()) plt.figure('mask_image') plt.imshow(mask_image[:, :, (2, 1, 0)]) plt.figure('src_image') plt.imshow(image_vis[:, :, (2, 1, 0)]) plt.figure('instance_image') plt.imshow(embedding_image[:, :, (2, 1, 0)]) plt.figure('binary_image') plt.imshow(binary_seg_image[0] * 255, cmap='gray') plt.show() cv2.imwrite('instance_mask_image.png', mask_image) cv2.imwrite('source_image.png', postprocess_result['source_image']) cv2.imwrite('binary_mask_image.png', binary_seg_image[0] * 255) sess.close() return
def __init__(self): """ initialize lanenet multi gpu trainner """ # define solver params and dataset self._train_dataset = lanenet_data_feed_pipline.LaneNetDataFeeder( flags='train') self._val_dataset = lanenet_data_feed_pipline.LaneNetDataFeeder( flags='val') self._steps_per_epoch = len(self._train_dataset) self._val_steps_per_epoch = len(self._val_dataset) self._model_name = '{:s}_{:s}'.format(CFG.MODEL.FRONT_END, CFG.MODEL.MODEL_NAME) self._train_epoch_nums = CFG.TRAIN.EPOCH_NUMS self._batch_size = CFG.TRAIN.BATCH_SIZE self._val_batch_size = CFG.TRAIN.VAL_BATCH_SIZE self._snapshot_epoch = CFG.TRAIN.SNAPSHOT_EPOCH self._model_save_dir = ops.join(CFG.TRAIN.MODEL_SAVE_DIR, self._model_name) self._tboard_save_dir = ops.join(CFG.TRAIN.TBOARD_SAVE_DIR, self._model_name) self._enable_miou = CFG.TRAIN.COMPUTE_MIOU.ENABLE if self._enable_miou: self._record_miou_epoch = CFG.TRAIN.COMPUTE_MIOU.EPOCH self._input_tensor_size = [int(tmp) for tmp in CFG.AUG.TRAIN_CROP_SIZE] self._gpu_devices = CFG.TRAIN.MULTI_GPU.GPU_DEVICES self._gpu_nums = len(self._gpu_devices) self._chief_gpu_index = CFG.TRAIN.MULTI_GPU.CHIEF_DEVICE_INDEX self._batch_size_per_gpu = int(self._batch_size / self._gpu_nums) self._init_learning_rate = CFG.SOLVER.LR self._moving_ave_decay = CFG.SOLVER.MOVING_AVE_DECAY self._momentum = CFG.SOLVER.MOMENTUM self._lr_polynimal_decay_power = CFG.SOLVER.LR_POLYNOMIAL_POWER self._optimizer_mode = CFG.SOLVER.OPTIMIZER.lower() if CFG.TRAIN.RESTORE_FROM_SNAPSHOT.ENABLE: self._initial_weight = CFG.TRAIN.RESTORE_FROM_SNAPSHOT.SNAPSHOT_PATH else: self._initial_weight = None if CFG.TRAIN.WARM_UP.ENABLE: self._warmup_epoches = CFG.TRAIN.WARM_UP.EPOCH_NUMS self._warmup_init_learning_rate = self._init_learning_rate / 1000.0 else: self._warmup_epoches = 0 # define tensorflow session sess_config = tf.ConfigProto(allow_soft_placement=True) sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' self._sess = tf.Session(config=sess_config) # define graph input tensor with tf.variable_scope(name_or_scope='graph_input_node'): self._input_src_image_list = [] self._input_binary_label_image_list = [] self._input_instance_label_image_list = [] for i in range(self._gpu_nums): src_imgs, binary_label_imgs, instance_label_imgs = self._train_dataset.next_batch( batch_size=self._batch_size_per_gpu) self._input_src_image_list.append(src_imgs) self._input_binary_label_image_list.append(binary_label_imgs) self._input_instance_label_image_list.append( instance_label_imgs) self._val_input_src_image, self._val_input_binary_label_image, self._val_input_instance_label_image = \ self._val_dataset.next_batch(batch_size=self._val_batch_size) # define model self._model = lanenet.LaneNet(phase='train') self._val_model = lanenet.LaneNet(phase='test') # define average container tower_grads = [] tower_total_loss = [] tower_binary_seg_loss = [] tower_instance_seg_loss = [] batchnorm_updates = None # define learning rate with tf.variable_scope('learning_rate'): self._global_step = tf.Variable(1.0, dtype=tf.float32, trainable=False, name='global_step') self._val_global_step = tf.Variable(1.0, dtype=tf.float32, trainable=False, name='val_global_step') self._val_global_step_update = tf.assign_add( self._val_global_step, 1.0) warmup_steps = tf.constant(self._warmup_epoches * self._steps_per_epoch, dtype=tf.float32, name='warmup_steps') train_steps = tf.constant(self._train_epoch_nums * self._steps_per_epoch, dtype=tf.float32, name='train_steps') self._learn_rate = tf.cond( pred=self._global_step < warmup_steps, true_fn=lambda: self._compute_warmup_lr( warmup_steps=warmup_steps, name='warmup_lr'), false_fn=lambda: tf.train.polynomial_decay( learning_rate=self._init_learning_rate, global_step=self._global_step, decay_steps=train_steps, end_learning_rate=0.000000001, power=self._lr_polynimal_decay_power)) self._learn_rate = tf.identity(self._learn_rate, 'lr') # define optimizer if self._optimizer_mode == 'sgd': optimizer = tf.train.MomentumOptimizer( learning_rate=self._learn_rate, momentum=self._momentum) elif self._optimizer_mode == 'adam': optimizer = tf.train.AdamOptimizer( learning_rate=self._learn_rate, ) else: raise NotImplementedError( 'Not support optimizer: {:s} for now'.format( self._optimizer_mode)) # define distributed train op with tf.variable_scope(tf.get_variable_scope()): is_network_initialized = False for i in range(self._gpu_nums): with tf.device('/gpu:{:d}'.format(i)): with tf.name_scope('tower_{:d}'.format(i)) as _: input_images = self._input_src_image_list[i] input_binary_labels = self._input_binary_label_image_list[ i] input_instance_labels = self._input_instance_label_image_list[ i] tmp_loss, tmp_grads = self._compute_net_gradients( input_images, input_binary_labels, input_instance_labels, optimizer, is_net_first_initialized=is_network_initialized) is_network_initialized = True # Only use the mean and var in the chief gpu tower to update the parameter if i == self._chief_gpu_index: batchnorm_updates = tf.get_collection( tf.GraphKeys.UPDATE_OPS) tower_grads.append(tmp_grads) tower_total_loss.append(tmp_loss['total_loss']) tower_binary_seg_loss.append( tmp_loss['binary_seg_loss']) tower_instance_seg_loss.append( tmp_loss['discriminative_loss']) grads = self._average_gradients(tower_grads) self._loss = tf.reduce_mean(tower_total_loss, name='reduce_mean_tower_total_loss') self._binary_loss = tf.reduce_mean( tower_binary_seg_loss, name='reduce_mean_tower_binary_loss') self._instance_loss = tf.reduce_mean( tower_instance_seg_loss, name='reduce_mean_tower_instance_loss') ret = self._val_model.compute_loss( input_tensor=self._val_input_src_image, binary_label=self._val_input_binary_label_image, instance_label=self._val_input_instance_label_image, name='LaneNet', reuse=True) self._val_loss = ret['total_loss'] self._val_binary_loss = ret['binary_seg_loss'] self._val_instance_loss = ret['discriminative_loss'] # define moving average op with tf.variable_scope(name_or_scope='moving_avg'): if CFG.TRAIN.FREEZE_BN.ENABLE: train_var_list = [ v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name ] else: train_var_list = tf.trainable_variables() moving_ave_op = tf.train.ExponentialMovingAverage( self._moving_ave_decay).apply(train_var_list + tf.moving_average_variables()) # group all the op needed for training batchnorm_updates_op = tf.group(*batchnorm_updates) apply_gradient_op = optimizer.apply_gradients( grads, global_step=self._global_step) self._train_op = tf.group(apply_gradient_op, moving_ave_op, batchnorm_updates_op) # define prediction self._binary_prediciton, self._instance_prediciton = self._model.inference( input_tensor=self._input_src_image_list[self._chief_gpu_index], name='LaneNet', reuse=True) self._binary_prediciton = tf.identity( self._binary_prediciton, name='binary_segmentation_result') self._val_binary_prediction, self._val_instance_prediciton = self._val_model.inference( input_tensor=self._val_input_src_image, name='LaneNet', reuse=True) self._val_binary_prediction = tf.identity( self._val_binary_prediction, name='val_binary_segmentation_result') # define miou if self._enable_miou: with tf.variable_scope('miou'): pred = tf.reshape(self._binary_prediciton, [ -1, ]) gt = tf.reshape( self._input_binary_label_image_list[self._chief_gpu_index], [ -1, ]) indices = tf.squeeze( tf.where(tf.less_equal(gt, CFG.DATASET.NUM_CLASSES - 1)), 1) gt = tf.gather(gt, indices) pred = tf.gather(pred, indices) self._miou, self._miou_update_op = tf.metrics.mean_iou( labels=gt, predictions=pred, num_classes=CFG.DATASET.NUM_CLASSES) val_pred = tf.reshape(self._val_binary_prediction, [ -1, ]) val_gt = tf.reshape(self._val_input_binary_label_image, [ -1, ]) indices = tf.squeeze( tf.where(tf.less_equal(val_gt, CFG.DATASET.NUM_CLASSES - 1)), 1) val_gt = tf.gather(val_gt, indices) val_pred = tf.gather(val_pred, indices) self._val_miou, self._val_miou_update_op = tf.metrics.mean_iou( labels=val_gt, predictions=val_pred, num_classes=CFG.DATASET.NUM_CLASSES) # define saver and loader with tf.variable_scope('loader_and_saver'): self._net_var = [ vv for vv in tf.global_variables() if 'lr' not in vv.name ] self._loader = tf.train.Saver(self._net_var) self._saver = tf.train.Saver(max_to_keep=10) # define summary with tf.variable_scope('summary'): summary_merge_list = [ tf.summary.scalar("learn_rate", self._learn_rate), tf.summary.scalar("total_loss", self._loss), tf.summary.scalar('binary_loss', self._binary_loss), tf.summary.scalar('instance_loss', self._instance_loss), ] val_summary_merge_list = [ tf.summary.scalar('val_total_loss', self._val_loss), tf.summary.scalar('val_binary_loss', self._val_binary_loss), tf.summary.scalar('val_instance_loss', self._val_instance_loss), ] if self._enable_miou: with tf.control_dependencies([self._miou_update_op]): summary_merge_list_with_miou = [ tf.summary.scalar("learn_rate", self._learn_rate), tf.summary.scalar("total_loss", self._loss), tf.summary.scalar('binary_loss', self._binary_loss), tf.summary.scalar('instance_loss', self._instance_loss), tf.summary.scalar('miou', self._miou) ] self._write_summary_op_with_miou = tf.summary.merge( summary_merge_list_with_miou) with tf.control_dependencies( [self._val_miou_update_op, self._val_global_step_update]): val_summary_merge_list_with_miou = [ tf.summary.scalar("total_loss", self._loss), tf.summary.scalar('binary_loss', self._binary_loss), tf.summary.scalar('instance_loss', self._instance_loss), tf.summary.scalar('val_miou', self._val_miou), ] self._val_write_summary_op_with_miou = tf.summary.merge( val_summary_merge_list_with_miou) if ops.exists(self._tboard_save_dir): shutil.rmtree(self._tboard_save_dir) os.makedirs(self._tboard_save_dir, exist_ok=True) model_params_file_save_path = ops.join( self._tboard_save_dir, CFG.TRAIN.MODEL_PARAMS_CONFIG_FILE_NAME) with open(model_params_file_save_path, 'w', encoding='utf-8') as f_obj: CFG.dump_to_json_file(f_obj) self._write_summary_op = tf.summary.merge(summary_merge_list) self._val_write_summary_op = tf.summary.merge( val_summary_merge_list) self._summary_writer = tf.summary.FileWriter( self._tboard_save_dir, graph=self._sess.graph) LOG.info('Initialize tusimple lanenet multi gpu trainner complete')
def evaluate_lanenet_accuracy_ckpt(data_path, weights_path): input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') postprocessor = lanenet_postprocess.LaneNetPostProcessor() saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.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) gt_image_path = ops.join(data_path, 'gt_image') gt_binary_path = ops.join(data_path, 'gt_binary_image') accuracy_sum = 0.0 fp_sum = 0.0 fn_sum = 0.0 with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) image_path_list = os.listdir(gt_image_path) i = 0 for image_path in image_path_list: image_path = ops.join(gt_image_path, image_path) print(image_path) image = cv2.imread(image_path, cv2.IMREAD_COLOR) # image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 # 归一化 (只归一未改变维数) binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) image_name = ops.split(image_path)[1] gt_image = cv2.imread(ops.join(gt_binary_path, image_name), cv2.IMREAD_COLOR) image_vis = gt_image postprocess_result = postprocessor.postprocess_for_test( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis) plt.figure("result") plt.imshow(postprocess_result['mask_image']) plt.show() # binary_seg_image = binary_seg_image.astype((np.float32)) # print(binary_seg_image.shape) # binary_seg_image = binary_seg_image.reshape(1, 256, 512, 1) # binary_seg_image = cv2.resize(postprocess_result['source_image'], (512, 256), interpolation=cv2.INTER_LINEAR) binary_seg_image = postprocess_result['source_image'].reshape( 1, 720, 1280, 1) # gt_image = cv2.resize(gt_image, (512, 256), interpolation=cv2.INTER_LINEAR) gt_image = gt_image[:, :, 0].reshape(720, 1280, 1) binary_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 720, 1280, 1], name='binary_tensor') gt_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 720, 1280, 1], name='gt_tensor') accuracy, count, idx = evaluate_model_utils.calculate_model_precision_for_test( binary_tensor, gt_tensor) fn = evaluate_model_utils.calculate_model_fn_for_test( binary_tensor, gt_tensor) fp = evaluate_model_utils.calculate_model_fp_for_test( binary_tensor, gt_tensor) accuracy_result, count_result, idx = sess.run( [accuracy, count, idx], feed_dict={ binary_tensor: binary_seg_image, gt_tensor: [gt_image] }) fn_result = sess.run(fn, feed_dict={ binary_tensor: binary_seg_image, gt_tensor: [gt_image] }) fp_result = sess.run(fp, feed_dict={ binary_tensor: binary_seg_image, gt_tensor: [gt_image] }) print("accuracy:", accuracy_result, count_result, idx) print("fn:", fn_result) print("fp:", fp_result) i += 1 accuracy_sum += accuracy_result fp_sum += fp_result fn_sum += fn_result if i == 100: print("average_accuracy: ", accuracy_sum / 100.0) print("average_fn: ", fn_sum / 100.0) print("average_fp: ", fp_sum / 100.0) break return
ax4.set_title("Result Image") plt.show() """ return res_img sess.close() if __name__ == '__main__': weights_path = "model/tusimple_lanenet/tusimple_lanenet.ckpt" # placeholder 자료형의 변수 생성(텐서를 placeholder에 맵핑) -> 입력텐서 input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', cfg=CFG) # 딥러닝 네트워크 그래프 설정, 딥러닝 추론위한 텐서 생성 # 텐서플로우 그래프 빌드, 그래프는 세션 실행시 실제로 실행되는 내용 정의 binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='LaneNet') postprocessor = lanenet_postprocess.LaneNetPostProcessor( cfg=CFG) # 이미지 후처리 하기위해 postprocessor 객체 생성 # Set sess configuration (텐서플로우 GPU메모리 할당 설정) sess_config = tf.ConfigProto() # 세션설정정보 저장하기 위한 객체생성 sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH # 상황에 맞게 자동으로 메모리 할당 sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) # 세션 설정 이용하여 세션 생성하여 저장
def test_lanenet(video_path, weights_path): """ :param image_path: :param weights_path: :return: """ assert ops.exists(video_path), '{:s} not exist'.format(video_path) obj_w = 1920 obj_h = 320 input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, obj_h, obj_w, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', cfg=CFG) binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='LaneNet') postprocessor = lanenet_postprocess.LaneNetPostProcessor(cfg=CFG) # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) # define moving average version of the learned variables for eval with tf.variable_scope(name_or_scope='moving_avg'): variable_averages = tf.train.ExponentialMovingAverage( CFG.SOLVER.MOVING_AVE_DECAY) variables_to_restore = variable_averages.variables_to_restore() # define saver saver = tf.train.Saver(variables_to_restore) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) cap = cv2.VideoCapture(video_path) fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) log.info('width: %d, height: %d', width, height) out = cv2.VideoWriter('out.mp4', fourcc, 15, (obj_w, obj_h), isColor=True) nn = 1000 #while(cap.isOpened() and nn >= 0): while (cap.isOpened()): nn = nn - 1 ret, image = cap.read() if ret: #image = cv2.imread('/workspace/lanenet-lane-detection/1410486660.jpg') # image = cv2.resize(image, (1280, 720), interpolation=cv2.INTER_LINEAR) image = image[442:762, :] image_vis = image.copy() image = image / 127.5 - 1.0 t_start = time.time() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) t_cost = time.time() - t_start LOG.info( 'Single image inference cost time: {:.5f}s'.format(t_cost)) #log.info(instance_seg_image.shape) #mask_image = gen_color_img(binary_seg_image[0], instance_seg_image[0], 5) #cv2.imwrite('test.png', mask_image) #return postprocess_result = postprocessor.postprocess2( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis) binary_image = postprocess_result['binary_image'] instance_image = postprocess_result['instance_image'] ''' np.set_printoptions(suppress = True, precision = 2, threshold = np.inf) print('before: ') print(instance_seg_image[0]) for i in range(CFG.TRAIN.EMBEDDING_FEATS_DIMS): instance_seg_image[0][:, :, i] = minmax_scale(instance_seg_image[0][:, :, i]) embedding_image = np.array(instance_seg_image[0], np.uint8) print('after: ') print(embedding_image) ''' out.write(np.uint8(instance_image)) #out.write(np.uint8(mask_image)) #cv2.imwrite('test.png', binary_image) #return else: break cap.release() out.release() sess.close() return
def test_lanenet_batch(src_dir, weights_path, save_dir, net_flag): """ :param src_dir: :param weights_path: :param save_dir: :return: """ assert ops.exists(src_dir), '{:s} not exist'.format(src_dir) os.makedirs(save_dir, exist_ok=True) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag=net_flag) binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') postprocessor = lanenet_postprocess.LaneNetPostProcessor() saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.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) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) image_list = glob.glob('{:s}/**/*.jpg'.format(src_dir), recursive=True) avg_time_cost = [] for index, image_path in tqdm.tqdm(enumerate(image_list), total=len(image_list)): image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 t_start = time.time() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) avg_time_cost.append(time.time() - t_start) postprocess_result = postprocessor.postprocess( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis) if index % 100 == 0: log.info( 'Mean inference time every single image: {:.5f}s'.format( np.mean(avg_time_cost))) avg_time_cost.clear() input_image_dir = ops.split(image_path.split('clips')[1])[0][1:] input_image_name = ops.split(image_path)[1] output_image_dir = ops.join(save_dir, input_image_dir) os.makedirs(output_image_dir, exist_ok=True) output_image_path = ops.join(output_image_dir, input_image_name) if ops.exists(output_image_path): continue cv2.imwrite(output_image_path, postprocess_result['source_image']) return
def test_lanenet_batch(image_list, weights_path, batch_size, use_gpu, net_flag='vgg'): """ :param image_list: :param weights_path: :param batch_size: :param use_gpu: :param net_flag: :return: """ assert ops.exists(image_list), '{:s} not exist'.format(image_list) log.info('开始加载数据集列表...') test_dataset = lanenet_data_processor.DataSet(image_list, traing=False) # ============================== gt_label_binary_list = [] with open(image_list, 'r') as file: for _info in file: info_tmp = _info.strip(' ').split() gt_label_binary_list.append(info_tmp[1]) # ============================== input_tensor = tf.placeholder(dtype=tf.float32, shape=[None, 256, 512, 3], name='input_tensor') binary_label_tensor = tf.placeholder(dtype=tf.int64, shape=[None, 256, 512, 1], name='binary_input_label') phase_tensor = tf.constant('test', tf.string) net = lanenet.LaneNet(phase=phase_tensor, net_flag=net_flag) binary_seg_ret, instance_seg_ret, recall_ret, false_positive, false_negative, precision_ret, accuracy_ret = \ net.compute_acc(input_tensor=input_tensor, binary_label_tensor=binary_label_tensor, name='lanenet_model') saver = tf.train.Saver() # ============================== # Set sess configuration if use_gpu: sess_config = tf.ConfigProto(device_count={'GPU': 1}) else: sess_config = tf.ConfigProto(device_count={'GPU': 0}) sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.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) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) epoch_nums = int(math.ceil(test_dataset._dataset_size / batch_size)) mean_accuracy = 0.0 mean_recall = 0.0 mean_precision = 0.0 mean_fp = 0.0 mean_fn = 0.0 total_num = 0 t_start = time.time() for epoch in range(epoch_nums): gt_imgs, binary_gt_labels, instance_gt_labels = test_dataset.next_batch(batch_size) if net_flag == 'vgg': image_list_epoch = [tmp / 127.5 - 1.0 for tmp in gt_imgs] elif net_flag == 'mobilenet_v2': image_list_epoch = [tmp - [103.939, 116.779, 123.68] for tmp in gt_imgs] binary_seg_images, instance_seg_images, recall, fp, fn, precision, accuracy = sess.run( [binary_seg_ret, instance_seg_ret, recall_ret, false_positive, false_negative, precision_ret, accuracy_ret], feed_dict={input_tensor: image_list_epoch, binary_label_tensor: binary_gt_labels}) # ============================== out_dir = 'H:/Other_DataSets/TuSimple/out/' dst_binary_image_path = ops.join(out_dir,gt_label_binary_list[epoch]) root_dir = ops.dirname(ops.abspath(dst_binary_image_path)) if not os.path.exists(root_dir): os.makedirs(root_dir) cv2.imwrite(dst_binary_image_path, binary_seg_images[0] * 255) # ============================== print(recall, fp, fn) mean_accuracy += accuracy mean_precision += precision mean_recall += recall mean_fp += fp mean_fn += fn total_num += len(gt_imgs) t_cost = time.time() - t_start mean_accuracy = mean_accuracy / epoch_nums mean_precision = mean_precision / epoch_nums mean_recall = mean_recall / epoch_nums mean_fp = mean_fp / epoch_nums mean_fn = mean_fn / epoch_nums print('测试 {} 张图片,耗时{},{}_recall = {}, precision = {}, accuracy = {}, fp = {}, fn = {}, '.format( total_num, t_cost, net_flag, mean_recall, mean_precision, mean_accuracy, mean_fp, mean_fn)) sess.close()
except Exception as e: print("Unable to show image: "+str(e)) cv2.waitKey(1) detect_object_alive = False """ :param image_path: :param weights_path: :return: """ input_tensor = tf.placeholder(dtype=tf.float32, shape=[ 1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference( input_tensor=input_tensor, name='lanenet_model') saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.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) saver.restore(sess=sess, save_path=weights_path) def inference_net(): global detect_lane_alive,image_0 detect_lane_alive = True # preprocess
def train_lanenet(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 = lanenet_data_feed_pipline.LaneNetDataFeeder( dataset_dir=dataset_dir, flags='train') val_dataset = lanenet_data_feed_pipline.LaneNetDataFeeder( dataset_dir=dataset_dir, flags='val') with tf.device('/gpu:1'): # set lanenet train_net = lanenet.LaneNet(net_flag=net_flag, phase='train', reuse=False) val_net = lanenet.LaneNet(net_flag=net_flag, phase='val', reuse=True) # set compute graph node for training train_images, train_binary_labels, train_instance_labels = train_dataset.inputs( CFG.TRAIN.BATCH_SIZE, 1) train_compute_ret = train_net.compute_loss( input_tensor=train_images, binary_label=train_binary_labels, instance_label=train_instance_labels, name='lanenet_model') train_total_loss = train_compute_ret['total_loss'] train_binary_seg_loss = train_compute_ret['binary_seg_loss'] train_disc_loss = train_compute_ret['discriminative_loss'] train_pix_embedding = train_compute_ret['instance_seg_logits'] train_prediction_logits = train_compute_ret['binary_seg_logits'] train_prediction_score = tf.nn.softmax(logits=train_prediction_logits) train_prediction = tf.argmax(train_prediction_score, axis=-1) train_accuracy = evaluate_model_utils.calculate_model_precision( train_compute_ret['binary_seg_logits'], train_binary_labels) train_fp = evaluate_model_utils.calculate_model_fp( train_compute_ret['binary_seg_logits'], train_binary_labels) train_fn = evaluate_model_utils.calculate_model_fn( train_compute_ret['binary_seg_logits'], train_binary_labels) train_binary_seg_ret_for_summary = evaluate_model_utils.get_image_summary( img=train_prediction) train_embedding_ret_for_summary = evaluate_model_utils.get_image_summary( img=train_pix_embedding) train_cost_scalar = tf.summary.scalar(name='train_cost', tensor=train_total_loss) train_accuracy_scalar = tf.summary.scalar(name='train_accuracy', tensor=train_accuracy) train_binary_seg_loss_scalar = tf.summary.scalar( name='train_binary_seg_loss', tensor=train_binary_seg_loss) train_instance_seg_loss_scalar = tf.summary.scalar( name='train_instance_seg_loss', tensor=train_disc_loss) train_fn_scalar = tf.summary.scalar(name='train_fn', tensor=train_fn) train_fp_scalar = tf.summary.scalar(name='train_fp', tensor=train_fp) train_binary_seg_ret_img = tf.summary.image( name='train_binary_seg_ret', tensor=train_binary_seg_ret_for_summary) train_embedding_feats_ret_img = tf.summary.image( name='train_embedding_feats_ret', tensor=train_embedding_ret_for_summary) train_merge_summary_op = tf.summary.merge([ train_accuracy_scalar, train_cost_scalar, train_binary_seg_loss_scalar, train_instance_seg_loss_scalar, train_fn_scalar, train_fp_scalar, train_binary_seg_ret_img, train_embedding_feats_ret_img ]) # set compute graph node for validation val_images, val_binary_labels, val_instance_labels = val_dataset.inputs( CFG.TRAIN.VAL_BATCH_SIZE, 1) val_compute_ret = val_net.compute_loss( input_tensor=val_images, binary_label=val_binary_labels, instance_label=val_instance_labels, name='lanenet_model') val_total_loss = val_compute_ret['total_loss'] val_binary_seg_loss = val_compute_ret['binary_seg_loss'] val_disc_loss = val_compute_ret['discriminative_loss'] val_pix_embedding = val_compute_ret['instance_seg_logits'] val_prediction_logits = val_compute_ret['binary_seg_logits'] val_prediction_score = tf.nn.softmax(logits=val_prediction_logits) val_prediction = tf.argmax(val_prediction_score, axis=-1) val_accuracy = evaluate_model_utils.calculate_model_precision( val_compute_ret['binary_seg_logits'], val_binary_labels) val_fp = evaluate_model_utils.calculate_model_fp( val_compute_ret['binary_seg_logits'], val_binary_labels) val_fn = evaluate_model_utils.calculate_model_fn( val_compute_ret['binary_seg_logits'], val_binary_labels) val_binary_seg_ret_for_summary = evaluate_model_utils.get_image_summary( img=val_prediction) val_embedding_ret_for_summary = evaluate_model_utils.get_image_summary( img=val_pix_embedding) val_cost_scalar = tf.summary.scalar(name='val_cost', tensor=val_total_loss) val_accuracy_scalar = tf.summary.scalar(name='val_accuracy', tensor=val_accuracy) val_binary_seg_loss_scalar = tf.summary.scalar( name='val_binary_seg_loss', tensor=val_binary_seg_loss) val_instance_seg_loss_scalar = tf.summary.scalar( name='val_instance_seg_loss', tensor=val_disc_loss) val_fn_scalar = tf.summary.scalar(name='val_fn', tensor=val_fn) val_fp_scalar = tf.summary.scalar(name='val_fp', tensor=val_fp) val_binary_seg_ret_img = tf.summary.image( name='val_binary_seg_ret', tensor=val_binary_seg_ret_for_summary) val_embedding_feats_ret_img = tf.summary.image( name='val_embedding_feats_ret', tensor=val_embedding_ret_for_summary) val_merge_summary_op = tf.summary.merge([ val_accuracy_scalar, val_cost_scalar, val_binary_seg_loss_scalar, val_instance_seg_loss_scalar, val_fn_scalar, val_fp_scalar, val_binary_seg_ret_img, val_embedding_feats_ret_img ]) # set optimizer global_step = tf.Variable(0, trainable=False) learning_rate = tf.train.polynomial_decay( learning_rate=CFG.TRAIN.LEARNING_RATE, global_step=global_step, decay_steps=CFG.TRAIN.EPOCHS, power=0.9) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): optimizer = tf.train.MomentumOptimizer( learning_rate=learning_rate, momentum=CFG.TRAIN.MOMENTUM).minimize( loss=train_total_loss, var_list=tf.trainable_variables(), global_step=global_step) # Set tf model save path model_save_dir = 'model/tusimple_lanenet_{:s}'.format(net_flag) 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 = 'tusimple_lanenet_{:s}_{:s}.ckpt'.format( net_flag, str(train_start_time)) model_save_path = ops.join(model_save_dir, model_name) saver = tf.train.Saver() # Set tf summary save path tboard_save_path = 'tboard/tusimple_lanenet_{:s}'.format(net_flag) if not ops.exists(tboard_save_path): os.makedirs(tboard_save_path) # 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(): 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) if net_flag == 'vgg' and weights_path is None: load_pretrained_weights(tf.trainable_variables(), './data/vgg16.npy', sess) train_cost_time_mean = [] for epoch in range(train_epochs): # training part t_start = time.time() _, train_c, train_accuracy_figure, train_fn_figure, train_fp_figure, lr, train_summary, train_binary_loss, \ train_instance_loss, train_embeddings, train_binary_seg_imgs, train_gt_imgs, \ train_binary_gt_labels, train_instance_gt_labels = \ sess.run([optimizer, train_total_loss, train_accuracy, train_fn, train_fp, learning_rate, train_merge_summary_op, train_binary_seg_loss, train_disc_loss, train_pix_embedding, train_prediction, train_images, train_binary_labels, train_instance_labels]) if math.isnan(train_c) or math.isnan( train_binary_loss) or math.isnan(train_instance_loss): log.error('cost is: {:.5f}'.format(train_c)) log.error('binary cost is: {:.5f}'.format(train_binary_loss)) log.error( 'instance cost is: {:.5f}'.format(train_instance_loss)) return if epoch % 100 == 0: record_training_intermediate_result( gt_images=train_gt_imgs, gt_binary_labels=train_binary_gt_labels, gt_instance_labels=train_instance_gt_labels, binary_seg_images=train_binary_seg_imgs, pix_embeddings=train_embeddings) summary_writer.add_summary(summary=train_summary, global_step=epoch) if epoch % CFG.TRAIN.DISPLAY_STEP == 0: log.info( 'Epoch: {:d} total_loss= {:6f} binary_seg_loss= {:6f} ' 'instance_seg_loss= {:6f} accuracy= {:6f} fp= {:6f} fn= {:6f}' ' lr= {:6f} mean_cost_time= {:5f}s '.format( epoch + 1, train_c, train_binary_loss, train_instance_loss, train_accuracy_figure, train_fp_figure, train_fn_figure, lr, np.mean(train_cost_time_mean))) del train_cost_time_mean[:] # validation part val_c, val_accuracy_figure, val_fn_figure, val_fp_figure, val_summary, val_binary_loss, \ val_instance_loss, val_embeddings, val_binary_seg_imgs, val_gt_imgs, \ val_binary_gt_labels, val_instance_gt_labels = \ sess.run([val_total_loss, val_accuracy, val_fn, val_fp, val_merge_summary_op, val_binary_seg_loss, val_disc_loss, val_pix_embedding, val_prediction, val_images, val_binary_labels, val_instance_labels]) if math.isnan(val_c) or math.isnan(val_binary_loss) or math.isnan( val_instance_loss): log.error('cost is: {:.5f}'.format(val_c)) log.error('binary cost is: {:.5f}'.format(val_binary_loss)) log.error('instance cost is: {:.5f}'.format(val_instance_loss)) return if epoch % 100 == 0: record_training_intermediate_result( gt_images=val_gt_imgs, gt_binary_labels=val_binary_gt_labels, gt_instance_labels=val_instance_gt_labels, binary_seg_images=val_binary_seg_imgs, pix_embeddings=val_embeddings, flag='val') cost_time = time.time() - t_start train_cost_time_mean.append(cost_time) summary_writer.add_summary(summary=val_summary, global_step=epoch) if epoch % CFG.TRAIN.VAL_DISPLAY_STEP == 0: log.info( 'Epoch_Val: {:d} total_loss= {:6f} binary_seg_loss= {:6f} ' 'instance_seg_loss= {:6f} accuracy= {:6f} fp= {:6f} fn= {:6f}' ' mean_cost_time= {:5f}s '.format( epoch + 1, val_c, val_binary_loss, val_instance_loss, val_accuracy_figure, val_fp_figure, val_fn_figure, np.mean(train_cost_time_mean))) del train_cost_time_mean[:] if epoch % 2000 == 0: saver.save(sess=sess, save_path=model_save_path, global_step=global_step) return
def train_lanenet_multi_gpu(dataset_dir, weights_path=None, net_flag='vgg'): """ train lanenet with multi gpu :param dataset_dir: :param weights_path: :param net_flag: :return: """ # set lanenet dataset train_dataset = lanenet_data_feed_pipline.LaneNetDataFeeder( dataset_dir=dataset_dir, flags='train') val_dataset = lanenet_data_feed_pipline.LaneNetDataFeeder( dataset_dir=dataset_dir, flags='val') # set lanenet train_net = lanenet.LaneNet(net_flag=net_flag, phase='train', reuse=False) val_net = lanenet.LaneNet(net_flag=net_flag, phase='val', reuse=True) # set compute graph node train_images, train_binary_labels, train_instance_labels = train_dataset.inputs( CFG.TRAIN.BATCH_SIZE, 1) val_images, val_binary_labels, val_instance_labels = val_dataset.inputs( CFG.TRAIN.VAL_BATCH_SIZE, 1) # set average container tower_grads = [] train_tower_loss = [] val_tower_loss = [] batchnorm_updates = None train_summary_op_updates = None # set lr global_step = tf.Variable(0, trainable=False) learning_rate = tf.train.polynomial_decay( learning_rate=CFG.TRAIN.LEARNING_RATE, global_step=global_step, decay_steps=CFG.TRAIN.EPOCHS, power=0.9) # set optimizer optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate, momentum=CFG.TRAIN.MOMENTUM) # set distributed train op with tf.variable_scope(tf.get_variable_scope()): for i in range(CFG.TRAIN.GPU_NUM): with tf.device('/gpu:{:d}'.format(i)): with tf.name_scope('tower_{:d}'.format(i)) as _: train_loss, grads = compute_net_gradients( train_images, train_binary_labels, train_instance_labels, train_net, optimizer) # Only use the mean and var in the first gpu tower to update the parameter if i == 0: batchnorm_updates = tf.get_collection( tf.GraphKeys.UPDATE_OPS) train_summary_op_updates = tf.get_collection( tf.GraphKeys.SUMMARIES) tower_grads.append(grads) train_tower_loss.append(train_loss) with tf.name_scope('validation_{:d}'.format(i)) as _: val_loss, _ = compute_net_gradients( val_images, val_binary_labels, val_instance_labels, val_net, optimizer) val_tower_loss.append(val_loss) grads = average_gradients(tower_grads) avg_train_loss = tf.reduce_mean(train_tower_loss) avg_val_loss = tf.reduce_mean(val_tower_loss) # Track the moving averages of all trainable variables variable_averages = tf.train.ExponentialMovingAverage( CFG.TRAIN.MOVING_AVERAGE_DECAY, num_updates=global_step) variables_to_average = tf.trainable_variables( ) + tf.moving_average_variables() variables_averages_op = variable_averages.apply(variables_to_average) # Group all the op needed for training batchnorm_updates_op = tf.group(*batchnorm_updates) apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step) train_op = tf.group(apply_gradient_op, variables_averages_op, batchnorm_updates_op) # Set tf summary save path tboard_save_path = 'tboard/tusimple_lanenet_multi_gpu_{:s}'.format( net_flag) if not os.path.exists(tboard_save_path): os.makedirs(tboard_save_path) summary_writer = tf.summary.FileWriter(tboard_save_path) avg_train_loss_scalar = tf.summary.scalar(name='average_train_loss', tensor=avg_train_loss) avg_val_loss_scalar = tf.summary.scalar(name='average_val_loss', tensor=avg_val_loss) learning_rate_scalar = tf.summary.scalar(name='learning_rate_scalar', tensor=learning_rate) train_merge_summary_op = tf.summary.merge( [avg_train_loss_scalar, learning_rate_scalar] + train_summary_op_updates) val_merge_summary_op = tf.summary.merge([avg_val_loss_scalar]) # set tensorflow saver saver = tf.train.Saver() model_save_dir = 'model/tusimple_lanenet_multi_gpu_{:s}'.format(net_flag) if not os.path.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 = 'tusimple_lanenet_{:s}_{:s}.ckpt'.format( net_flag, str(train_start_time)) model_save_path = ops.join(model_save_dir, model_name) # set sess config sess_config = tf.ConfigProto(device_count={'GPU': CFG.TRAIN.GPU_NUM}, 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' # Set the training parameters train_epochs = CFG.TRAIN.EPOCHS log.info('Global configuration is as follows:') log.info(CFG) sess = tf.Session(config=sess_config) summary_writer.add_graph(sess.graph) 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) train_cost_time_mean = [] val_cost_time_mean = [] for epoch in range(train_epochs): # training part t_start = time.time() _, train_loss_value, train_summary, lr = \ sess.run( fetches=[train_op, avg_train_loss, train_merge_summary_op, learning_rate] ) if math.isnan(train_loss_value): log.error('Train loss is nan') return 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 t_start_val = time.time() val_loss_value, val_summary = \ sess.run(fetches=[avg_val_loss, val_merge_summary_op]) 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_Train: {:d} total_loss= {:6f} ' 'lr= {:6f} mean_cost_time= {:5f}s '.format( epoch + 1, train_loss_value, lr, np.mean(train_cost_time_mean))) del train_cost_time_mean[:] if epoch % CFG.TRAIN.VAL_DISPLAY_STEP == 0: log.info('Epoch_Val: {:d} total_loss= {:6f}' ' mean_cost_time= {:5f}s '.format( epoch + 1, val_loss_value, np.mean(val_cost_time_mean))) del val_cost_time_mean[:] if epoch % 2000 == 0: saver.save(sess=sess, save_path=model_save_path, global_step=epoch) return
def train_lanenet(weights_path=None, net_flag='vgg', version_flag='', scratch=False): """ :param weights_path: :param net_flag: choose which base network to use :param version_flag: exp flag :return: """ # ========================== placeholder ========================= # with tf.name_scope('train_input'): train_input_tensor = tf.placeholder(dtype=tf.float32, name='input_image', shape=[None, None, None, 3]) train_binary_label_tensor = tf.placeholder(dtype=tf.float32, name='binary_input_label', shape=[None, None, None, 1]) train_instance_label_tensor = tf.placeholder( dtype=tf.float32, name='instance_input_label', shape=[None, None, None, 1]) with tf.name_scope('val_input'): val_input_tensor = tf.placeholder(dtype=tf.float32, name='input_image', shape=[None, None, None, 3]) val_binary_label_tensor = tf.placeholder(dtype=tf.float32, name='binary_input_label', shape=[None, None, None, 1]) val_instance_label_tensor = tf.placeholder(dtype=tf.float32, name='instance_input_label', shape=[None, None, None, 1]) # ================================================================ # # Define Network # # ================================================================ # train_net = lanenet.LaneNet(net_flag=net_flag, phase='train', reuse=tf.AUTO_REUSE) val_net = lanenet.LaneNet(net_flag=net_flag, phase='val', reuse=True) # ---------------------------------------------------------------- # # ================================================================ # # Train Input & Output # # ================================================================ # trainset = DataSet('train') # trainset = MergeDataSet('train_lane') train_compute_ret = train_net.compute_loss( input_tensor=train_input_tensor, binary_label=train_binary_label_tensor, instance_label=train_instance_label_tensor, name='lanenet_model') train_total_loss = train_compute_ret['total_loss'] train_binary_seg_loss = train_compute_ret['binary_seg_loss'] # 语义分割 loss train_disc_loss = train_compute_ret[ 'discriminative_loss'] # embedding loss train_pix_embedding = train_compute_ret[ 'instance_seg_logits'] # embedding feature, HxWxN train_l2_reg_loss = train_compute_ret['l2_reg_loss'] train_prediction_logits = train_compute_ret[ 'binary_seg_logits'] # 语义分割结果,HxWx2 train_prediction_score = tf.nn.softmax(logits=train_prediction_logits) train_prediction = tf.argmax(train_prediction_score, axis=-1) # 语义分割二值图 train_accuracy = evaluate_model_utils.calculate_model_precision( train_compute_ret['binary_seg_logits'], train_binary_label_tensor) train_fp = evaluate_model_utils.calculate_model_fp( train_compute_ret['binary_seg_logits'], train_binary_label_tensor) train_fn = evaluate_model_utils.calculate_model_fn( train_compute_ret['binary_seg_logits'], train_binary_label_tensor) train_binary_seg_ret_for_summary = evaluate_model_utils.get_image_summary( img=train_prediction) # (I - min) * 255 / (max -min), 归一化到0-255 train_embedding_ret_for_summary = evaluate_model_utils.get_image_summary( img=train_pix_embedding) # (I - min) * 255 / (max -min), 归一化到0-255 # ---------------------------------------------------------------- # # ================================================================ # # Define Optimizer # # ================================================================ # # set optimizer global_step = tf.Variable(0, trainable=False, name='global_step') # learning_rate = tf.train.cosine_decay_restarts( # 余弦衰减 # learning_rate=cfg.TRAIN.LEARNING_RATE, # 初始学习率 # global_step=global_step, # 当前迭代次数 # first_decay_steps=cfg.TRAIN.STEPS/3, # 首次衰减周期 # t_mul=2.0, # 随后每次衰减周期倍数 # m_mul=1.0, # 随后每次初始学习率倍数 # alpha = 0.1, # 最小的学习率=alpha*learning_rate # ) learning_rate = tf.train.polynomial_decay( # 多项式衰减 learning_rate=cfg.TRAIN.LEARNING_RATE, # 初始学习率 global_step=global_step, # 当前迭代次数 decay_steps=cfg.TRAIN.STEPS / 4, # 在迭代到该次数实际,学习率衰减为 learning_rate * dacay_rate end_learning_rate=cfg.TRAIN.LEARNING_RATE / 10, # 最小的学习率 power=0.9, cycle=True) learning_rate_scalar = tf.summary.scalar(name='learning_rate', tensor=learning_rate) update_ops = tf.get_collection( tf.GraphKeys.UPDATE_OPS) # for batch normalization with tf.control_dependencies(update_ops): optimizer = tf.train.MomentumOptimizer( learning_rate=learning_rate, momentum=cfg.TRAIN.MOMENTUM).minimize( loss=train_total_loss, var_list=tf.trainable_variables(), global_step=global_step) # ---------------------------------------------------------------- # # ================================================================ # # Train Summary # # ================================================================ # train_loss_scalar = tf.summary.scalar(name='train_cost', tensor=train_total_loss) train_accuracy_scalar = tf.summary.scalar(name='train_accuracy', tensor=train_accuracy) train_binary_seg_loss_scalar = tf.summary.scalar( name='train_binary_seg_loss', tensor=train_binary_seg_loss) train_instance_seg_loss_scalar = tf.summary.scalar( name='train_instance_seg_loss', tensor=train_disc_loss) train_fn_scalar = tf.summary.scalar(name='train_fn', tensor=train_fn) train_fp_scalar = tf.summary.scalar(name='train_fp', tensor=train_fp) train_binary_seg_ret_img = tf.summary.image( name='train_binary_seg_ret', tensor=train_binary_seg_ret_for_summary) train_embedding_feats_ret_img = tf.summary.image( name='train_embedding_feats_ret', tensor=train_embedding_ret_for_summary) train_merge_summary_op = tf.summary.merge([ train_accuracy_scalar, train_loss_scalar, train_binary_seg_loss_scalar, train_instance_seg_loss_scalar, train_fn_scalar, train_fp_scalar, train_binary_seg_ret_img, train_embedding_feats_ret_img, learning_rate_scalar ]) # ---------------------------------------------------------------- # # ================================================================ # # Val Input & Output # # ================================================================ # valset = DataSet('val', net_flag) # valset = MergeDataSet('test_lane') val_compute_ret = val_net.compute_loss( input_tensor=val_input_tensor, binary_label=val_binary_label_tensor, instance_label=val_instance_label_tensor, name='lanenet_model') val_total_loss = val_compute_ret['total_loss'] val_binary_seg_loss = val_compute_ret['binary_seg_loss'] val_disc_loss = val_compute_ret['discriminative_loss'] val_pix_embedding = val_compute_ret['instance_seg_logits'] val_prediction_logits = val_compute_ret['binary_seg_logits'] val_prediction_score = tf.nn.softmax(logits=val_prediction_logits) val_prediction = tf.argmax(val_prediction_score, axis=-1) val_accuracy = evaluate_model_utils.calculate_model_precision( val_compute_ret['binary_seg_logits'], val_binary_label_tensor) val_fp = evaluate_model_utils.calculate_model_fp( val_compute_ret['binary_seg_logits'], val_binary_label_tensor) val_fn = evaluate_model_utils.calculate_model_fn( val_compute_ret['binary_seg_logits'], val_binary_label_tensor) val_binary_seg_ret_for_summary = evaluate_model_utils.get_image_summary( img=val_prediction) val_embedding_ret_for_summary = evaluate_model_utils.get_image_summary( img=val_pix_embedding) # ---------------------------------------------------------------- # # ================================================================ # # VAL Summary # # ================================================================ # val_loss_scalar = tf.summary.scalar(name='val_cost', tensor=val_total_loss) val_accuracy_scalar = tf.summary.scalar(name='val_accuracy', tensor=val_accuracy) val_binary_seg_loss_scalar = tf.summary.scalar(name='val_binary_seg_loss', tensor=val_binary_seg_loss) val_instance_seg_loss_scalar = tf.summary.scalar( name='val_instance_seg_loss', tensor=val_disc_loss) val_fn_scalar = tf.summary.scalar(name='val_fn', tensor=val_fn) val_fp_scalar = tf.summary.scalar(name='val_fp', tensor=val_fp) val_binary_seg_ret_img = tf.summary.image( name='val_binary_seg_ret', tensor=val_binary_seg_ret_for_summary) val_embedding_feats_ret_img = tf.summary.image( name='val_embedding_feats_ret', tensor=val_embedding_ret_for_summary) val_merge_summary_op = tf.summary.merge([ val_accuracy_scalar, val_loss_scalar, val_binary_seg_loss_scalar, val_instance_seg_loss_scalar, val_fn_scalar, val_fp_scalar, val_binary_seg_ret_img, val_embedding_feats_ret_img ]) # ---------------------------------------------------------------- # # ================================================================ # # Config Saver & Session # # ================================================================ # # Set tf model save path model_save_dir = 'model/tusimple_lanenet_{:s}_{:s}'.format( net_flag, version_flag) os.makedirs(model_save_dir, exist_ok=True) train_start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time())) model_name = 'tusimple_lanenet_{:s}_{:s}.ckpt'.format( net_flag, str(train_start_time)) model_save_path = ops.join(model_save_dir, model_name) # ============================== if scratch: """ 删除 Momentum 的参数, 注意这里保存的 meta 文件也会删了 tensorflow 在 save model 的时候,如果选择了 global_step 选项,会 global_step 值也保存下来, 然后 restore 的时候也就会接着这个 global_step 继续训练下去,因此需要去掉 """ variables = tf.contrib.framework.get_variables_to_restore() variables_to_resotre = [ v for v in variables if 'Momentum' not in v.name.split('/')[-1] ] variables_to_resotre = [ v for v in variables_to_resotre if 'global_step' not in v.name.split('/')[-1] ] # remove global step restore_saver = tf.train.Saver(variables_to_resotre) else: restore_saver = tf.train.Saver() saver = tf.train.Saver(max_to_keep=10) # ============================== # Set tf summary save path tboard_save_path = 'tboard/tusimple_lanenet_{:s}_{:s}'.format( net_flag, version_flag) os.makedirs(tboard_save_path, exist_ok=True) # Set sess configuration # ============================== config GPU 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 import math one_epoch2step = math.ceil(cfg.TRAIN.TRAIN_SIZE / cfg.TRAIN.BATCH_SIZE) # 训练一个 epoch 需要的 batch 数量 total_epoch = math.ceil(cfg.TRAIN.STEPS / one_epoch2step) # 一共需要训练多少 epoch log.info('Global configuration is as follows:') log.info(cfg) max_acc = 0.9 save_num = 0 val_step = 0 # ================================================================ # # Train & Val # # ================================================================ # with sess.as_default(): # ============================== load pretrain model # if weights_path is None: # log.info('Training from scratch') # sess.run(tf.global_variables_initializer()) # elif net_flag == 'vgg' and weights_path is None: # load_pretrained_weights(tf.trainable_variables(), './data/vgg16.npy', sess) # elif scratch: # 从头开始训练,类似 Caffe 的 --weights # sess.run(tf.global_variables_initializer()) # log.info('Restore model from last model checkpoint {:s}, scratch'.format(weights_path)) # try: # restore_saver.restore(sess=sess, save_path=weights_path) # except: # log.info('model maybe is not exist!') # else: # 继续训练,类似 Caffe 的 --snapshot # log.info('Restore model from last model checkpoint {:s}'.format(weights_path)) # try: # restore_saver.restore(sess=sess, save_path=weights_path) # except: # log.info('model maybe is not exist!') sess.run(tf.global_variables_initializer()) # ============================== for epoch in range(total_epoch): # ================================================================ # # Train # # ================================================================ # train_epoch_loss = [] pbar_train = tqdm(trainset) train_t_start = time.time() for gt_imgs, binary_gt_labels, instance_gt_labels in pbar_train: _, global_step_val, train_loss, train_accuracy_figure, train_fn_figure, train_fp_figure, \ lr, train_summary, train_binary_loss, train_instance_loss, \ train_embeddings, train_binary_seg_imgs, train_l2_loss = \ sess.run([optimizer, global_step, train_total_loss, train_accuracy, train_fn, train_fp, learning_rate, train_merge_summary_op, train_binary_seg_loss, train_disc_loss, train_pix_embedding, train_prediction, train_l2_reg_loss], feed_dict={train_input_tensor: gt_imgs, train_binary_label_tensor: binary_gt_labels, train_instance_label_tensor: instance_gt_labels} ) # ============================== 透心凉,心飞扬 if math.isnan(train_loss) or math.isnan( train_binary_loss) or math.isnan(train_instance_loss): log.error('cost is: {:.5f}'.format(train_loss)) log.error( 'binary cost is: {:.5f}'.format(train_binary_loss)) log.error( 'instance cost is: {:.5f}'.format(train_instance_loss)) return # ============================== train_epoch_loss.append(train_loss) summary_writer.add_summary(summary=train_summary, global_step=global_step_val) pbar_train.set_description( ("train loss: %.4f, learn rate: %e") % (train_loss, lr)) train_cost_time = time.time() - train_t_start mean_train_loss = np.mean(train_epoch_loss) log.info( 'MEAN Train: total_loss= {:6f} mean_cost_time= {:5f}s'.format( mean_train_loss, train_cost_time)) # ---------------------------------------------------------------- # # ================================================================ # # Val # # ================================================================ # # 每隔 epoch 次,测试整个验证集 pbar_val = tqdm(valset) val_epoch_loss = [] val_epoch_binary_loss = [] val_epoch_instance_loss = [] val_epoch_accuracy_figure = [] val_epoch_fp_figure = [] val_epoch_fn_figure = [] val_t_start = time.time() for val_images, val_binary_labels, val_instance_labels in pbar_val: # validation part val_step += 1 val_summary, \ val_loss, val_binary_loss, val_instance_loss, \ val_accuracy_figure, val_fn_figure, val_fp_figure = \ sess.run([val_merge_summary_op, val_total_loss, val_binary_seg_loss, val_disc_loss, val_accuracy, val_fn, val_fp], feed_dict={val_input_tensor: val_images, val_binary_label_tensor: val_binary_labels, val_instance_label_tensor: val_instance_labels} ) # ============================== 透心凉,心飞扬 if math.isnan(val_loss) or math.isnan( val_binary_loss) or math.isnan(val_instance_loss): log.error('cost is: {:.5f}'.format(val_loss)) log.error('binary cost is: {:.5f}'.format(val_binary_loss)) log.error( 'instance cost is: {:.5f}'.format(val_instance_loss)) return # ============================== summary_writer.add_summary(summary=val_summary, global_step=val_step) pbar_val.set_description(("val loss: %.4f, accuracy: %.4f") % (val_loss, val_accuracy_figure)) val_epoch_loss.append(val_loss) val_epoch_binary_loss.append(val_binary_loss) val_epoch_instance_loss.append(val_instance_loss) val_epoch_accuracy_figure.append(val_accuracy_figure) val_epoch_fp_figure.append(val_fp_figure) val_epoch_fn_figure.append(val_fn_figure) val_cost_time = time.time() - val_t_start mean_val_loss = np.mean(val_epoch_loss) mean_val_binary_loss = np.mean(val_epoch_binary_loss) mean_val_instance_loss = np.mean(val_epoch_instance_loss) mean_val_accuracy_figure = np.mean(val_epoch_accuracy_figure) mean_val_fp_figure = np.mean(val_epoch_fp_figure) mean_val_fn_figure = np.mean(val_epoch_fn_figure) # ============================== if mean_val_accuracy_figure > max_acc: max_acc = mean_val_accuracy_figure if save_num < 3: # 前三次不算 max_acc = 0.9 log.info( 'MAX_ACC change to {}'.format(mean_val_accuracy_figure)) model_save_path_max = ops.join( model_save_dir, 'tusimple_lanenet_{}.ckpt'.format( mean_val_accuracy_figure)) saver.save(sess=sess, save_path=model_save_path_max, global_step=global_step) save_num += 1 # ============================== log.info( '=> Epoch: {}, MEAN Val: total_loss= {:6f} binary_seg_loss= {:6f} ' 'instance_seg_loss= {:6f} accuracy= {:6f} fp= {:6f} fn= {:6f}' ' mean_cost_time= {:5f}s '.format( epoch, mean_val_loss, mean_val_binary_loss, mean_val_instance_loss, mean_val_accuracy_figure, mean_val_fp_figure, mean_val_fn_figure, val_cost_time)) # ---------------------------------------------------------------- # return
def test_lanenet(image_path, weights_path): """ :param image_path: :param weights_path: :return: """ assert ops.exists(image_path), '{:s} not exist'.format(image_path) LOG.info('Start reading image and preprocessing') t_start = time.time() image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 LOG.info('Image load complete, cost time: {:.5f}s'.format(time.time() - t_start)) tf.compat.v1.disable_eager_execution() input_tensor = tf.compat.v1.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', cfg=CFG) binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='LaneNet') postprocessor = lanenet_postprocess.LaneNetPostProcessor(cfg=CFG) # Set sess configuration sess_config = tf.compat.v1.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' sess = tf.compat.v1.Session(config=sess_config) # define moving average version of the learned variables for eval with tf.compat.v1.variable_scope(name_or_scope='moving_avg'): variable_averages = tf.train.ExponentialMovingAverage( CFG.SOLVER.MOVING_AVE_DECAY) variables_to_restore = variable_averages.variables_to_restore() # define saver saver = tf.compat.v1.train.Saver(variables_to_restore) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) t_start = time.time() loop_times = 500 for i in range(loop_times): binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) t_cost = time.time() - t_start t_cost /= loop_times LOG.info('Single imgae inference cost time: {:.5f}s'.format(t_cost)) postprocess_result = postprocessor.postprocess( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis) mask_image = postprocess_result['mask_image'] for i in range(CFG.MODEL.EMBEDDING_FEATS_DIMS): instance_seg_image[0][:, :, i] = minmax_scale(instance_seg_image[0][:, :, i]) embedding_image = np.array(instance_seg_image[0], np.uint8) plt.figure('mask_image') plt.imshow(mask_image[:, :, (2, 1, 0)]) plt.savefig(fname="results/mask.png") plt.figure('src_image') plt.imshow(image_vis[:, :, (2, 1, 0)]) plt.savefig(fname="results/src.png") plt.figure('instance_image') plt.imshow(embedding_image[:, :, (2, 1, 0)]) plt.savefig(fname="results/instance.png") plt.figure('binary_image') plt.imshow(binary_seg_image[0] * 255, cmap='gray') plt.savefig(fname="results/binary.png") sess.close() return
def custom_test_lanenet(src_dir , weights_path, save_dir): """ :param src_dir: :param weights_path: :param save_dir: :return: """ assert ops.exists(src_dir), '{:s} not exist'.format(src_dir) os.makedirs(save_dir, exist_ok=True) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', cfg=CFG) binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='LaneNet') postprocessor = lanenet_postprocess.LaneNetPostProcessor(cfg=CFG) saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) with sess.as_default(): save_img_lane_dict = {} saver.restore(sess=sess, save_path=weights_path) img_names = os.listdir(src_dir) image_list = [os.path.join(src_dir,i) for i in img_names] avg_time_cost = [] for index, image_path in tqdm.tqdm(enumerate(image_list), total=len(image_list)): image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 t_start = time.time() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]} ) avg_time_cost.append(time.time() - t_start) postprocess_result = postprocessor.postprocess( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis ) lanes = postprocess_result['lane_points_on_img'] if index % 100 == 0: LOG.info('Mean inference time every single image: {:.5f}s'.format(np.mean(avg_time_cost))) avg_time_cost.clear() input_image_dir = src_dir input_image_name = ops.split(image_path)[-1] output_image_dir = save_dir os.makedirs(output_image_dir, exist_ok=True) output_image_path = ops.join(output_image_dir, input_image_name) # if ops.exists(output_image_path): # continue save_img_lane_dict[input_image_name] = lanes cv2.imwrite(output_image_path, postprocess_result['source_image']) output_json_name = os.path.join(save_dir,'img_lane_points.json') with open(output_json_name, 'w') as fp: json.dump(save_img_lane_dict, fp) return
def test_lanenet(image_path, weights_path): """ :param image_path: :param weights_path: :return: """ assert ops.exists(image_path), '{:s} not exist'.format(image_path) log.info('Start reading image and preprocessing') t_start = time.time() image = cv2.imread(image_path, cv2.IMREAD_COLOR) image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 log.info('Image load complete, cost time: {:.5f}s'.format(time.time() - t_start)) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', net_flag='vgg') binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model') postprocessor = lanenet_postprocess.LaneNetPostProcessor() saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.TEST.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) with sess.as_default(): saver.restore(sess=sess, save_path=weights_path) t_start = time.time() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]}) t_cost = time.time() - t_start log.info('Single imgae inference cost time: {:.5f}s'.format(t_cost)) postprocess_result = postprocessor.postprocess( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis) mask_image = postprocess_result['mask_image'] for i in range(CFG.TRAIN.EMBEDDING_FEATS_DIMS): instance_seg_image[0][:, :, i] = minmax_scale(instance_seg_image[0][:, :, i]) embedding_image = np.array(instance_seg_image[0], np.uint8) plt.figure('mask_image') plt.imshow(mask_image[:, :, (2, 1, 0)]) plt.figure('src_image') plt.imshow(image_vis[:, :, (2, 1, 0)]) plt.figure('instance_image') plt.imshow(embedding_image[:, :, (2, 1, 0)]) plt.figure('binary_image') plt.imshow(binary_seg_image[0] * 255, cmap='gray') plt.show() cv2.imwrite('instance_mask_image.png', mask_image) cv2.imwrite('source_image.png', postprocess_result['source_image']) cv2.imwrite('binary_mask_image.png', binary_seg_image[0] * 255) sess.close() return
def custom_test_lanenet_video(src_dir , weights_path, save_dir): """ :param src_dir: :param weights_path: :param save_dir: :return: """ assert ops.exists(src_dir), '{:s} not exist'.format(src_dir) os.makedirs(save_dir, exist_ok=True) input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor') net = lanenet.LaneNet(phase='test', cfg=CFG) binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='LaneNet') postprocessor = lanenet_postprocess.LaneNetPostProcessor(cfg=CFG) saver = tf.train.Saver() # Set sess configuration sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = CFG.GPU.GPU_MEMORY_FRACTION sess_config.gpu_options.allow_growth = CFG.GPU.TF_ALLOW_GROWTH sess_config.gpu_options.allocator_type = 'BFC' sess = tf.Session(config=sess_config) fourcc = cv2.VideoWriter_fourcc(*'MP4V') # out = cv2.VideoWriter('output.mp4', fourcc, 4.0, (640,480)) out = cv2.VideoWriter(os.path.join(save_dir,'output.mp4'),0x7634706d, 4, (1280,720)) with sess.as_default(): save_img_lane_dict = {} saver.restore(sess=sess, save_path=weights_path) cap = cv2.VideoCapture(src_dir) avg_time_cost = [] frame_id = 0 while(True): ret, image = cap.read() frame_id+=1 if not ret or frame_id>50: break image_vis = image image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 127.5 - 1.0 t_start = time.time() binary_seg_image, instance_seg_image = sess.run( [binary_seg_ret, instance_seg_ret], feed_dict={input_tensor: [image]} ) avg_time_cost.append(time.time() - t_start) postprocess_result = postprocessor.postprocess( binary_seg_result=binary_seg_image[0], instance_seg_result=instance_seg_image[0], source_image=image_vis ) lanes = postprocess_result['lane_points_on_img'] if frame_id % 100 == 0: LOG.info('Mean inference time every single image: {:.5f}s'.format(np.mean(avg_time_cost))) avg_time_cost.clear() output_image_dir = save_dir os.makedirs(output_image_dir, exist_ok=True) output_image_path = ops.join(output_image_dir, str(frame_id)+'.jpg') # if ops.exists(output_image_path): # continue save_img_lane_dict[frame_id] = lanes cv2.imwrite(output_image_path, postprocess_result['source_image']) out.write(postprocess_result['source_image']) out.release() cap.release() output_json_name = os.path.join(save_dir,'img_lane_points.json') with open(output_json_name, 'w') as fp: json.dump(save_img_lane_dict, fp) return