物體檢測項目


1、項目介紹

1.1 項目架構設計        

        實現基於tensorflow的物體檢測。項目框架主要分為三部分:數據采集層、深度模型層、用戶層。其中,數據采集層用於對數據進行標記以及轉換成TFRecords格式數據文件。深度模型層的功能是讀取數據采集層輸出的TFRecords數據進行數據的預處理以及對深度模型的訓練,其中深度模型可以使用不同的框架(例如SSD、YOLO等),通過模型工廠進行選擇,本項目中使用SSD物體檢測框架。訓練得到的模型通過tensorflow serving進行部署,提供給后台。用戶層通過前端和后台業務交互得到想要的結果。項目結構如下:

                                                  圖1 物體檢測項目框架

 

使用TensorflowServing進行模型部署有以下幾個好處:

1、可以進行模型的熱更新:只要上傳模型文件到服務器上即可,TensorFlow會自動識別模型並使用,不需要重啟serving 服務。

2、導出模型和使用模型進行解耦合

                                                                                        圖2 TensorflowServing模型部署邏輯

 整個項目開發流程主要分為兩大部分:

1.模型的訓練與測試

    訓練

        數據集處理(將數據轉換成TFRecords格式文件)

        數據讀取

        preprocess(數據預處理)

        網絡構建預測結果

        損失計算並訓練

        模型保存

    測試
        測試數據

        preprocess(數據預處理)

        模型加載

        postprocess(預測結果后期處理)

        預測結果顯示(matplotlib)

2、模型部署與小程序

    模型導出

    TensorFlow Serving部署模型

    Serving客戶端+Flask Web

    小程序前端

 

1.2 項目代碼訓練架構設計

                                                     圖2 項目代碼訓練架構設計

其中:

1.數據集工廠(data factory)

為了使項目能夠讀取不同的數據集

2.預處理工廠(preprocess factory)

為了處理不同模型要求的處理需求

3.模型工廠(model factory)

為了項目訓練數據能夠使用不同的模型

 

1.3 訓練代碼架構設計意義

1.網絡模型和網絡模型之間不交叉,模型和數據之間解耦合,數據集與預處理邏輯之間解耦合;

2.訓練代碼可以調用不同的模型與不同的數據集訓練不同的模型結果。 

 

2. 數據模塊接口
        獲取到的圖片數據集,保存在IMAGE/commodity/JPEGImages文件下。使用圖片標記工具(本項目使用labelimg)將圖片進行標記,輸出XML格式文件,保存在 IMAGE/commodity/Annotatons文件下。這樣的數據集類似PASCAL VOC數據集,數據集的圖片和標記文件分布在不同的文件中,並且圖片和標簽沒有一一對應,后續項目中不方便處理,也不方便項目的解耦合。tensorflow提供了TFRecord個數來統一存儲數據,TFRecord格式是一種將圖像數據和標簽數據存放在一起的二進制文件,在tensorflow中能夠快速處理。因此項目中需要將數據集轉換成TFRecords文件。TFRecord文件中的數據是通過tf.train.Example Protocol Buffer格式存儲的。每個我想ample對應一張圖片,其中包括圖片的各種信息。特點是:

1)體積小,消息大小只需要xml文件的1/10~1/3;

2)解析速度快:解析速度比xml塊20~100倍。

其中,tf.train.Example的定義見本博客的《TFRecord數據處理》一節。

2.1 數據轉換成TFRecord格式文件

2.1.1 轉換步驟:

1)設定每個tfrecord文件中保存多的樣本個數

2)讀取每張圖片內容以及xml文件

3)將每次去讀內容寫入tfrecord文件

2.1.2 數據轉換成TFRecord文件

代碼結構如圖所示:

            圖3 圖片轉換成tfrecord文件

 

        其中,datasets文件夾下的utils存放讀取數據的公用組件;dataset_config.py存放數據讀取的配置;dataset_to_tfrecords.py為主要的數據轉換邏輯。dataset_to_tfrecord.py文件執行dataset_to_tfrecords.py中的run()函數完成數據轉換。具體代碼如下:

2.1.2.1 配置文件dataset_config.py如下:

"""
數據集轉換配置文件
"""

# 指定原始圖片的XML和圖片的文件夾名字
DIRECTORY_ANNOTATIONS = "Annotations/"
DIRECTORY_IMAGES = "JPEGImages/"

# 指定每個TFRecord文件存儲example的數量
SAMPLER_PER_FILES = 200

# 定義字典,保存數據集的類別
# 字典的key是類別,字典的value是一個元組
# 元組的元素不能修改,元組中是類別代表的數字和類別
VOC_LABELS = {
    'none': (0, 'Background'),
    'clothes': (1, 'clothes'),
    'pants': (2, 'pants'),
    'shoes': (3, 'shoes'),
    'watch': (4, 'watch'),
    'phone': (5, 'phone'),
    'audio': (6, 'audio'),
    'computer': (7, 'computer'),
    'books': (8, 'books')
}

 

2.1.2.2 utils文件下的dataset_utils.py文件中,編寫定義tf Example需要的feature轉換公式,代碼如下:

import tensorflow as tf


# 生成整數型的屬性
def int64_feature(value):
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


# 生成浮點型的屬性
def float_feature(value):
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))


# 生成字符串類型的屬性
def bytes_feature(value):
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))

 

2.1.2.3 dataset_to_tfrecords.py文件下主要編寫編寫轉換邏輯,代碼如下:

import tensorflow as tf
import os
import xml.etree.ElementTree as ET
from datasets.dataset_config import DIRECTORY_ANNOTATIONS, DIRECTORY_IMAGES, SAMPLER_PER_FILES, VOC_LABELS
from datasets.utils.dataset_utils import int64_feature, float_feature, bytes_feature


# 獲取輸出的TFRecord文件名字,格式如下:commodity_2018_train_xxx.tfrecord
# xxx代表序號,從000開始
def _get_output_filename(outputdir, dataset_name, fdx):
    """
    獲取輸出的TFRecord文件的名字
    :param outputdir: 輸出路徑
    :param dataset_name: 數據集名字
    :param fdx: 文件id
    :return:
    """
    return "%s/%s_%03d.tfrecord" % (outputdir, dataset_name, fdx)


def _process_image(dataset_dir, image_name):
    """
    處理一張圖片的數據:獲取圖片數據以及xml文件中的內容。根據需要獲取
    :param dataset_dir: 數據集路徑
    :param img_name: 圖片名字
    :return:
    """
    # 圖片路徑 + 圖片名字
    filename = dataset_dir + DIRECTORY_IMAGES + image_name + '.jpg'

    # 讀取圖片數據
    image_data = tf.gfile.FastGFile(filename, 'rb').read()

    # 讀取xml數據,使用ET工具
    # 構造xml文件名字
    filename_xml = dataset_dir + DIRECTORY_ANNOTATIONS + image_name + '.xml'

    # 將文件內容轉換成樹狀結構tree
    tree = ET.parse(filename_xml)

    # 獲取root節點
    root = tree.getroot()

    # 獲取root節點下面的子節點
    # 1、獲取size信息
    size = root.find('size')
    # 把height、width、depth存放在一個shape里面
    shape = [int(size.find('height').text),
             int(size.find('width').text),
             int(size.find('depth').text)]

    # 用於存儲object對應的label的編號
    labels = []
    labels_text = []
    difficults = []
    truncated = []
    bboxes = []

    # 2、獲取 object信息
    for obj in root.findall('object'):
        # 解析每一個object,包含name、difficult、truncated、bndbox[xmin, ymin, xmax, ymax]
        # 取出label和與之對應的數字
        label = obj.find('name').text
        labels.append(int(VOC_LABELS[label][0]))
        labels_text.append(label.encode('ascii'))

        # 取出difficult
        if obj.find('difficult'):
            difficults.append(int(obj.find('difficult').text))
        else:
            # 不存在,默認difficult為0
            difficults.append(0)

        # 取出truncated
        if obj.find('truncated'):
            truncated.append(int(obj.find('truncated').text))
        else:
            # 不存在,默認truncated為0
            truncated.append(0)

        # 取出bndbox
        bbox = obj.find('bndbox')
        bboxes.append([float(bbox.find('ymin').text)/shape[0],
                       float(bbox.find('xmin').text) / shape[1],
                       float(bbox.find('ymax').text) / shape[0],
                       float(bbox.find('xmax').text) / shape[1]])
    return image_data, shape, labels, labels_text, difficults, truncated, bboxes


def _convert_to_example(image_data, shape, labels, labels_text, difficults, truncated, bboxes):
    """
    將圖片數據轉換成example protocol buffer格式
    :param image_data:
    :param shape:
    :param labels:
    :param difficults:
    :param truncated:
    :param bboxes:
    :return:
    """
    # bboxes存儲格式如下:[[a0, b0, c0, d0], [a1, b1, c1, d1]]轉換成
    # ymin[a0, a1], xmin[b0, b1], ymax[c0, c1], xmax[d0, d1]
    ymin = []
    xmin = []
    ymax = []
    xmax = []

    for b in bboxes:
        ymin.append(b[0])
        xmin.append(b[1])
        ymax.append(b[2])
        xmax.append(b[3])

    # 將所有信息封裝成example
    image_format = b'JPEG'
    example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': int64_feature(shape[0]),
        'image/width': int64_feature(shape[1]),
        'image/channels': int64_feature(shape[2]),
        'image/shape': int64_feature(shape),
        'image/object/bbox/ymin': float_feature(ymin),
        'image/object/bbox/xmin': float_feature(xmin),
        'image/object/bbox/ymax': float_feature(ymax),
        'image/object/bbox/xmax': float_feature(xmax),
        'image/object/bbox/label': int64_feature(labels),
        'image/object/bbox/difficult': int64_feature(difficults),
        'image/object/bbox/truncated': int64_feature(truncated),
        'image/object/bbox/label_text': bytes_feature(labels_text),
        'image/format': bytes_feature(image_format),
        'image/encoded': bytes_feature(image_data)}))
    return example


def _add_to_tfrecord(dataset_dir, image_name, tfrecord_writer):
    """
    添加一個圖片文件和xml內容寫入文件中
    :param dataset_dir: 數據集目錄
    :param img_name: 圖片名
    :param tfrecord_writer: 文件寫入實例
    :return:
    """
    # 1、讀取每張圖片內容及其對應的xml文件的內容
    image_data, shape, labels, labels_text, difficults, truncated, bboxes = _process_image(dataset_dir, image_name)

    # 2、將每張圖片的數據封裝成一個example
    example = _convert_to_example(image_data, shape, labels, labels_text, difficults, truncated, bboxes)

    # 3、使用tfrecord_writer將example序列化結果寫入TFRecord文件
    tfrecord_writer.write(example.SerializeToString())
    return None


def run(dataset_dir, output_dir, dataset_name="data"):
    """
    運行轉換代碼邏輯:存入tfrecord文件,每個文件固定N個樣本
    :param dataset_dir: 數據集目錄
    :param output_dir: TFRecord存儲目錄
    :param dataset_name: 數據集名字,指定名字以及train_or_test
    :return:
    """
    # 1、判斷數據集目錄是否存在,不存在則創建一個目錄
    if not tf.gfile.Exists(dataset_dir):
        tf.gfile.MakeDirs(dataset_dir)
    # 2、讀取某個文件夾下的所有文件名字列表
    path = os.path.join(dataset_dir, DIRECTORY_ANNOTATIONS)

    # 讀取所有文件,返回所有文件名字列表。但是會打亂順序,需要使用sorted函數進行排序
    filenames = sorted(os.listdir(path))

    # 3、循環遍歷列表,每N張圖片和XML信息存儲到一個tfrecord文件中
    i = 0
    fdx = 0
    while i < len(filenames):
        # 1、創建TFRecord文件
        tf_filename = _get_output_filename(output_dir, dataset_name, fdx)

        # 每N個文件存儲一次
        # 新建tfrecord的存儲器
        with tf.python_io.TFRecordWriter(tf_filename) as tfrecord_writer:
            j = 0
            while i < len(filenames) and j < SAMPLER_PER_FILES:
                print("轉換圖片進度%d/%d" % (i+1, len(filenames)))

                # 取出圖片以及xml的名字
                single_filename = filenames[i]
                image_name = single_filename[:-4]

                # 讀取圖片和xml內容,存入圖片,每次構造一個圖片文件存儲指定文件
                _add_to_tfrecord(dataset_dir, image_name, tfrecord_writer)

                i += 1
                j += 1

            # 每N個數據,文件id增加計數
            fdx += 1
    print("數據集 %s 轉換成功" % dataset_name)

 

2.1.2.4 dataset_to_tfrecords.py文件代碼

from datasets import dataset_to_tfrecords

if __name__ == '__main__':
    dataset_to_tfrecords.run('./IMAGE/commodity/', './IMAGE/tfrecords/commodity_tfrecords/', 'commodity_2018_train')

 

為了實現數據格式的轉換,需要在圖3的IMAGE文件夾下分別放置如下目錄:

commodity/Annotations/

commodity/JPEGImages/

tfrecords/commodity_tfrecords/

其中,commodity/Annotations/路徑下存放標記過的xml格式文件;commodity/JPEGImages/路徑下存放於xml格式對應的圖片數據;tfrecords/commodity_tfrecords/路徑用於存放轉換好的tfrecord格式數據。

 

2.2 TFRecord格式文件讀取

TFRecord文件讀取有兩種方法:

1)使用tensorflow進行實現

2)使用tensorflow.slim庫進行實現

本項目使用tensorflow.slim進行實現,具體步驟如下:

1、定義解碼器decoder

decoder = tf.slim.tfexample_decoder.TFExampleDecoder()

其中,定義解碼器時,需要制定兩個參數:keys_to_features,和items_to_handlers兩個字典參數。key_to_features這個字典需要和TFrecord文件中定義的字典項匹配。items_to_handlers中的關鍵字可以是任意值,但是它的handler的初始化參數必須要來自於keys_to_features中的關鍵字。

2、定義dataset

dataset= tf.slim.dataset.Dataset()

其中,定義dataset時需要將datasetsource、reader、decoder、num_samples等參數

3、定義provider

provider = slim.dataset_data_provider.DatasetDataProvider

其中,需要的參數為:dataset, num_readers, reader_kwargs, shuffle, num_epochs,common_queue_capacity,common_queue_min, record_key=',seed, scope等。

4、調用provider的get方法

獲取items_to_handlers中定義的關鍵字

5、利用分好的batch建立一個prefetch_queue

6、prefetch_queue中有一個dequeue的op,每執行一次dequeue則返回一個batch的數據。

具體代碼如下(這里先只介紹到通過provider的get函數獲取數據,后面步驟5和步驟6的隊列處理先不介紹,在實際項目代碼中會使用到):

import os
import tensorflow as tf


slim = tf.contrib.slim


def get_dataset(dataset_dir):
    """
    獲取commodity2018數據集
    :param dataset_dir: 數據集目錄
    :return: Dataset
    """
    # 1.准備 tf.slim.dataset.Dataset()的參數
    # 1.1第一個參數:dataset
    file_pattern = os.path.join(dataset_dir, "commodity_2018_train_*.tfrecord")

    # 1.2第二個參數:reader
    reader = tf.TFRecordReader

    # 1.3第三個參數:decoder
    # 創建decoder需要兩個參數:keys_to_features和items_to_handlers
    # 1.3.1 定義keys_to_features,反序列化的格式
    keys_to_features = {
        'image/encoded': tf.FixedLenFeature((), tf.string, default_value=''),
        'image/format': tf.FixedLenFeature((), tf.string, default_value='jpeg'),
        'image/height': tf.FixedLenFeature([1], tf.int64),
        'image/width': tf.FixedLenFeature([1], tf.int64),
        'image/channels': tf.FixedLenFeature([1], tf.int64),
        'image/shape': tf.FixedLenFeature([3], tf.int64),
        'image/object/bbox/xmin': tf.VarLenFeature(dtype=tf.float32),
        'image/object/bbox/ymin': tf.VarLenFeature(dtype=tf.float32),
        'image/object/bbox/xmax': tf.VarLenFeature(dtype=tf.float32),
        'image/object/bbox/ymax': tf.VarLenFeature(dtype=tf.float32),
        'image/object/bbox/label': tf.VarLenFeature(dtype=tf.int64),
        'image/object/bbox/difficult': tf.VarLenFeature(dtype=tf.int64),
        'image/object/bbox/truncated': tf.VarLenFeature(dtype=tf.int64),
    }

    # 1.3.2 items_to_handlers,反序列化成高級的格式
    items_to_handlers = {
        'image': slim.tfexample_decoder.Image('image/encoded', 'image/format'),
        'shape': slim.tfexample_decoder.Tensor('image/shape'),
        'object/bbox': slim.tfexample_decoder.BoundingBox(
            ['ymin', 'xmin', 'ymax', 'xmax'], 'image/object/bbox/'),
        'object/label': slim.tfexample_decoder.Tensor('image/object/bbox/label'),
        'object/difficult': slim.tfexample_decoder.Tensor('image/object/bbox/difficult'),
        'object/truncated': slim.tfexample_decoder.Tensor('image/object/bbox/truncated'),
    }

    # 1.3.3構造decoder
    decoder = slim.tfexample_decoder.TFExampleDecoder(keys_to_features, items_to_handlers)

    # 2.tf.slim.dataset.Dataset()並返回
    return slim.dataset.Dataset(data_sources=file_pattern,
                                reader=reader,
                                decoder=decoder,
                                num_samples=88,
                                items_to_descriptions={
                                    'image': 'A color image of varying height and width.',
                                    'shape': 'Shape of the image',
                                    'object/bbox': 'A list of bounding boxes, one per each object.',
                                    'object/label': 'A list of labels, one per each object.'
                                },  # 數據集返回的格式描述字典
                                num_classes=8)

 

from datasets.dataset_init import commodity_2018
import tensorflow as tf

slim = tf.contrib.slim

if __name__ == '__main__':
    # 獲取dataset
    dataset = commodity_2018.get_dataset("./IMAGE/tfrecords/commodity_tfrecords/")

    # 通過provider取出數據
    provider = slim.dataset_data_provider.DatasetDataProvider(dataset=dataset,
                                                              num_readers=3)

    # 通過get方法獲取指定名稱的數據(名稱在准備規范數據dataset時高級格式的名稱,即items_to_handlers中定義的名稱)
    [image, shape, bbox, label, difficult, truncated] = provider.get(
        ['image', 'shape', 'object/bbox', 'object/label', 'object/difficult', 'object/truncated'])

    print(image, shape, bbox, label, difficult, truncated)

最后得到如下輸出結果:

 

                                                                                                  圖4 輸出tfrecord文件

 

2.3 數據模塊接口——數據工廠的實現

功能需求:

1)原始數據集(圖片+XML)轉換成TFRecords文件格式

2)讀取TFRecords數據

數據模塊設計的目錄如下:

                     圖5 數據模塊接口 

其中:

dataset_factory:數據模塊工廠,找到不同的數據集讀取邏輯;

dataset_init:保存不同數據集的TFRecords格式讀取功能;

utils:數據模塊的共用組件

dataset_config‘:數據模塊的一些數據集配置文件

dataset_to_tfrecords:原始數據集格式轉換邏輯

 

2.3.1 格式轉換

上一節以及介紹了將數據集轉換成TFRecord格式文件,這里就不再贅述。

2.3.2 讀取TFRecord文件數據

2.3.2.1 讀取代碼框架設計

數據模塊需要實現對不同數據集類型進行讀取操作,因此可以定義一個基類,同時不同數據集繼承這個基類。類的設計如下:

                                             圖6 數據讀取基類設計

 

2.3.2.2 數據讀取代碼

1.在dataset_utils.py中新建一個基類,該文件下的代碼如下:

import tensorflow as tf


# 定義數據集TFRecord文件讀取基類
class TFRecordsReaderBase(object):
    """
    數據集讀取基類
    """
    def __init__(self, param):
        # param是給不同數據集使用的屬性配置
        self.param = param

    def get_dataset(self, train_or_test, dataset_dir):
        """
        獲取數據
        :param train_or_test: 訓練還是測試
        :param dataset_dir: 數據集目錄
        :return:
        """
        return None

 

2. 因為在讀取TFRecord數據時,不同的數據集,都會有自己特有的參數(比如:文件名、樣本數、類別數等)。因此在dataset_config.py文件中定義不同數據集的參數,作為繼承類的參數。這里使用命名字典:

"""
數據集讀取
"""
from collections import namedtuple

# 創建命名字典,用於存放讀取數據類中的param參數
DataSetParams = namedtuple("DataSetParamters", ['FILE_PATTERN',
                                                'NUM_CLASSES',
                                                'SPLITS_TO_SIZES',
                                                'ITEMS_TO_DESCRIPTIONS'
                                                ])

# 定義commodity_2018屬性配置
Cmd2018 = DataSetParams(
    FILE_PATTERN='commodity_2018_%s_*.tfrecord',
    NUM_CLASSES=8,
    SPLITS_TO_SIZES={
        'train': 88,
        'test': 0
    },
    ITEMS_TO_DESCRIPTIONS={
        'image': '圖片數據',
        'shape': '圖片形狀',
        'object/bbox': '若干物體對象的bbox框組成的列表',
        'object/label': '若干物體對應的label編號'
    }
)

 

3. 繼承基類來定義派生類用於處理不同數據集

繼承的基類存放在dataset/dataset_init/目錄下。對於不同數據集,定義不同的文件繼承基類,本項目值處理commodity數據集,因此僅創建commodity_2018.py繼承基類,代碼如下:

import os
import tensorflow as tf
from datasets.utils import dataset_utils

slim = tf.contrib.slim


class CommodityTFRecords(dataset_utils.TFRecordsReaderBase):
    """
    商品數據集讀取類
    """
    def __init__(self, param):
        self.param = param

    def get_dataset(self, train_or_test, dataset_dir):
        """
        獲取commodity2018數據集
        :param train_or_test: train or test
        :param dataset_dir: 數據集目錄
        :return:
        """
        # 參數檢查,異常拋出
        if train_or_test not in ['train', 'test']:
            raise ValueError("訓練/測試的名字 %s 錯誤" % train_or_test)

        if not tf.gfile.Exists(dataset_dir):
            raise ValueError("數據集目錄 %s 不存在" % dataset_dir)

        # 1.准備 tf.slim.dataset.Dataset()的參數
        # 1.1第一個參數:dataset
        file_pattern = os.path.join(dataset_dir, self.param.FILE_PATTERN % train_or_test)

        # 1.2第二個參數:reader
        reader = tf.TFRecordReader

        # 1.3第三個參數:decoder
        # 創建decoder需要兩個參數:keys_to_features和items_to_handlers
        # 1.3.1 定義keys_to_features,反序列化的格式
        keys_to_features = {
            'image/encoded': tf.FixedLenFeature((), tf.string, default_value=''),
            'image/format': tf.FixedLenFeature((), tf.string, default_value='jpeg'),
            'image/height': tf.FixedLenFeature([1], tf.int64),
            'image/width': tf.FixedLenFeature([1], tf.int64),
            'image/channels': tf.FixedLenFeature([1], tf.int64),
            'image/shape': tf.FixedLenFeature([3], tf.int64),
            'image/object/bbox/xmin': tf.VarLenFeature(dtype=tf.float32),
            'image/object/bbox/ymin': tf.VarLenFeature(dtype=tf.float32),
            'image/object/bbox/xmax': tf.VarLenFeature(dtype=tf.float32),
            'image/object/bbox/ymax': tf.VarLenFeature(dtype=tf.float32),
            'image/object/bbox/label': tf.VarLenFeature(dtype=tf.int64),
            'image/object/bbox/difficult': tf.VarLenFeature(dtype=tf.int64),
            'image/object/bbox/truncated': tf.VarLenFeature(dtype=tf.int64),
        }

        # 1.3.2 items_to_handlers,反序列化成高級的格式
        items_to_handlers = {
            'image': slim.tfexample_decoder.Image('image/encoded', 'image/format'),
            'shape': slim.tfexample_decoder.Tensor('image/shape'),
            'object/bbox': slim.tfexample_decoder.BoundingBox(
                ['ymin', 'xmin', 'ymax', 'xmax'], 'image/object/bbox/'),
            'object/label': slim.tfexample_decoder.Tensor('image/object/bbox/label'),
            'object/difficult': slim.tfexample_decoder.Tensor('image/object/bbox/difficult'),
            'object/truncated': slim.tfexample_decoder.Tensor('image/object/bbox/truncated'),
        }

        # 1.3.3構造decoder
        decoder = slim.tfexample_decoder.TFExampleDecoder(keys_to_features, items_to_handlers)

        # 2.tf.slim.dataset.Dataset()並返回
        return slim.dataset.Dataset(data_sources=file_pattern,
                                    reader=reader,
                                    decoder=decoder,
                                    num_samples=self.param.SPLITS_TO_SIZES[train_or_test],
                                    items_to_descriptions=self.param.ITEMS_TO_DESCRIPTIONS,  # 數據集返回的格式描述字典
                                    num_classes=self.param.NUM_CLASSES)

 

2.3.3 定義數據工廠

在datasets根目錄下創建dataset_factory.py文件,定義數據工廠獲取數據,代碼如下:

from datasets.dataset_init import commodity_2018
from datasets.dataset_config import Cmd2018

# 定義dataset種類的字典,目前只是有commodity數據集,后續可以添加
datasets_maps = {
    'commodity_2018': commodity_2018.CommodityTFRecords
}

# 定義參數種類的字典,不同數據集,param參數不一樣,目前只是有commodity的參數,后續可以添加
param_map = {
    'commodity_2018': Cmd2018
}


def get_dataset(dataset_name, train_or_test, dataset_dir):
    """
    獲取指定數據名稱的數據文件
    :param dataset_name: 數據集名稱(數據當中要存在
    :param train_or_test: train or test數據集
    :param dataset_dir: 數據集目錄
    :return: Dataset 數據規范
    """
    if dataset_name not in datasets_maps:
        raise ValueError("數據集名稱 %s 不存在" % dataset_name)

    param = param_map[dataset_name]

    return datasets_maps[dataset_name](param).get_dataset(train_or_test, dataset_dir)

最終對外只提供dataset_factory.py文件用於讀取TFRecord文件。

 

3. 模型接口

本項目使用SSD模型。

項目文件結構如下:

          圖7 網絡模型接口文件格式

 

其中的公共組件的源碼都是已知的,本項目使用的ssd網絡模型實現文件ssd_vgg_300.py相關代碼都是現有代碼。對於SSD模型以及其代碼實現,將在另外章節介紹。

3.1 網絡工廠nets_factory實現

類似數據工廠,我們定義模型工廠nets_factory.py文件,代碼如下:

from nets.nets_model import ssd_vgg_300

nets_maps = {
    'ssd_vgg_300': ssd_vgg_300.SSDNet
}


def get_network(network_name):
    """
    獲取不同網絡模型
    :param network_name: 網絡模型名稱
    :return: 網絡
    """
    if network_name not in nets_maps:
        raise ValueError("網絡名稱 %s 不存在" % network_name)
    
    return nets_maps[network_name]

 

4.預處理模塊

目的:

1)在圖像的深度學習中,對輸入數據進行數據增強(Data Augmentation),為了豐富圖像的訓練集,更好地提取圖像特征,泛化模型(防止過擬合)。

通過一系列圖像的操作(比如:剪切、翻轉、偏移、縮放等圖像變換),增加數據集的大小,防止過擬合。

2)還有一個根本目的就是把圖片變成符合大小要求的格式:

RCNN網絡對於輸入圖片沒有要求,但是網絡當中卷積之前需要的大小為227×227;

YOLO算法:輸入圖片大小為448×448;

SSD算法:輸入圖片大小為300×300;

 

4.1 預處理模塊代碼實現

首先,預處理模塊的結構如圖所示:

                 圖8 預處理模塊結構

       

        其中,需要創建一個preprocessing目錄,該目錄下的文件用於數據預處理。該目錄下的processing目錄中的ssd_vgg_preprocessing.py是對於SSD模型的預處理的。如果后續需要增加網絡模型,需要在這個文件夾下增加預處理的文件。utils中是預處理需要用到的公共組件。這些相關代碼都是公開的代碼,這里不做介紹。有了上面的基礎文件,下面就來完成數據預處理工廠代碼的編寫,在preprocessing_factory.py文件中實現:

from preprocessing.processing import ssd_vgg_preprocessing

# 目前只有sdd_vgg_300,后續可以增加
preprocessing_maps = {
    'ssd_vgg_300': ssd_vgg_preprocessing
}


def get_preprocessing(name, is_trainning=True):
    """
    預處理工廠獲取不同的數據增強方法
    :param name: 預處理名稱
    :param is_trainning: 是否是訓練
    :return: 返回預處理的函數,后續再調用函數
    """
    if name not in preprocessing_maps:
        raise ValueError("數據預處理名稱 %s 不存在" % name)

    # 定義一個預處理函數,用於函數返回,后續再調用該預處理函數
    def preprocessing_fn(image, labels, bboxes, out_shape,
                         data_format, **kwargs):
        return preprocessing_maps[name].preprocess_image(image, labels, bboxes, out_shape,
                                                         data_format=data_format,
                                                         is_training=is_trainning, **kwargs)

    return preprocessing_fn

 

5.訓練不同模塊接口參數

對於2、3、4章節,只是分別單獨介紹了數據模塊接口、模型接口以及數據預處理接口。現在需要統一每一個模塊接口提供給訓練的參數,整理成文檔。這樣以后就直接查看文檔即可調用相關模塊。總結如圖9所示:

 

                                                                                                                   圖9 訓練不同模型參數

 

6. 多GPU訓練

終於到了模型訓練這一步了。這里介紹多GPU訓練。

        對於深度學習來說,大量的計算量導致CPU會顯得十分乏力耗時。所以需要GPU來進行提供幫助計算,那么他們的主要任務就是計算得出結果,與CPU之間會進行分工,CPU會做一些基本工作,變量存儲,更新參數,輸入數據變量等等。如圖10所示。在TensorFlow當中會通過標號來區別不同的GPU和CPU,如 ,''/device:CPU:0", "/device:CPU:1","/device:GPU:0","/device:GPU:1","/device:GPU:2",那么這些標號都是程序自動給的編號,指的具體哪塊計算設備。

                                                 圖10 CPU與GPU之間的分工合作

 

6.1 訓練步驟

  • 步驟
    • 數據讀取
    • preprocess(數據預處理)
    • 網絡構建預測結果
    • 損失計算
    • 添加變量到TensorBoard
    • 模型訓練、保存
  • 部署需求:訓練整個模型需要在多GPU、多計算機的環境下進行

那么接下來首先我們要講模型訓練的設備邏輯原理弄清楚,如圖11所示:

                                                                                              圖11 模型訓練的設備邏輯原理

       

        訓練主要是在設備(GPU/CPU)上訓練,但是如果我們利用目前簡單的TensorFlow提供的API去進行指定設備訓練會比較繁瑣。所以在這里需要介紹一個TensorFlow提供的最新的專門用於多GPU,多計算機的設備部署模塊——model_deploy。

 

6.2 model_deploy介紹

model_deploy位於TensorFlow slim模塊的deployment目錄下,可以使得用多個 GPU / CPU在同一台機器或多台機器上執行同步或異步訓練變得更簡單。可以從如下官方地址下載:

https://github.com/tensorflow/models/blob/master/research/slim/deployment/model_deploy.py

首先我們要介紹:

replica:使用多機訓練時,一台機器對應一個replica(復本);

clone:由於tensorflow里多GPU訓練一般都是每個GPU上都有完整的模型,各自進行前向傳播計算,得到的梯度交給CPU平均后統一反向計算,每個GPU上的模型叫做一個clone;

parameter server:多機訓練時,計算梯度平均值並執行反向傳播操作的參數,功能類似於單機多GPU的CPU;

worker server:一般指單機多卡中的GPU,用於訓練。

6.2.1 DeploymentConfig

1. DeploymentConfig為文件中的一個類,主要用於給變量配置選擇的設備。

  • class DeploymentConfig(object):
    • 配置參數
    • num_clones=1:每一個計算設備上的模型克隆數(每台計算機的GPU/CPU總數)
    • clone_on_cpu=False:如果為True,將只在CPU上訓練
    • replica_id=0:指定某個計算機去部署,默認第0台計算機(TensorFlow會給個默認編號)
    • num_replicas=1:多少台可用計算機
    • num_ps_tasks=0:用於參數服務器的計算機數量,0為不適用計算機作為參數服務器
    • worker_job_name='worker':工作服務器名稱
    • ps_job_name='ps':參數服務器名稱
  • config.variables_device()
    • 作為tf.device(func)的參數,返回默認創建變量的設備
    • 一般用於指定全局步數變量的設備,默認運行計算機的"/device:CPU:0"
  • config.inputs_device()
    • 作為tf.device(func)的參數,返回用於構建數據輸入變量所在的設備。
    • 默認運行計算機的"/device:CPU:0"
  • config.optimizer_device()
    • 作為tf.device(func)的參數,返回學習率、優化器所在的設備。
    • 默認運行計算機的"/device:CPU:0"
  • config.clone_scope(self, clone_index):
    • 返回指定編號的設備命名空間
    • 按照這樣編號,clone_0,clone_1...

 

6.2.2 model_deploy定義的相關函數,主要用於為每一個clone創建一個復制的模型(在GPU)

  • model_deploy.create_clones(config, model_fn, args=None, kwargs=None):
    • 作用:每個clone創建一個復制的模型,給GPU進行clone模型
    • config:一個DeploymentConfig的配置對象
    • model_fn:用於回調的函數model_fn,
    • args=None, kwargs=None:回調函數model_fn的參數
    • 返回元組組成的列表,列表個數大小為指定的num_clones數量
      • Clone(outputs, scope, device)
        • outputs:網絡模型的每一層節點
        • scope: 第i個GPU設備的命名空間,config.clone_scope(i)
        • clone_device:第i個GPU設備
  • model_deploy.optimize_clones(clones, optimizer,regularization_losses=None, **kwargs)
    • 作用:計算所有給定的clones的總損失以及每個需要優化的變量的總梯度
    • clones: 元組列表,每個元素Clone(outputs, scope, device)
    • optimizer:選擇的優化器
    • **kwargs:可選參數,優化器優化的變量
    • 返回:
      • total_loss:總損失
      • grads_and_vars:每個需要優化變量的總梯度組成的列表

 

源碼介紹使用:

# Set up DeploymentConfig
config = model_deploy.DeploymentConfig(num_clones=2, clone_on_cpu=True)
# Create the global step on the device storing the variables.
with tf.device(config.variables_device()):
    global_step = slim.create_global_step()
# Define the inputs
with tf.device(config.inputs_device()):
    images, labels = LoadData(...)
    inputs_queue = slim.data.prefetch_queue((images, labels))
# Define the optimizer.
with tf.device(config.optimizer_device()):
    optimizer = tf.train.MomentumOptimizer(FLAGS.learning_rate, FLAGS.momentum)


# Define the model including the loss.
def model_fn(inputs_queue):
    images, labels = inputs_queue.dequeue()
    predictions = CreateNetwork(images)
    slim.losses.log_loss(predictions, labels)


model_dp = model_deploy.deploy(config, model_fn, [inputs_queue],
                               optimizer=optimizer)
# Run training.
slim.learning.train(model_dp.train_op, my_log_dir,
                    summary_op=model_dp.summary_op)

 

6.3 訓練邏輯

1)DeploymentConfig

    需要在訓練之前配置所有的設備信息

    定義全局步數

2)獲取圖片隊列

    在config.inputs_device()指定

3)數據輸入、網絡計算結果、定義損失並復制模型到clones,添加變量到tensorboard

    model_deploy.create_clones

4)定義學習率、優化器

   config.optimizer_device()指定

5)計算所有GPU/CPU設備的平均損失和每個變量的梯度總和、定義訓練OP、summaries OP

    model_deploy.optimize_clones

6)配置訓練的config,進行訓練

slim.learning.train

 

代碼框架如下:

 

圖中,pre_trained文件下存放的是預訓練好的ssd_vgg_300網絡的預訓練模型,fine_tuning是訓練存放模型的路徑。

根目錄下的utils是公共組件,最后訓練的文件是train_ssd_network.py。

訓練代碼如下:

"""
訓練初始化參數

PRE_TRAINED_PATH=./ckpt/pre_trained/ssd_vgg_300.ckpt
TRAIN_MODEL_PDIR=./ckpt/fine_tuning/
DATASET_DIR=./IMAGE/tfrecords/commodity_tfrecords/

每批次訓練樣本數:32或者更小
懲罰項:0.005
學習率:0.001
優化器選擇:adam
模型名稱:ssd_vgg_300
"""

import tensorflow as tf
from datasets import dataset_factory
from preprocessing import preprocessing_factory
from nets import nets_factory
from utils import train_tools
from deployment import model_deploy

slim = tf.contrib.slim

DATA_FORMAT = 'NHWC'

# 命令行參數
# 設備相關的命令行參數
tf.app.flags.DEFINE_integer('num_clones', 1, "可用GPU數量")
tf.app.flags.DEFINE_boolean('clone_on_cpu', False, "是否只在CPU上運行")
tf.app.flags.DEFINE_integer('replica_id', 0, "復本id")

# 數據集相關命令行參數
tf.app.flags.DEFINE_string('dataset_dir', ' ', "訓練數據集目錄")
tf.app.flags.DEFINE_string('dataset_name', 'commodity_2018', "數據集名稱")
tf.app.flags.DEFINE_string('train_or_test', 'train', "訓練還是測試")

# 網絡相關命令行參數
tf.app.flags.DEFINE_string('network_name', 'ssd_vgg_300', "網絡名稱")
tf.app.flags.DEFINE_integer('batch_size', 32, "每批次獲取樣本換數量")
tf.app.flags.DEFINE_float('weight_decay', 0.0001, "網絡誤差懲罰項")

# 訓練相關參數
tf.app.flags.DEFINE_string(
    'optimizer', 'rmsprop', '優化器種類 可選"adadelta", "adagrad", "adam","ftrl", "momentum", "sgd" or "rmsprop".')
tf.app.flags.DEFINE_string(
    'learning_rate_decay_type', 'exponential', '學習率種類 "fixed", "exponential", "polynomial".')
tf.app.flags.DEFINE_float('learning_rate', 0.01, '模型初始學習率')
tf.app.flags.DEFINE_float('end_learning_rate', 0.0001, '模型終止學習率')

tf.app.flags.DEFINE_integer('max_number_of_steps', None, '訓練的最大步數')
tf.app.flags.DEFINE_string('train_model_dir', ' ', '訓練輸出的模型目錄')
tf.app.flags.DEFINE_string('pre_trained_model', None, '預訓練模型目錄')

FLAGS = tf.app.flags.FLAGS


def main(_):

    if not FLAGS.dataset_dir:
        raise ValueError("必須指定一個TFRecord的數據集目錄")

    # 設置打印級別
    tf.logging.set_verbosity(tf.logging.DEBUG)

    # 在默認圖中進行訓練
    with tf.Graph().as_default():
        # 1.DeploymentConfig配置
        deploy_config = model_deploy.DeploymentConfig(num_clones=FLAGS.num_clones,
                                                      clone_on_cpu=FLAGS.clone_on_cpu,
                                                      replica_id=0,
                                                      num_replicas=1,
                                                      num_ps_tasks=0)

        # 在variables_device定義全局步長(網絡訓練一般都這么配置)
        with tf.device(deploy_config.variables_device()):
            global_step = tf.train.create_global_step()

        # 2.獲取圖片數據,做一些預處理
        # image, shape, bbox, label
        # 不是直接進行訓練,而是需要進行正負樣本標記(輸出的anchor和GT進行IOU計算選擇)

        # 2.1步驟如下:
        # (1)通過數據工廠獲取DataSet規范,不是真正的數據,需要通過后續操作去獲取數
        dataset = dataset_factory.get_dataset(dataset_name=FLAGS.dataset_name,
                                              train_or_test=FLAGS.train_or_test,
                                              dataset_dir=FLAGS.dataset_dir)

        # (2)通過網絡計算獲取的anchors結果
        # 通過網絡工廠獲取網絡
        ssd_class = nets_factory.get_network(FLAGS.network_name)

        # 獲取默認網絡參數
        ssd_params = ssd_class.default_params._replace(num_classes=9)

        # 初始化網絡init函數
        ssd_net = ssd_class(ssd_params)

        # 獲取shape
        ssd_shape = ssd_net.params.img_shape

        # 獲取anchors, SSD網絡中6層的所有計算出來的默認候選框default boxes
        ssd_anchors = ssd_net.anchors(ssd_shape)

        # (3)獲取預處理函數
        image_preprocessing_fn = preprocessing_factory.get_preprocessing(name=FLAGS.network_name,
                                                                         is_training=True)

        # 打印網絡相關參數
        train_tools.print_configuration(ssd_params, dataset.data_sources)

        # 2.2
        # (1)通過slim.dataset_data_provider.DatasetDataProvider獲取圖像數據
        # (2)進行數據預處理
        # (3)對獲取出來的GT標簽和bbox進行編碼
        # (4)獲取的單個樣本數據,要進行批處理以及返回隊列
        with tf.device(deploy_config.inputs_device()):
            with tf.name_scope(FLAGS.network_name + "_data_provider"):
                provider = slim.dataset_data_provider.DatasetDataProvider(
                    dataset,
                    num_readers=4,
                    common_queue_capacity=20 * FLAGS.batch_size,
                    common_queue_min=10 * FLAGS.batch_size,
                    shuffle=True)

                # get獲取數據(真正獲取參數)
                [image, shape, glabels, gbboxes] = provider.get(['image', 'shape', 'object/label', 'object/bbox'])

                # 數據預處理 [?, ?, 3]-->[300, 300, 3]
                image, glabels, gbboxes = image_preprocessing_fn(image, glabels, gbboxes, ssd_shape, DATA_FORMAT)

                # 原始anchor boxes進行正負樣本標記
                # gclasses: 目標類別
                # glocalizations: 目標類別的真實位置
                # gscores: 目標結果(概率值)
                gclasses, glocalizations, gscores = ssd_net.bboxes_encode(glabels, gbboxes, ssd_anchors)

                # 批處理、隊列處理
                # tensor_list:tensor組成的類別 [tensor, tensor, tensor, ...]
                # r是1個tensor組成的列表
                r = tf.train.batch(tensors=train_tools.reshape_list([image, gclasses, glocalizations, gscores]),
                                   batch_size=FLAGS.batch_size,
                                   num_threads=4,
                                   capacity=5 * FLAGS.batch_size)

                batch_queue = slim.prefetch_queue.prefetch_queue(r, capacity=deploy_config.num_clones)

        # 3.數據輸入、網絡計算結果、定義損失並復制模型到clones,添加變量到tensorboard
        summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES))
        # batch_shape:獲取的默認隊列大小,即上面r的大小
        batch_shape = [1] + 3 * [len(ssd_anchors)]
        update_ops, first_clone_scope, clones = train_tools.deploy_loss_summary(deploy_config,
                                                                                batch_queue,
                                                                                ssd_net,
                                                                                summaries,
                                                                                batch_shape,
                                                                                FLAGS)

        # 4.定義學習率、優化器
        # 初始學習率:0.001
        # 終止學習率:0.0001
        # 優化器選擇:adam
        with tf.device(deploy_config.optimizer_device()):
            # 定義學習率和優化器
            learning_rate = train_tools.configure_learning_rate(FLAGS, dataset.num_samples, global_step)

            # 定義優化器
            optimizer = train_tools.configure_optimizer(FLAGS, learning_rate)

            # 觀察學習的變化情況添加到summaries中
            summaries.add(tf.summary.scalar('learning_rate', learning_rate))

        # 5.計算所有GPU/CPU設備的平均損失和每個變量的梯度總和、定義訓練OP、summaries OP
        train_op, summaries_op = train_tools.get_trainop(optimizer,
                                                         summaries,
                                                         clones,
                                                         global_step,
                                                         first_clone_scope, update_ops)

        # 6.配置訓練的config,進行訓練
        # 6.1 配置config和saver
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8)
        config = tf.ConfigProto(log_device_placement=False,  # 若果打印會有許多變量的設備信息出現
                                gpu_options=gpu_options)

        saver = tf.train.Saver(max_to_keep=5,  # 默認保留最近幾個模型文件
                               keep_checkpoint_every_n_hours=1.0,
                               write_version=2,
                               pad_step_number=False)

        # 6.2 訓練
        slim.learning.train(
            train_op,  # 訓練優化器tensor
            logdir=FLAGS.train_model_dir,  # 模型存儲目錄
            master='',
            is_chief=True,
            init_fn=train_tools.get_init_fn(FLAGS),  # 初始化參數的邏輯,預訓練模型的讀取和微調模型判斷
            summary_op=summaries_op,  # 摘要
            number_of_steps=FLAGS.max_number_of_steps,  # 最大步數
            log_every_n_steps=10,  # 打印頻率
            save_summaries_secs=60,  # 保存摘要頻率
            saver=saver,  # 保存模型參數
            save_interval_secs=600,  # 保存模型間隔
            session_config=config,  # 會話參數配置
            sync_optimizer=None)


if __name__ == '__main__':
    tf.app.run()

 

訓練模型:

訓練的過程使用技嘉RTX2070Super顯卡。

切換到ObjectDetection目錄,執行如下命令(參數可以自己設定):

PRE_TRAINED_PATH=./ckpt/pre_trained/ssd_300_vgg.ckpt
TRAIN_MODEL_DIR=./ckpt/fine_tuning/
DATASET_DIR=./IMAGE/tfrecords/commodity_tfrecords/
python train_ssd_network.py --train_model_dir=${TRAIN_MODEL_DIR} --dataset_dir=${DATASET_DIR} --dataset_name="commodity_2018" --train_or_test=train --model_name=ssd_vgg_300 --pre_trained_path=${PRE_TRAINED_PATH} --weight_decay=0.0005 --optimizer=adam --learning_rate=0.001 --batch_size=16

 

此時可以學習。

 

同時在ckpt/fine_tuning文件夾下,執行如下命令,可以使用tensorboard查看已經添加到tensorboard中的相關參數。

tensorboard --logdir=./

 

訓練過程如下圖所示:

 

7.測試過程

7.1測試流程

1)測試數據准備

2)preprocessing數據預處理--測試過程的數據預處理就是需要圖片的resize

3)模型加載

4)postprocess(預測結果后期處理)--訓練過程中是不需要后期處理的

    通過scores篩選bbox

    使用NMS篩選box

    注意bbox邊界與原始圖片的bbox,按需修改bbox

5)預測結果顯示(使用matplotlib)

 

7.1 測試框架:

其中,test文件夾用於測試使用,visualization.py文件里面是顯示結果的代碼,test_image.py文件中文最終存放的測試代碼。

 

7.2 測試代碼

7.2.1顯示圖片代碼

visualization.py中的顯示結果的代碼如下:

import cv2
import random

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import matplotlib.cm as mpcm

VOC_LABELS = {
    '0': 'Background',
    '1': 'clothes',
    '2': 'pants',
    '3': 'shoes',
    '4': 'watch',
    '5': 'phone',
    '6': 'audio',
    '7': 'computer',
    '8': 'books'
}

# =========================================================================== #
# Matplotlib 顯示圖
# =========================================================================== #
def plt_bboxes(img, classes, scores, bboxes, figsize=(10,10), linewidth=1.5):
    """顯示bounding boxes.
    """
    fig = plt.figure(figsize=figsize)
    plt.imshow(img)
    height = img.shape[0]
    width = img.shape[1]
    colors = dict()
    for i in range(classes.shape[0]):
        cls_id = int(classes[i])
        if cls_id >= 0:
            score = scores[i]
            if cls_id not in colors:
                colors[cls_id] = (random.random(), random.random(), random.random())
            ymin = int(bboxes[i, 0] * height)
            xmin = int(bboxes[i, 1] * width)
            ymax = int(bboxes[i, 2] * height)
            xmax = int(bboxes[i, 3] * width)
            rect = plt.Rectangle((xmin, ymin), xmax - xmin,
                                 ymax - ymin, fill=False,
                                 edgecolor=colors[cls_id],
                                 linewidth=linewidth)
            plt.gca().add_patch(rect)
            class_name = str(cls_id)
            plt.gca().text(xmin, ymin - 2,
                           '{:s} | {:.3f}'.format(VOC_LABELS[class_name], score),
                           bbox=dict(facecolor=colors[cls_id], alpha=0.5),
                           fontsize=12, color='white')

    plt.show()

 

7.2.1測試過程代碼

測試過程test_image.py代碼如下:

import numpy as np
import tensorflow as tf
from PIL import Image

import sys
sys.path.append('../')

import matplotlib.pyplot as plt
import matplotlib.image as mping
import visualization
from utils.basic_tools import np_methods

slim = tf.contrib.slim

from nets import nets_factory
from preprocessing import preprocessing_factory

# 1.定義輸入圖片數據的占位符
image_input = tf.placeholder(tf.uint8, shape=[None, None, 3])

# 定義輸出形狀,元組表示
net_shape = (300, 300)

data_format = 'NHWC'

# 2.數據輸入預處理工廠,進行預處理
preprocessing_fn = preprocessing_factory.get_preprocessing('ssd_vgg_300', is_training=False)
image_Pre, _, _, bbox_img = preprocessing_fn(image_input, None, None, net_shape, data_format)

# image_Pre是三維形狀--->(300, 300, 3)
# 卷積神經網絡要求都是四維的數據計算
# 維度的擴充--->(1, 300, 300, 3)
image_4d = tf.expand_dims(image_Pre, 0)

# 3.定義SSD模型,並輸出預測結果
# 網絡工廠獲取
ssd_class = nets_factory.get_network('ssd_vgg_300')
ssd_params = ssd_class.default_params._replace(num_classes=9)

reuse = True if 'ssd_net' in locals() else False

# 初始化網絡
ssd_net = ssd_class(ssd_params)

ssd_anchors = ssd_net.anchors(net_shape)

# 通過網絡的方法獲取結果
# 使用slim指定公有參數
with slim.arg_scope(ssd_net.arg_scope(data_format=data_format)):
    predictions, localizations, _, _ = ssd_net.net(image_4d, is_training=False, reuse=reuse)


config = tf.ConfigProto(log_device_placement=False)
sess = tf.InteractiveSession(config=config)
sess.run(tf.global_variables_initializer())

ckpt_filepath = '../ckpt/fine_tuning/model.ckpt-103480'

saver = tf.train.Saver()
saver.restore(sess, ckpt_filepath)

# 會話運行圖片,輸出結束
# 讀取一張圖片
img = Image.open('../IMAGE/commodity/JPEGImages/000080.jpg').convert('RGB')

img = np.array(img)

i, p, l, box_img = sess.run([image_4d, predictions, localizations, bbox_img], feed_dict={image_input:img})

# 進行結果篩選
classes, scores, bboxes = np_methods.ssd_bboxes_select(
    p, l, ssd_anchors, select_threshold=0.5, img_shape=(300, 300),
    num_classes=9, decode=True
)

# bbox邊框不能超過原圖片,默認原圖的相對於bbox大小比例 [0, 0, 1, 1]
bboxes = np_methods.bboxes_clip(box_img, bboxes)

# 根據 scores 從大到小排序,並改變classes rbboxes的順序
classes, scores, bboxes = np_methods.bboxes_sort(classes, scores, bboxes, top_k=400)

# 使用nms算法篩選bbox
classes, scores, bboxes = np_methods.bboxes_nms(classes, scores, bboxes, nms_threshold=.45)

# 根據原始圖片的bbox,修改所有bbox的范圍[.0, .0, .1, .1]
bboxes = np_methods.bboxes_resize(box_img, bboxes)

visualization.plt_bboxes(img, classes, scores, bboxes)

 

測試中使用訓練得到的ckpt/fine_tuning/model.ckpt-103480文件中的參數進行。測試結果如下圖所示:

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM