TensorFlow Lite+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

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提供的模型重训练,可以专注于工程不用重新构建网络,本文选用模型为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. $ 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界面:

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. # Strip out problematic nodes before even letting TOCO see the graphdef

  5. bazel run -c opt tensorflow/python/tools/optimize_for_inference -- \

  6. --input=$DETECT_PB  --output=$STRIPPED_PB --frozen_graph=True \

  7. --input_names=Preprocessor/sub --output_names=concat,concat_1 \

  8. --alsologtostderr

  9. # Run TOCO conversion.

  10. bazel run tensorflow/contrib/lite/toco:toco -- \

  11. --input_file=$STRIPPED_PB --output_file=$DETECT_FB \

  12. --input_format=TENSORFLOW_GRAPHDEF --output_format=TFLITE \

  13. --input_shapes=1,300,300,3 --input_arrays=Preprocessor/sub \

  14. --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.            float score = expit(data[i*NUM_CLASSES+j]);

  6.            if (0.8 < score) {

  7.                [resultArr addObject:@{@"score":@(score),

  8.                                              @"locationIndex":@(i),

  9.                                              @"classIndex":@(j)}];

  10.            }

  11.        }

  12.    }

  13. 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.        float ymin = ycenter - h / 2.f;

  8.        float xmin = xcenter - w / 2.f;

  9.        float ymax = ycenter + h / 2.f;

  10.        float xmax = xcenter + w / 2.f;

  11.        predictions[i * 4 + 0] = ymin;

  12.        predictions[i * 4 + 1] = xmin;

  13.        predictions[i * 4 + 2] = ymax;

  14.        predictions[i * 4 + 3] = xmax;

  15.    }

  16. }

通过上述方法处理,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. NSLog(@"Predictions: %@", [NSString stringWithFormat:@"%s - %f", label_strings[classIndex].c_str(), score]);

1.4 非极大值抑制(NMS)

解析之后,一个物体会得到了多个定位的框,如何确定哪一个是我们需要的最准确的框呢?我们就要用到非极大值抑制,来抑制那些冗余的框:抑制的过程是一个迭代-遍历-消除的过程。

  • 1,将所有框的得分排序,选中最高分及其对应的框

  • 2,遍历其余的框,如果和当前最高分框的重叠面积(IOU)大于一定阈值,我们就将框删除。

  • 3,从未处理的框中继续选一个得分最高的,重复上述过程。

处理之后:

OpenCV去水印

Opencv去水印有两种方法:

一种是直接用inpainter函数(处理质量较低,可以处理不透明水印),另一种是基于像素的反色中和(处理质量较高,只能处理半透明水印,未验证)

inpainter函数:

算法理论:基于Telea在2004年提出的基于快速行进的修复算法(FMM算法),先处理待修复区域边缘上的像素点,然后层层向内推进,直到修复完所有的像素点

处理方式:获取到黑底白色水印且相同位置的水印蒙版图(必须单通道灰度图),然后使用inpaint方法处理原始图像,因为SSD得到的定位大小很难完全精确,具体使用时可把mask水印区适当放大,因为这个方法的处理是从边缘往内执行,这样可以保证水印能完全被mask覆盖

通过水印位置和水印样式生成如下mask图(大小与原图保持一致)

处理之后:

基于像素的反色中和:

这种方法可以针对固定位置半透明水印做去除,算法原理是使用水印mask图,对加水印的图片做反向运算,计算出水印位置原来的颜色值。

总结

TensorFlow lite在4月份才做了对SSD的支持,目前文档比较缺乏,并且官方只提供了安卓实例,iOS的C++代码对输入输出的处理需要根据安卓demo的代码来推测,比如对结果中classes的解析以及对输出的框位置的解析,还有需要进行nms算法取最优等。还有一个问题就是由于TensorFlow更新比较快,TensorFlow Object Detection API中很多方法参数和路径各版本存在差异,需要注意。

在实际的应用中,水印的位置基本会比较固定,在图片的4个角或居中,所以在ssd检测过程中,后续可以考虑添加规则或者尝试使用注意力模型来增加四个角以及中间部分的处理权重,来提高效率和准确率。目前这个方法还存在一个问题,就是必须要提前知道水印的具体样式,并将包含这些水印的图片做训练,如果有新的水印就无法对其做出正确的识别和去除,后期我们会尝试通过GAN来直接修复图片的方式去水印,有做过相关尝试的欢迎一起探讨。

闲鱼技术
闲鱼技术

加入闲鱼,一起玩些酷的。(阿里巴巴集团闲鱼官方技术号,欢迎同道者技术交流。) 简历投递:guicai.gxy@alibaba-inc.com

工程OpenCVTensorFlow
6
相关数据
激活函数技术

在 计算网络中, 一个节点的激活函数定义了该节点在给定的输入或输入的集合下的输出。标准的计算机芯片电路可以看作是根据输入得到"开"(1)或"关"(0)输出的数字网络激活函数。这与神经网络中的线性感知机的行为类似。 一种函数(例如 ReLU 或 S 型函数),用于对上一层的所有输入求加权和,然后生成一个输出值(通常为非线性值),并将其传递给下一层。

权重技术

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

TensorFlow技术

TensorFlow是一个开源软件库,用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品,如语音识别、Gmail、Google 相册和搜索,其中许多产品曾使用过其前任软件DistBelief。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

验证集技术

验证数据集是用于调整分类器超参数(即模型结构)的一组数据集,它有时也被称为开发集(dev set)。

准确率技术

分类模型的正确预测所占的比例。在多类别分类中,准确率的定义为:正确的预测数/样本总数。 在二元分类中,准确率的定义为:(真正例数+真负例数)/样本总数

OpenCV技术

OpenCV的全称是Open Source Computer Vision Library,是一个跨平台的计算机视觉库。OpenCV是由英特尔公司发起并参与开发,以BSD许可证授权发行,可以在商业和研究领域中免费使用。OpenCV可用于开发实时的图像处理、计算机视觉以及模式识别程序。

暂无评论
暂无评论~