TensorFlow Lite+OpenCV實現移動端水印的檢測與去除

閒魚技術發表於2018-10-22
OpenCV

概要：

本篇文章介紹了TensorFlow Lite與OpenCV配合使用的一個應用場景，並詳細介紹了其中用到的SSD模型從訓練到端上使用的整個鏈路流程。在APP中的使用場景為，使用者在釋出圖片時，在端上實現水印的檢測和定位，並提供去水印的功能。

具體步驟有：

  • 1，使用TensorFlow Object Detection API進行SSD模型的訓練

  • 2，模型的優化和轉換，模型在端上的解析使用（本篇主要使用iOS端的C++程式碼作為示例）

  • 3，將輸出locations值通過NMS(非極大值抑制)演算法得到最優的框

  • 4，使用OpenCV去除水印

使用的庫及工具：

TensorFlow v:1.8r +

TensorFlowLite v:0.0.2 +

OpenCV

labelImg

TensorFlow Lite+OpenCV實現移動端水印的檢測與去除

SSD檢測並定位水印

SSD簡介

SSD，全稱Single Shot MultiBox Detector，是Wei Liu在ECCV 2016上提出的一種目標檢測演算法，截至目前是主要的檢測框架之一，相比Faster RCNN有明顯的速度優勢，相比YOLO又有明顯的mAP優勢（不過已經被CVPR 2017的YOLO9000超越）。SSD具有如下主要特點：

1，從YOLO中繼承了將detection轉化為regression的思路，同時一次即可完成網路訓練

2，基於Faster RCNN中的anchor，提出了相似的prior box

3，加入基於特徵金字塔（Pyramidal Feature Hierarchy）的檢測方式，相當於半個FPN思路

TensorFlow Object Detection API提供了多種目標檢測的網路結構預訓練的權重，全部是用COCO資料集進行訓練，各個模型的精度和計算所需時間如下：

TensorFlow Lite+OpenCV實現移動端水印的檢測與去除

我們直接使用TensorFlow提供的模型重訓練，可以專注於工程不用重新構建網路，本文選用模型為SSD-300 mobilenet-based

1.1 模型的訓練

1，配置環境

1.1 下載TensorFlow Object Detection API程式碼庫，Git地址：https://github.com/tensorflow/models.git

1.2 編譯protobuf庫，用來配置模型和訓練引數，下載直接編譯好的pb庫（https://github.com/google/protobuf/releases ），解壓壓縮包後，新增環境變數：

  1. $ cd tensorflow/models
  2. 

  3. $ protoc object_detection/protos/*.proto --python_out=.

1.3 將models和slim加入python環境變數：

  1. PYTHONPATH=$PYTHONPATH:/your/path/to/tensorflow/models:/your/path/to/tensorflow/models/slim

2，資料準備

TensorFlow Object Detection API訓練需要標註好的影象，推薦使用labelImg，是一個開源的影象標註工具，下載連結：https://github.com/tzutalin/labelImg。標註完樣本之後會生成一個xml的標註檔案，這些xml檔案我們需要最終轉換為訓練用的TFRecord型別檔案，GitHub上有個demo提供了很方便的轉換指令碼（https://github.com/datitran/raccoon_dataset）。我們把這些標註的xml檔案，按訓練集與驗證集分別放置到兩個目錄下，通過下載的xml_to_csv.py指令碼轉換為csv結構資料。然後使用轉換為TFRerord格式的指令碼：generate_tfrecord.py把對應的csv格式轉換成.record格式。

  1. python generate_tfrecord.py --csv_input=test_labels.csv --output_path=test.record

labelImg介面：

TensorFlow Lite+OpenCV實現移動端水印的檢測與去除

3，訓練

開啟下載後的coco資料集預訓練模型的資料夾，把model.ckpt檔案放置在待訓練的目錄，修改ssd_mobilenet_v1_pets.config檔案中的兩個地方：

1，num_classes:修改為自己的classes num

2，將所有PATH_TO_BE_CONFIGURED的地方修改為自己之前設定的路徑

呼叫train.py開始訓練：

  1. python object_detection/train.py \
  2. --logtostderr \
  3. --pipeline_config_path= /your/path/training-sets /data-translate/training/ssd_mobilenet_v1_pets.config \
  4. --train_dir= /your/path/training-sets/data-translate/training

pipelineconfigpath是訓練的配置檔案路徑

train_dir是訓練輸出的路徑

1.2 模型的優化和轉換

最後將訓練得到的pb模型，使用官方的optimize_for_inference優化，再用toco轉換為tflite模型（路徑需要修改），參照官方GitHub更新的這個Issues：

  1. DETECT_PB=$PWD/ssd_mobilenet_v1_coco_2017_11_17/frozen_inference_graph.pb
  2. STRIPPED_PB=$PWD/frozen_inference_graph_stripped.pb
  3. DETECT_FB=$PWD/tensorflow/contrib/lite/examples/android/assets/mobilenet_ssd.tflite
  4. 

  5. # Strip out problematic nodes before even letting TOCO see the graphdef
  6. bazel run -c opt tensorflow/python/tools/optimize_for_inference -- \
  7. --input=$DETECT_PB  --output=$STRIPPED_PB --frozen_graph=True \
  8. --input_names=Preprocessor/sub --output_names=concat,concat_1 \
  9. --alsologtostderr
  10. 

  11. # Run TOCO conversion.
  12. bazel run tensorflow/contrib/lite/toco:toco -- \
  13. --input_file=$STRIPPED_PB --output_file=$DETECT_FB \
  14. --input_format=TENSORFLOW_GRAPHDEF --output_format=TFLITE \
  15. --input_shapes=1,300,300,3 --input_arrays=Preprocessor/sub \
  16. --output_arrays=concat,concat_1 --inference_type=FLOAT --logtostderr

1.3 tflite 端上執行ssd

我們在這個案例中使用的ssd_mobilenet.tflite模型，輸入輸出資料型別為float32。SSD中沒有全連線層，可適應各種大小的圖片，我們的這個模型取的shape是{1, 300, 300, 3}。

圖片輸入的程式碼如下：

  1. NSString* image_path = FilePathForResourceName(@"test_img", @"jpg");
  2.   int image_width;
  3.   int image_height;
  4.   int image_channels;
  5.   std::vector<uint8_t> image_data =
  6.       LoadImageFromFile([image_path UTF8String], &image_width, &image_height, &image_channels);
  7.   const int wanted_width = 300;
  8.   const int wanted_height = 300;
  9.   const int wanted_channels = 3;
  10.   const float input_mean = 128.0f;
  11.   const float input_std = 128.0f;
  12.   assert(image_channels >= wanted_channels);
  13.   uint8_t* in = image_data.data();
  14.   float* out = interpreter->typed_tensor<float>(input);
  15.   for (int y = 0; y < wanted_height; ++y) {
  16.     const int in_y = (y * image_height) / wanted_height;
  17.     uint8_t* in_row = in + (in_y * image_width * image_channels);
  18.     float* out_row = out + (y * wanted_width * wanted_channels);
  19.     for (int x = 0; x < wanted_width; ++x) {
  20.       const int in_x = (x * image_width) / wanted_width;
  21.       uint8_t* in_pixel = in_row + (in_x * image_channels);
  22.       float* out_pixel = out_row + (x * wanted_channels);
  23.       for (int c = 0; c < wanted_channels; ++c) {
  24.         out_pixel[c] = (in_pixel[c] - input_mean) / input_std;
  25.       }
  26.     }
  27.   }

輸出的結構是包含Locations和Classes的陣列，程式碼如下：

  1. if (interpreter->Invoke() == kTfLiteOk) {
  2.     const std::vector<int>& results = interpreter->outputs();
  3.     TfLiteTensor* outputLocations = interpreter->tensor(results[0]);
  4.     TfLiteTensor* outputClasses = interpreter->tensor(results[1]);
  5.     float *data = tflite::GetTensorData<float>(outputClasses); 
  6. }

通過遍歷輸出，並使用sigmoid啟用函式，得到score，儲存大於0.8時的class與location的index

  1. for(int i=0;i<NUM_RESULTS;i++)
  2.     {
  3.         for(int j=1;j<NUM_CLASSES;j++)
  4. 

  5.         {
  6.             float score = expit(data[i*NUM_CLASSES+j]);
  7.             if (0.8 < score) {
  8.                 [resultArr addObject:@{@"score":@(score),
  9.                                               @"locationIndex":@(i),
  10.                                               @"classIndex":@(j)}];
  11.             }
  12.         }
  13.     }
  14. 

  15. decodeCenterSizeBoxes(outputLocations->data.f);//對outputLocations解析

outputLocations解析：

  1. static void decodeCenterSizeBoxes(float* predictions) {
  2.     for (int i = 0; i < NUM_RESULTS; ++i) {
  3.         float ycenter = predictions[i * 4 + 0] / Y_SCALE * boxPriorsArr[2][i] + boxPriorsArr[0][i];
  4.         float xcenter = predictions[i * 4 + 1] / X_SCALE * boxPriorsArr[3][i]  + boxPriorsArr[1][i];
  5.         float h = (float) std::exp(predictions[i * 4 + 2] / H_SCALE) * boxPriorsArr[2][i];
  6.         float w = (float) exp(predictions[i * 4 + 3] / W_SCALE) * boxPriorsArr[3][i];
  7. 

  8.         float ymin = ycenter - h / 2.f;
  9.         float xmin = xcenter - w / 2.f;
  10.         float ymax = ycenter + h / 2.f;
  11.         float xmax = xcenter + w / 2.f;
  12. 

  13.         predictions[i * 4 + 0] = ymin;
  14.         predictions[i * 4 + 1] = xmin;
  15.         predictions[i * 4 + 2] = ymax;
  16.         predictions[i * 4 + 3] = xmax;
  17.     }
  18. }

通過上述方法處理，outputLocations->data.f 4個值一組表示輸出的矩形框左上角和右下角座標，然後遍歷resultArr取score大於0.8時對應的classIndex與locationIndex，再通過如下程式碼得到框的座標並輸出識別出的類別與分數：

  1. int top = (outputLocations->data.f)[locationIndex * 4 + 0]  * 300;
  2. int left = (outputLocations->data.f)[locationIndex * 4 + 1] * 300;
  3. int right = (outputLocations->data.f)[locationIndex * 4 + 2] * 300;
  4. int bottom = (outputLocations->data.f)[locationIndex * 4 + 3] * 300;
  5. 

  6. NSLog(@"Predictions: %@", [NSString stringWithFormat:@"%s - %f", label_strings[classIndex].c_str(), score]);

1.4 非極大值抑制（NMS）

TensorFlow Lite+OpenCV實現移動端水印的檢測與去除

解析之後，一個物體會得到了多個定位的框，如何確定哪一個是我們需要的最準確的框呢？我們就要用到非極大值抑制，來抑制那些冗餘的框：抑制的過程是一個迭代-遍歷-消除的過程。

  • 1，將所有框的得分排序，選中最高分及其對應的框

  • 2，遍歷其餘的框，如果和當前最高分框的重疊面積(IOU)大於一定閾值，我們就將框刪除。

  • 3，從未處理的框中繼續選一個得分最高的，重複上述過程。

處理之後：

TensorFlow Lite+OpenCV實現移動端水印的檢測與去除

OpenCV去水印

Opencv去水印有兩種方法：

一種是直接用inpainter函式（處理質量較低，可以處理不透明水印），另一種是基於畫素的反色中和（處理質量較高，只能處理半透明水印，未驗證）

inpainter函式：

演算法理論：基於Telea在2004年提出的基於快速行進的修復演算法（FMM演算法），先處理待修復區域邊緣上的畫素點，然後層層向內推進，直到修復完所有的畫素點

處理方式：獲取到黑底白色水印且相同位置的水印蒙版圖(必須單通道灰度圖)，然後使用inpaint方法處理原始影象，因為SSD得到的定位大小很難完全精確，具體使用時可把mask水印區適當放大，因為這個方法的處理是從邊緣往內執行，這樣可以保證水印能完全被mask覆蓋

通過水印位置和水印樣式生成如下mask圖（大小與原圖保持一致）

TensorFlow Lite+OpenCV實現移動端水印的檢測與去除

處理之後：

TensorFlow Lite+OpenCV實現移動端水印的檢測與去除

基於畫素的反色中和：

這種方法可以針對固定位置半透明水印做去除，演算法原理是使用水印mask圖，對加水印的圖片做反向運算，計算出水印位置原來的顏色值。

總結

TensorFlow lite在4月份才做了對SSD的支援，目前文件比較缺乏，並且官方只提供了安卓例項，iOS的C++程式碼對輸入輸出的處理需要根據安卓demo的程式碼來推測，比如對結果中classes的解析以及對輸出的框位置的解析，還有需要進行nms演算法取最優等。還有一個問題就是由於TensorFlow更新比較快，TensorFlow Object Detection API中很多方法引數和路徑各版本存在差異，需要注意。

在實際的應用中，水印的位置基本會比較固定，在圖片的4個角或居中，所以在ssd檢測過程中，後續可以考慮新增規則或者嘗試使用注意力模型來增加四個角以及中間部分的處理權重，來提高效率和準確率。目前這個方法還存在一個問題，就是必須要提前知道水印的具體樣式，並將包含這些水印的圖片做訓練，如果有新的水印就無法對其做出正確的識別和去除，後期我們會嘗試通過GAN來直接修復圖片的方式去水印，有做過相關嘗試的歡迎一起探討。

相關文章