用deeplabV3+训练自己的数据集并导出.pb模型进行可视化测试 |
您所在的位置：网站首页 › deep可视化 › 用deeplabV3+训练自己的数据集并导出.pb模型进行可视化测试 |
用deeplabV3+训练自己的数据集并导出.pb模型进行可视化测试
|
用deeplabV3+训练自己的数据集并导出.pb模型进行可视化测试
本文介绍deeplabv3+训练模型并做可视化测试，这里有一篇ckpt模型转tflite并在Android端部署的文章可供参考 一.环境配置和前期准备1.下载tensorflow的deeplabV3+开源项目 2.安装CUDA10.0 3.安装cudnn7.6.5(需和CUDA的版本对应) 4.安装anaconda3 5.用anaconda3新建虚拟环境(python用3.7),教程 6.在虚拟环境中安装tensorflow-gpu pip install tensorflow-gpu==1.15.3 7.在虚拟环境中进到第一步下载的项目的*/research/slim文件夹下删除BUILD文件，运行： python setup.py build python setup.py install并添加环境变量：运行export PYTHONPATH=$PYTHONPATH:pwd:pwd/slim 二.数据集的处理1.数据集文件夹格式： 4.修改build_voc2012_data.py tf.app.flags.DEFINE_string('image_folder', './VOC/JPEGImages', 'Folder containing images.') tf.app.flags.DEFINE_string( 'semantic_segmentation_folder', './VOC/SegmentationClass', 'Folder containing semantic segmentation annotations.') tf.app.flags.DEFINE_string( 'list_folder', './VOC/ImageSets/Segmentation', 'Folder containing lists for training and validation') tf.app.flags.DEFINE_string( 'output_dir', './VOC/tfrecord', 'Path to save converted SSTable of TensorFlow examples.') _NUM_SHARDS = 4 #生成tfrecord文件夹中的文件数 三.模型的训练1.下载预训练模型 2.修改data_generator.py 文件，将自己的数据集加入到里面 _MYDATA = DatasetDescriptor( splits_to_sizes={ 'train': 20281, # num of samples in images/training 'val': 2798, # num of samples in images/validation 'trainval': 23079 }, num_classes=2, ignore_label=255, ) _DATASETS_INFORMATION = { 'cityscapes': _CITYSCAPES_INFORMATION, 'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION, 'ade20k': _ADE20K_INFORMATION, 'mydata':_MYDATA, }3.修改train_utils.py exclude_list = ['global_step', 'logits'] if not initialize_last_layer: exclude_list.extend(last_layers)4.修改train.py flags.DEFINE_string('train_logdir', "./datasets/VOC/trainout", 'Where the checkpoint and logs are stored.') flags.DEFINE_integer('training_number_of_steps', 5000, 'The number of steps used for training') ########### 迭代次数 flags.DEFINE_integer('train_batch_size', 12, 'The number of images in each batch during training.') flags.DEFINE_list('train_crop_size', '513,513', 'Image crop size [height, width] during training.') flags.DEFINE_string('tf_initial_checkpoint', "{PATH}/model.ckpt", 'The initial checkpoint in tensorflow format.')#预训练模型的位置 flags.DEFINE_boolean('initialize_last_layer', False, 'Initialize the last layer.') flags.DEFINE_boolean('last_layers_contain_logits_only', True, 'Only consider logits as last layers or not.') flags.DEFINE_string('dataset', 'mydata', 'Name of the segmentation dataset.') flags.DEFINE_string('dataset_dir', "./datasets/VOC/tfrecord", 'Where the dataset reside.')5.运行train.py 6.训练过程loss值可视化 运行tensorboard --logdir={PATH}/trainout，会出现一个网址，用Google浏览器打开给出的网址. 四.导出.pb模型1.修改export_model.py文件 flags.DEFINE_string('checkpoint_path', "./datasets/VOC/trainout/model.ckpt", 'Checkpoint path') flags.DEFINE_string('export_path', "./datasets/VOC/trainout/pb/frozen_inference_graph.pb", 'Path to output Tensorflow frozen graph.') flags.DEFINE_integer('num_classes', 2, 'Number of classes.') flags.DEFINE_multi_integer('crop_size', [513, 513], 'Crop size [height, width].')运行export_model.py 五.可视化测试1.在frozen_inference_graph.pb所在目录的上一级目录pb所在处，执行打包命令 tar -czf pb.tar.gz pb 2.新建并运行deeplab_demo_test.py文件： import os import tarfile from matplotlib import gridspec import matplotlib.pyplot as plt import numpy as np from PIL import Image import tempfile from six.moves import urllib import tensorflow as tf class DeepLabModel(object): """ 加载 DeepLab 模型； 推断 Inference. """ INPUT_TENSOR_NAME = 'ImageTensor:0' OUTPUT_TENSOR_NAME = 'SemanticPredictions:0' INPUT_SIZE = 513 ########和转换为.pb时的crop_size对应 FROZEN_GRAPH_NAME = 'frozen_inference_graph' def __init__(self, tarball_path): """ Creates and loads pretrained deeplab model. """ self.graph = tf.Graph() graph_def = None # Extract frozen graph from tar archive. tar_file = tarfile.open(tarball_path) for tar_info in tar_file.getmembers(): if self.FROZEN_GRAPH_NAME in os.path.basename(tar_info.name): file_handle = tar_file.extractfile(tar_info) graph_def = tf.GraphDef.FromString(file_handle.read()) break tar_file.close() if graph_def is None: raise RuntimeError('Cannot find inference graph in tar archive.') with self.graph.as_default(): tf.import_graph_def(graph_def, name='') self.sess = tf.Session(graph=self.graph) def run(self, image): """ Runs inference on a single image. Args: image: A PIL.Image object, raw input image. Returns: resized_image: RGB image resized from original input image. seg_map: Segmentation map of `resized_image`. """ width, height = image.size resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height) target_size = (int(resize_ratio * width), int(resize_ratio * height)) resized_image = image.convert('RGB').resize(target_size, Image.ANTIALIAS) batch_seg_map = self.sess.run(self.OUTPUT_TENSOR_NAME, feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]}) seg_map = batch_seg_map[0] return resized_image, seg_map def create_pascal_label_colormap(): """ Creates a label colormap used in PASCAL VOC segmentation benchmark. Returns: A Colormap for visualizing segmentation results. """ colormap = np.zeros((256, 3), dtype=int) ind = np.arange(256, dtype=int) for shift in reversed(range(8)): for channel in range(3): colormap[:, channel] |= ((ind >> channel) & 1) = 3 return colormap def label_to_color_image(label): """ Adds color defined by the dataset colormap to the label. Args: label: A 2D array with integer type, storing the segmentation label. Returns: result: A 2D array with floating type. The element of the array is the color indexed by the corresponding element in the input label to the PASCAL color map. Raises: ValueError: If label is not of rank 2 or its value is larger than color map maximum entry. """ if label.ndim != 2: raise ValueError('Expect 2-D input label') colormap = create_pascal_label_colormap() if np.max(label) >= len(colormap): raise ValueError('label value too large.') return colormap[label] def vis_segmentation(image, seg_map): """Visualizes input image, segmentation map and overlay view.""" plt.figure(figsize=(15, 5)) grid_spec = gridspec.GridSpec(1, 4, width_ratios=[6, 6, 6, 1]) plt.subplot(grid_spec[0]) plt.imshow(image) plt.axis('off') plt.title('input image') plt.subplot(grid_spec[1]) seg_image = label_to_color_image(seg_map).astype(np.uint8) plt.imshow(seg_image) plt.axis('off') plt.title('segmentation map') plt.subplot(grid_spec[2]) plt.imshow(image) plt.imshow(seg_image, alpha=0.7) plt.axis('off') plt.title('segmentation overlay') unique_labels = np.unique(seg_map) ax = plt.subplot(grid_spec[3]) plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation='nearest') ax.yaxis.tick_right() plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels]) plt.xticks([], []) ax.tick_params(width=0.0) plt.grid('off') plt.show() ## LABEL_NAMES = np.asarray( ['background', 'sky', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv']) FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1) FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP) ## Tensorflow 提供的模型下载 MODEL_NAME = 'xception_coco_voctrainval' # ['mobilenetv2_coco_voctrainaug', 'mobilenetv2_coco_voctrainval', 'xception_coco_voctrainaug', 'xception_coco_voctrainval'] _DOWNLOAD_URL_PREFIX = 'http://download.tensorflow.org/models/' _MODEL_URLS = {'mobilenetv2_coco_voctrainaug': 'deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz', 'mobilenetv2_coco_voctrainval': 'deeplabv3_mnv2_pascal_trainval_2018_01_29.tar.gz', 'xception_coco_voctrainaug': 'deeplabv3_pascal_train_aug_2018_01_04.tar.gz', 'xception_coco_voctrainval': 'deeplabv3_pascal_trainval_2018_01_04.tar.gz', } # _TARBALL_NAME = 'deeplab_model.tar.gz' # model_dir = tempfile.mkdtemp() # tf.gfile.MakeDirs(model_dir) # # download_path = os.path.join(model_dir, _TARBALL_NAME) # print('downloading model, this might take a while...') # urllib.request.urlretrieve(_DOWNLOAD_URL_PREFIX + _MODEL_URLS[MODEL_NAME], download_path) # print('download completed! loading DeepLab model...') download_path = '/home/dreamdeck/Downloads/Tensorflow/models-master/research/deeplab/datasets/VOC/trainout/pb.tar.gz' #模型所在位置 #download_path = '/home/dreamdeck/Downloads/Tensorflow/models-master/research/deeplab/datasets/VOC2012/test_model/pb_53506.tar.gz' #download_path = '/home/dreamdeck/Downloads/Tensorflow/models-master/research/deeplab/deeplabv3_cityscapes_train/deeplabv3_mnv2_pascal_train_aug_8bit/pb.tar.gz' #模型所在位置 MODEL = DeepLabModel(download_path) print('model loaded successfully!') ## def run_visualization(imagefile): """ DeepLab 语义分割，并可视化结果. """ orignal_im = Image.open(imagefile) print('running deeplab on image %s...' % imagefile) resized_im, seg_map = MODEL.run(orignal_im) vis_segmentation(resized_im, seg_map) images_dir = '/home/dreamdeck/Downloads/Tensorflow/models-master/research/deeplab/datasets/VOC/JPEGImages' # 测试图片目录所在位置 #images_dir = '/home/dreamdeck/Downloads/Tensorflow/models-master/research/deeplab/datasets/VOC2012/JPEGImages' images = sorted(os.listdir(images_dir)) for imgfile in images: run_visualization(os.path.join(images_dir, imgfile)) print('Done.') |
其中train.txt文件为训练的图片的名称，不带后缀名 val.txt文件为验证的图片的名称，不带后缀名 trainval.txt文件为所有的图片的名称，不带后缀名 JPEGImages存放所有的RGB图片（.jpg格式） SegmentationClass存放所有与JPEGImages对应的标签图片（.png格式，单通道） 2.在VOC下新建tfrecord文件夹 3.修改build_data.py文件的image_format的png为jpg，如下：【本文地址】