win10系统下使用deeplabv3训练自己的数据

版权所有，翻版必究。https://blog.csdn.net/fightingxyz/article/details/105512849

环境：win10，Anaconda3，tensorflow1.15.0，tensorflow-gpu1.14.0，pycharm。

再看这篇博客的时候建议先看上一篇博客，然后再看就简单了！https://blog.csdn.net/fightingxyz/article/details/105488802

第一步：准备数据。使用labelme来进行标注。安装方法pip install labelme

然后对你的数据进行标注，并保存。可以参考文章https://blog.csdn.net/heiheiya/article/details/88342597。该作者说的很详细。

补充：将数据的输入变成一样的大小，后面的测试需要。

第二步：制作类似pasval_cov_2012的数据。

我的数据最终形式：

那么该怎么做呢？

首先将你第一步准备好的数据进行处理。

import argparse
import json
import os
import os.path as osp
import warnings
import copy
import shutil
import numpy as np
import PIL.Image
from skimage import io
import yaml

from labelme import utils

NAME_LABEL_MAP = {
    '_background_': 0,
    '第二类': 1,
}

def main():
    json_file = 'D:\\segement_data\\train_image\\'
    # json_file = 'D:\\segement_data\\json\\'
    out_dir = 'D:\\segement_data\\output'
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    print(out_dir)
    list = os.listdir(json_file)
    for i in range(0, len(list)):
        path = os.path.join(json_file, list[i])
        if (list[i].split(".")[-1]) != "json":
            continue
        filename = list[i][:-5]  # .json
        print(filename)
        # label_name_to_value = {'_background_': 0}
        if os.path.isfile(path):

            data = json.load(open(path))
            img = utils.image.img_b64_to_arr(data['imageData'])
            lbl, lbl_names = utils.shape.labelme_shapes_to_label(img.shape, data['shapes'])  # labelme_shapes_to_label
            # lbl, lbl_names = utils.shape.shapes_to_label(img.shape, data['shapes'])  # labelme_shapes_to_label
            # lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)

            # modify labels according to NAME_LABEL_MAP
            lbl_tmp = copy.copy(lbl)
            for key_name in lbl_names:
                old_lbl_val = lbl_names[key_name]
                new_lbl_val = NAME_LABEL_MAP[key_name]
                lbl_tmp[lbl == old_lbl_val] = new_lbl_val
            lbl_names_tmp = {}
            for key_name in lbl_names:
                lbl_names_tmp[key_name] = NAME_LABEL_MAP[key_name]

            # Assign the new label to lbl and lbl_names dict
            lbl = np.array(lbl_tmp, dtype=np.int8)
            lbl_names = lbl_names_tmp
            captions = ['%d: %s' % (l, name) for l, name in enumerate(lbl_names)]
            # lbl_viz = utils.draw.draw_label(lbl, img, captions)
            utils.lblsave(osp.join(out_dir, '{}.png'.format(filename)), lbl)
            print('Saved to: %s' % out_dir)


if __name__ == '__main__':
    main()

上诉程序的目的是为了将你的json文件中的特征取出来，放在一个单独的文件夹中。train_image中存放的是你的数据极其对应的json文件。output即为输出。将你的数据（不包含json文件）放在JPEGImages中。把output中生成的文件放在SegmentationClass中。

那么ImageSets中应该放什么呢？其实存放一个\Segmentation就可以了。那\Segmentation中放什么呢？看下图。

我用的数据较少，直接输入的有点说明一下，一定不要多余的空格，否则运行的时候找不到你的文件位置。如果出现错了就找找txt是不是有问题。

当数据较多的时候（就是一个简单的数据写入和随机分配而已。打开文件写入名字，在关闭就行了）。此段为网上找的代码。

#coding:utf-8
import os
import random

trainval_percent = 1  #训练验证数据集的百分比
train_percent = 0.9 		#训练集的百分比
filepath = './JPEGImages'
total_img = os.listdir(filepath)
num=len(total_img)  				#列表的长度
list=range(num)
tv=int(num*trainval_percent)  #训练验证集的图片个数
tr=int(tv*train_percent)  	  #训练集的图片个数	# sample(seq, n) 从序列seq中选择n个随机且独立的元素；
trainval= random.sample(list,tv)
train=random.sample(trainval,tr)
#创建文件trainval.txt,test.txt,train.txt,val.txt
ftrain = open('./ImageSets/Segmentation/train.txt', 'w')
fval = open('./ImageSets/Segmentation/val.txt', 'w')
for i  in list:
    name=total_img[i][:-4]+'\n'
    if i in train:
        ftrain.write(name)
    else:
        fval.write(name)
ftrain.close()
fval.close()

哈哈哈，到这里数据基本准备好了。是不是知道该怎么处理了，和之前官网的文件一样。目前的标签数据并不是灰度图，所以还需要处理一下才行。使用remove_gt_colormap.py文件，将数据集转化为单通道的标签数据。

然后你的SegmentationClassRaw就有数据了哦。

第三步：制作tfrecord格式数据。

如果你看了我的上一篇博客你就知道怎么做了。简单说一下。使用build_voc2012_data.py文件将其转化为tfrcord格式的数据。

重点：保证数据的输入格式哦！！！！！正确的输入数据！！！！在build_data数据中。有点图像数据为.tif，.bmp等改一下格式即可。

上代码：

import json
import os
import os.path as osp
import copy
import numpy as np
from labelme import utils
import cv2

def rename():
    json_file = '输入位置'
    out_dir = '输出位置'
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    # print(out_dir)
    list = os.listdir(json_file)
    # print(list)
    for filename in list:
        newname = filename
        newname = newname.split(".")
        print(newname)
        if newname[-1] == "tif":   //改这个地方就行了
            newname[-1] = "jpg"
            newname = str.join(".", newname)  # 这里要用str.join
            print(newname)
            filename = json_file + filename
            newname = json_file + newname
            os.rename(filename, newname)
            print(newname, "updated successfully")



if __name__ == '__main__':
    rename()

第四步：开始训练你的模型

需要做的事？你的train.py导入的文件中是不是有一个叫data_generator的，打开修改。

//主要是添加你的训练和测试数据的numbers。还有num_classes的类别数（和你自己的数据有关）
_YOUR_DATA = DatasetDescriptor(
   splits_to_sizes={
        'train': 10,  # num of samples in images/training   
        'val': 10,  # num of samples in images/validation
    },
    num_classes=2,
    ignore_label=255,
)
_DATASETS_INFORMATION = {
    'cityscapes': _CITYSCAPES_INFORMATION,
    'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
    'ade20k': _ADE20K_INFORMATION,
    'your_data':_YOUR_DATA,
}

对于train.py文件的设置：

initialize_last_layer=True

last_layers_contain_logits_only=False

train_crop_size=[513,513]

其他设置和上个博客一样。有些参数可以根据你的需求进行修该。我这里没变。https://blog.csdn.net/fightingxyz/article/details/105488802

第五步：运行vis.py，vis_crop_size大小与输入的图像相同（相同大小的用处在这里）

（如果你要运行eval.py也是可以的）。设置和前一篇博客一样。https://blog.csdn.net/fightingxyz/article/details/105488802

在你的生成目录看你的结果即可。

第六步：生成,pb文件，便于下一步调用.pb文件

运行export_model.py文件

第一个需要修改的地方，别忘了你的类别。不用修改，需要调一下这一块

（第一个修改的地方不需要，用原来的即可）

第二个你的存放地址。我是设立一个pb文件夹存放的

第三个你的训练结果，图像大小设置：

其他设置：

结果：

第七步：对单张图片进行查看。

# -*- coding: utf-8 -*-
from matplotlib import gridspec
from matplotlib import pyplot as plt
import numpy as np
from PIL import Image

import tensorflow as tf
config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
sess = tf.Session(config=config)

#这个地方指定输出的模型路径
TEST_PB_PATH    = '你的模型位置，即.pb文件'

#这个地方指定需要测试的图片
TEST_IMAGE_PATH = "需要测试的图片地址"


import cv2

class DeepLabModel(object):
  """Class to load deeplab model and run inference."""
  INPUT_TENSOR_NAME  = 'ImageTensor:0'
  OUTPUT_TENSOR_NAME = 'SemanticPredictions:0'
  INPUT_SIZE         = 513
  FROZEN_GRAPH_NAME  = 'frozen_inference_graph'

  def __init__(self):
    """Creates and loads pretrained deeplab model."""
    self.graph = tf.Graph()

    graph_def = None

    with open(TEST_PB_PATH, 'rb') as fhandle:
        graph_def = tf.GraphDef.FromString(fhandle.read())


    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):
    width, height = image.size
    resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height)
    target_size = (int(resize_ratio * width), int(resize_ratio * height))
    print("target_size", target_size)
    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]

    #显示图片！！！！！使用opencv进行图片的显示！！！
    seg_map1 = seg_map.astype(np.uint8)
    seg_map1[seg_map1 > 0] = 255
    cv2.imshow(" ", seg_map1)
    cv2.waitKey()
    return resized_image, seg_map


def create_pascal_label_colormap():
  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) << shift
    ind >>= 3
  return colormap


def label_to_color_image(label):
  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.imshow(seg_image, alpha=0.5)
  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', '你的标签名字（可多个）']) ！！！！！！

FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)


MODEL = DeepLabModel()
print('model loaded successfully!')



def run_visualization(path):
    oringnal_im = Image.open(path)
    print('running deeplab on image %s...' % path)
    resized_im, seg_map = MODEL.run(oringnal_im)
    # print("resized_im:{0}, seg_map:{1}".format(resized_im, seg_map))
    vis_segmentation(resized_im, seg_map)

run_visualization(TEST_IMAGE_PATH)

保证数据的输入类别和你的输出一致！！！很多错误都是数据的输入输出问题注意！！！！！

欢迎参考，互相学习，引用请注明出处谢谢！不喜勿喷！