deeplabv3訓練自己的數據

版權所有，翻版必究。https://mp.csdn.net/console/editor/html/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中放什麼呢？看下圖。

我用的數據較少，直接輸入的😓

當數據較多的時候（就是一個簡單的數據寫入和隨機分配而已。打開文件寫入名字，在關閉就行了）。此段爲網上找的代碼。

#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)

保證數據的輸入類別和你的輸出一致！！！很多錯誤都是數據的輸入輸出問題注意！！！！！

歡迎參考，互相學習，引用請註明出處謝謝！不喜勿噴！