以圖書識別模型爲例，新手如何從01到1 完成tensorflow模型訓練

項目體驗地址：http://at.iunitv.cn/

效果預覽：

花絮：

讀書節馬上就要到了，相信很多小夥伴嘴上說着學不動了，其實身體還是很誠實的。

畢竟讀書還是有很多好處的：比如讓你的腦門散發智慧的光芒，再或者讓你有理由說因爲讀書太忙了所以沒有女朋友等等。所以在這個特殊的日子裏，你這一年的圖書我們承包了。不爲別的，只爲幫助在座的各位在2020年能夠遇見更好的自己!

今天的主題僅僅是送圖書，我們也想要藉助這個特殊的機會，普及一下Tensorflow相關的知識，我們會用TensorFlow.js做一個圖書識別的模型，並在Vue Application中運行，賦予網頁識別圖書的能力。

本文講述了AI相關的概念知識和如何運用SSD Mobile Net V1模型進行遷移學習的方法，從而幫助大家完成一個可以在網頁上運行的圖書識別模型。

正文：

什麼是遷移學習

遷移學習和域適應指的是在一種環境中學到的知識被用在另一個領域中來提高它的泛化性能。——《深度學習》，第 526 頁

再簡單一點理解，以今天圖書識別模型訓練爲例，我們利用前人訓練好的具備圖片識別能力的AI模型，保留AI模型中對圖片特徵提取的能力的基礎上再訓練，使AI模型具備識別圖書的能力。

遷移學習能夠大大提高模型訓練的速度，並達到相對不錯的正確率。

而我們今天所要遷移學習的對象就是SSD Mobile Net V1模型，初次接觸神經網絡的同學可以將其理解爲一種具備圖片識別的輕便小巧的AI模型，它能夠在移動設備上高效地運行。對這個模型具體的神經網絡設計結構感興趣的同學可以自行搜索。

瞭解了基本的概念之後，我們便開始動手吧！我們可以基於SSD Mobile Net模型去設計一個屬於自己的AI模型，並讓它在Vue Application中運行。

Object Detection（目標識別）

本次項目是爲了訓練一個Object Detection的模型，即目標識別的模型，該模型能夠識別並圈選出圖片中相應的目標對象。

準備工作

同步開發環境

爲了避免小夥伴因爲環境問題遇到各種各樣的坑，在工作開展之前，我們先跟大家同步一下運行的環境。大家如果要動手去做，也儘量跟我們的運行環境保持一致，這樣可以有效避免踩坑，規避“從入門到放棄”的現象。

開發環境：

• 系統Mac OS系統
• Python版本：3.7.3
• TensorFlow版本：1.15.2
• TensorFlowJS版本：1.7.2
• 開發工具：Pycharm和Webstorm

下載項目

同步完開發環境後，終於要開始動工了。首先我們需要在Github上下載幾個項目：

• 遷移模型訓練項目
• 圖片格式轉換項目
• 圖片圈選工具LabelImg

準備圖片素材

我們可以通過搜索引擎收集有關圖書的圖片素材：

其次，我們可以在Github上克隆LabelImg項目，並根據Github的使用說明，按照不同的環境安裝運行LabelImg項目，運行後的頁面如下：

然後我們按照以下步驟，將圖片格式轉換爲圈選區域後的XML文件：

1. 打開圖片存放的目錄
2. 選擇圈選後的存放目錄
3. 圈選圖片目標區域
4. 設置圈選區域的標籤
5. 保存成XML格式

存放完後我們在存放的目錄下會看到許多XML格式的文件，這個文件記錄了圖片的位置信息、圈選信息和標籤信息等，用於後續的模型訓練。

配置安裝Object Detection的環境

從Github克隆遷移模型訓練的項目遷移模型訓練項目，注意要在r1.5分支運行，並用PyCharm打開項目。

項目的目錄環境爲上圖，首先我們需要下載TensorFlow1.15.2版本：

pip install tensorflow==1.15.2

其次安裝依賴包：

sudo pip install pillow
sudo pip install lxml
sudo pip install jupyter
sudo pip install matplotlib

然後通過終端切換到research目錄，並執行幾行配置命令，具體請參考Github的使用說明：

cd ./research
protoc object_detection/protos/*.proto --python_out=.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

最後我們運行model_builder_test.py文件，如果在終端中看到OK字樣，表示配置成功。

python object_detection/builders/model_builder_test.py

將XML格式轉換爲TensorFlow需要的TFRecord格式

克隆並打開圖片格式轉換項目，然後我們對該項目加以小改造：

改造文件目錄:

1. 刪除annotations、data、training目錄中的內容
2. 增加一個xmls目錄，用以存放xml文件

改造文件：
接着，我們再改造以下2個文件並新增一個文件，方便我們轉換圖片格式

1. 改造xml_to_csv.py爲：

   import os
   import glob
   import pandas as pd
   import xml.etree.ElementTree as ET
   import random
   import time
   import shutil
   
   class Xml2Cvs:
       def __init__(self):
           self.xml_filepath = r'./xmls'
           self.save_basepath = r"./annotations"
           self.trainval_percent = 0.9
           self.train_percent = 0.85
   
       def xml_split_train(self):
   
           total_xml = os.listdir(self.xml_filepath)
           num = len(total_xml)
           list = range(num)
           tv = int(num * self.trainval_percent)
           tr = int(tv * self.train_percent)
           trainval = random.sample(list, tv)
           train = random.sample(trainval, tr)
           print("train and val size", tv)
           print("train size", tr)
           start = time.time()
           test_num = 0
           val_num = 0
           train_num = 0
           for i in list:
               name = total_xml[i]
               if i in trainval:
                   if i in train:
                       directory = "train"
                       train_num += 1
                       xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
                       if (not os.path.exists(xml_path)):
                           os.mkdir(xml_path)
                       filePath = os.path.join(self.xml_filepath, name)
                       newfile = os.path.join(self.save_basepath, os.path.join(directory, name))
                       shutil.copyfile(filePath, newfile)
                   else:
                       directory = "validation"
                       xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
                       if (not os.path.exists(xml_path)):
                           os.mkdir(xml_path)
                       val_num += 1
                       filePath = os.path.join(self.xml_filepath, name)
                       newfile = os.path.join(self.save_basepath, os.path.join(directory, name))
                       shutil.copyfile(filePath, newfile)
               else:
                   directory = "test"
                   xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
                   if (not os.path.exists(xml_path)):
                       os.mkdir(xml_path)
                   test_num += 1
                   filePath = os.path.join(self.xml_filepath, name)
                   newfile = os.path.join(self.save_basepath, os.path.join(directory, name))
                   shutil.copyfile(filePath, newfile)
   
           end = time.time()
           seconds = end - start
           print("train total : " + str(train_num))
           print("validation total : " + str(val_num))
           print("test total : " + str(test_num))
           total_num = train_num + val_num + test_num
           print("total number : " + str(total_num))
           print("Time taken : {0} seconds".format(seconds))
   
       def xml_to_csv(self, path):
           xml_list = []
           for xml_file in glob.glob(path + '/*.xml'):
               tree = ET.parse(xml_file)
               root = tree.getroot()
               print(root.find('filename').text)
               for object in root.findall('object'):
                   value = (root.find('filename').text,
                            int(root.find('size').find('width').text),
                            int(root.find('size').find('height').text),
                            object.find('name').text,
                            int(object.find('bndbox').find('xmin').text),
                            int(object.find('bndbox').find('ymin').text),
                            int(object.find('bndbox').find('xmax').text),
                            int(object.find('bndbox').find('ymax').text)
                            )
                   xml_list.append(value)
           column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
           xml_df = pd.DataFrame(xml_list, columns=column_name)
           return xml_df
   
       def main(self):
           for directory in ['train', 'test', 'validation']:
               xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
               xml_df = self.xml_to_csv(xml_path)
               xml_df.to_csv('data/mask_{}_labels.csv'.format(directory), index=None)
               print('Successfully converted xml to csv.')
   
   
   if __name__ == '__main__':
       Xml2Cvs().xml_split_train()
       Xml2Cvs().main()
   

2. 改造generate_tfrecord.py文件，將csv格式轉換爲TensorFlow需要的record格式：
   

將該區域的row_label改成我們LabelImg中的標籤名，因爲我們只有一個標籤，所以直接修改成book即可。

3. 新增一個generate_tfrecord.sh腳本，方便執行generate_tfrecord.py文件

   #!/usr/bin/env bash
   python generate_tfrecord.py --csv_input=data/mask_train_labels.csv  --output_path=data/mask_train.record --image_dir=images
   python generate_tfrecord.py --csv_input=data/mask_test_labels.csv  --output_path=data/mask_test.record --image_dir=images
   python generate_tfrecord.py --csv_input=data/mask_validation_labels.csv  --output_path=data/mask_validation.record --image_dir=images
   
   

配置Object Decation的環境：

export PYTHONPATH=$PYTHONPATH:你的models/research/slim所在的全目錄路徑

最後我們將圖片文件複製到images目錄，將xml文件複製到xmls目錄下，再執行xml_to_csv.py文件，我們會看到data目錄下產生了幾個csv格式結尾的文件；這時，我們在終端執行generate_tfrecord.sh文件，TensorFlow所需要的數據格式就大功告成啦。

遷移訓練模型：

在這個環節我們要做以下幾件事：

• 將剛剛生成好的record文件放到對應目錄下
• 下載SSD Mobile Net V1模型文件
• 配置book.pbtxt文件和book.config文件

放置record文件和SSD Mobile Net V1模型

爲了方便我直接將models/research/object_detection/test_data下的目錄清空，放置遷移訓練的文件。

首先我們下載[SSD Mobile Net V1模型文件]

我們下載第一個ssd_mobilenet_v1_coco模型即可，下載完畢後，我們解壓下載的模型壓縮包文件，並將模型相關的文件放在test_data的model目錄下。並將我們剛剛生成的record文件放置在test_data目錄下。

完成pbtxt和config配置文件

我們在test_data目錄下，新建一個book.pbtxt文件，並完成配置內容:

item {
  id: 1
  name: 'book'
}

由於我們只有一個標籤，我們就直接配置一個id值爲1，name爲book的item對象。

由於我們使用SSD Mobile Net V1模型進行遷移學習，因此我們到sample\configs目錄下複製一份ssd_mobilenet_v1_coco.config文件並重命名爲book.config文件。

接着我們修改book.config中的配置文件：

將num_classes修改爲當前的標籤數量：

由於我們只有一個book標籤，因此修改成1即可。

修改所有PATH_TO_BE_CONFIGURED的路徑：

我們將此處的模型文件地址設置成testdata/model/model.ckpt的全路徑地址。

我們將train_input_reader的input_path設置成mask_train.record的全路徑地址；將label_map_path設置成book.pbtxt的全路徑地址；將eval_input_reader的input_path設置成mask_test.record的全路徑地址。

到目前爲止我們所有配置都已經完成啦。接下來就是激動人心的訓練模型的時刻。

運行train.py文件訓練模型

我們在終端中運行train.py文件，開始遷移學習、訓練模型。

python3 train.py --logtostderr --train_dir=./test_data/training/ --pipeline_config_path=./test_data/book.config

其中train_dir爲我們訓練後的模型存放的目錄，pipeline_config_path爲我們book.config文件所在的相對路徑。

運行命令後，我們可以看到模型在進行一步一步的訓練：

並在/test_data/training目錄下存放訓練後的模型文件:

將ckpt文件轉換爲pb文件

我們通過export_inference_graph.py文件，將訓練好的模型轉換爲pb格式的文件，這個文件格式在後面我們要用來轉換爲TensorFlow.js能夠識別的文件格式。終於我們見到TensorFlow.js的影子啦。

我們執行命令，運行export_inference_graph.py文件：

python export_inference_graph.py --input_type image_tensor --pipeline_config_path ./test_data/book.config --trained_checkpoint_prefix ./test_data/training/model.ckpt-1989 --output_directory ./test_data/training/book_model_test

其中pipeline_config_path爲book.config的相對文件路徑，trained_checkpoint_prefix爲模型文件的路徑，例如我們選擇訓練了1989步的模型文件，output_directory爲我們輸出pb文件的目標目錄。

運行完後，我們可以看到一個生成了book_model_test目錄：

將pb文件轉換爲TensorFlowJs模型

首先我們需要依賴TensorFlowjs的依賴包

pip install tensorflowjs

然後通過命令行轉換剛剛生成的pb文件

tensorflowjs_converter --input_format=tf_saved_model --output_node_names='detection_boxes,detection_classes,detection_features,detection_multiclass_scores,detection_scores,num_detections,raw_detection_boxes,raw_detection_scores' --saved_model_tags=serve --output_format=tfjs_graph_model ./saved_model ./web_model

其中我們設置最後兩個參數，即saved_model的目錄與TensorFlow.js識別模型的輸出目錄。

運行結束後，我們可以看到一個新生成的web_model目錄，其中包括了我們遷移學習訓練後的模型。

到這裏，模型訓練的階段終於結束了。

在Vue中運行模型

準備工作

新建Vue項目，在Vue項目的public目錄下放入我們訓練好的模型，即web_model目錄。

接着我們藉助Tensorflow.js的依賴包，在package.json的dependencies中加入：

"@tensorflow/tfjs": "^1.7.2",
"@tensorflow/tfjs-core": "^1.7.2",
"@tensorflow/tfjs-node": "^1.7.2",

然後通過npm命令安裝依賴包。

加載模型

在我們的JS代碼部分引入TensorFlow的依賴包：

import * as tf from '@tensorflow/tfjs';
import {loadGraphModel} from '@tensorflow/tfjs-converter';

接着第一步，我們先加載模型文件中的model.json文件：

const MODEL_URL = process.env.BASE_URL+"web_model/model.json";
this.model = await loadGraphModel(MODEL_URL);

通過loadGraphModel方法，我們加載好訓練的模型，再將模型對象打印出來：

隨後，我們可以看到模型會輸出一個長度爲4的數組:

• detection_scores：表示識別對象模型的置信度，置信度越高，則代表模型認爲對應區域識別爲書本的可能性越高
• detection_classes：表示模型識別的區域對應的標籤，例如在本案例中，識別出來的是book
• num_detections：表示模型識別出目標對象的個數
• detection_boxes：表示模型識別出來目標對象的區域，爲一個長度爲4的數組，分別是:[x_pos,y_pos,x_width,y_height] 。第一個位代表圈選區域左上角的x座標，第二位代表圈選左上角的y座標，第三位代表圈選區域的寬度，第四位代表圈選區域的長度。

模型識別

知道了輸出值，我們就可以開始將圖片輸入到模型中，從而得到模型預測的結果：

const img = document.getElementById('img');
let modelres =await this.model.executeAsync(tf.browser.fromPixels(img).expandDims(0));

我們通過model.executeAsync方法，將圖片輸入到模型中，從而得到模型的輸出值。

結果是我們前文提到的一個長度爲4的數組。接着我們通過自定義方法，將得到的結果進行整理，從而輸出一個想要的結果格式：

buildDetectedObjects:function(scores, threshold, imageWidth, imageHeight, boxes, classes, classesDir) {
          const detectionObjects = [];
          scores.forEach((score, i) => {
              if (score > threshold) {
                  const bbox = [];
                  const minY = boxes[i * 4] * imageHeight;
                  const minX = boxes[i * 4 + 1] * imageWidth;
                  const maxY = boxes[i * 4 + 2] * imageHeight;
                  const maxX = boxes[i * 4 + 3] * imageWidth;
                  bbox[0] = minX;
                  bbox[1] = minY;
                  bbox[2] = maxX - minX;
                  bbox[3] = maxY - minY;
                  detectionObjects.push({
                      class: classes[i],
                      label: classesDir[classes[i]].name,
                      score: score.toFixed(4),
                      bbox: bbox
                  });
              }
          });

          return detectionObjects
}

我們通過調用buildDetectedObjects來整理和返回最後的結果。

• scores：輸入模型的detection_scores數組
• threshold：閾值，即結果score>threshold我們纔會將對應的結果放入結果對象detectionObjects中
• imageWidth：圖片的寬度
• imageHeight：圖片的長度
• boxes：輸入模型的detection_boxes數組
• classes：輸入模型的detection_classes數組
• classesDir：即模型標籤對象

調用buildDetectedObjects方法示例：

let classesDir = {
	1: {
		name: 'book',
		id: 1,
		}
	};
let res=this.buildDetectedObjects(modelres[0].dataSync(),0.20,img.width,img.height,modelres[3].dataSync(),modelres[1].dataSync(),classesDir);

我們通過modelres[0].dataSync()，來獲取對應結果的數組對象，再輸入到方法中，從而最終獲得res結果對象。

最後我們通過Canvas的API，將圖片根據bbox返回的數組對象，畫出對應的區域即可。由於篇幅原因，就不贅述了，最終效果如下：

最後

本案例的模型存在一定的不足，由於訓練時間較短，圖書的封面類型衆多，存在人像、風景圖等等的樣式，導致模型在識別過程中可能會將少部分的人臉、風景照等圖片錯誤地識別成圖書封面。各位小夥伴在訓練自己模型的過程中可以考慮優化此問題。

當然，本案例的模型在識別非圖書的場景會存在識別不準確的情況，一方面這是因爲本案例從網絡收集的圖書樣本具有一定侷限性，而且圖書的封面類型千差萬別，存在人像、風景圖等等的樣式；另一方面因爲本文在僅爲起到拋磚引玉的作用，爲各位前端小夥伴普及TensorFlow.js相關的知識，並提供訓練自己的模型的解決方案，所以在收集樣本和模型訓練時間較短。感興趣的小夥伴可以自己琢磨琢磨如何優化樣本和在避免過擬合的情況下提高訓練時長，從而提高模型對被識別物體的準確性。

我們寫下本文僅爲起到拋磚引玉的作用，爲各位前端小夥伴普及TensorFlow.js相關知識並提供一種AI的解決方案。

在世界讀書日，我們希望和廣大程序員一起學習新知、共同進步，個推技術學院也特地爲大家準備了微信讀書卡，願每位熱愛學習的開發者都能暢遊書海，遇見更好的自己！

活動獎品：

一等獎 ：極客時間充值卡1張，1名

二等獎：得到電子書vip年卡1張，3名

三等獎5名：《深度學習》1本，5名

抽獎方式：

掃描下方二維碼關注個推技術學院公衆號，後臺回覆“我愛讀書”獲取抽獎入口，點擊即可參與抽獎。

開獎時間：2020年4月27日 16:00，系統將隨機抽取出幸運粉絲。

領取方式：請中獎者於24小時內在抽獎助手中填寫收件信息，我們會在7個工作日之內爲您寄出。

注：活動解釋權歸個推所有。