TensorFlowLite+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地址:github/tensorflow/models.git
1.2 编译protobuf库,⽤来配置模型和训练参数,下载直接编译好的pb库(github/google/protobuf/releases ),解压压缩包后,添加环境变量:
1.3 将models和slim加⼊python环境变量:女孩好名字
2,数据准备
TensorFlow Object Detection API训练需要标注好的图像,推荐使⽤ labelImg,是⼀个开源的图像标注⼯具,下载链接:
github/tzutalin/labelImg。标注完样本之后会⽣成⼀个xml的标注⽂件,这些xml⽂件我们需要最终转换为训练⽤的TFRecord类型⽂件,GitHub上有个demo提供了很⽅便的转换脚本(github/datitran/raccoon_dataset)。我们把这些标注的xml⽂件,按训练集与验证集分别放置到两个⽬录下,通过下载的xml_to_csv.py脚本转换为csv结构数据。然后使⽤转换为TFRerord 格式的脚本:generate_tfrecord.py把对应的csv格式转换成.record格式。
labelImg界⾯:
3,训练
打开下载后的coco数据集预训练模型的⽂件夹,把model.ckpt⽂件放置在待训练的⽬录,修改ssd_mobilenet_fig⽂件中的两个地⽅:
1,num_classes:修改为⾃⼰的classes num
2,将所有PATH_TO_BE_CONFIGURED的地⽅修改为⾃⼰之前设置的路径
调⽤train.py开始训练:
力不从心什么意思pipelineconfigpath是训练的配置⽂件路径train_dir是训练输出的路径
1.2 模型的优化和转换
最后将训练得到的pb模型,使⽤官⽅的optimize_for_inference优化,再⽤toco转换为tflite模型(路径需要修改),参照官⽅GitHub更新的这个Issues:
1.3 tflite 端上执⾏ssd
我们在这个案例中使⽤的ssd_mobilenet.tflite模型,输⼊输出数据类型为float32。SSD中没有全连接层,可适应各种⼤⼩的图⽚,我们的这个模型取的shape是{1, 300, 300, 3}。
图⽚输⼊的代码如下:
输出的结构是包含Locations和Classes的数组,代码如下:
通过遍历输出,并使⽤sigmoid激活函数,得到score,保存⼤于0.8时的class与location的index
outputLocations解析:
快手宏楠视频通过上述⽅法处理,outputLocations->data.f 4个值⼀组表⽰输出的矩形框左上⾓和右下⾓坐标,然后遍历resultArr取score⼤于0.8时对应的classIndex与locationIndex,再通过如下代码得到框的坐标并输出识别出的类别与分数:
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来直接修复
图⽚的⽅式去⽔印,有做过相关尝试的欢迎⼀起探讨。