【OpenCV】OpenCV4调用tensorflow pb模型进行目标检测
先在visual studio 2015新建一个工程并配置好环境。
OpenCV调用tensorflow的pb模型,不仅需要一个.pb文件,还需要一个.pbtxt文件。
一些预训练好的pb模型的下载地址,在opencv_extra里有一些生成好的pbtxt文件。注意pb文件和pbtxt文件要对应,否则后面调用会出错。
如果没有pb文件对应的pbtxt文件,则在opencv里提供了tf_text_graph_faster_rcnn.py,tf_text_graph_mask_rcnn.py,tf_text_graph_ssd.py三个脚本来生成对应网络结构的pbtxt文件。用这个脚本生成pbtxt文件不仅要输入pb文件,还需要输入.config文件,这里有一些config文件。
使用方式如下:
python tf_text_graph_ssd.py \
--input /path/to/xxx.pb \
--config /path/to/xxx.config \
--output /path/to/xxx.pbtxt
准备就绪之后,可以开始写OpenCV的代码了。这里有几个关键点。
OpenCV的dnn模块提供了readNetFromTensorflow来读取tensorflow的pb模型,当然还有readNetFromDarknet,readNetFromCaffe,readNetFromTorch等等。这里就只关注tensorflow了。
dnn::Net net = dnn::readNetFromTensorflow(pbFile, pbTxt);
vector<Mat> outs;
vector<String> outNames = net.getUnconnectedOutLayersNames();
Mat inputBlob = dnn::blobFromImage(cropImg, 1.0 / 255, Size(cropWidth, cropHeight), Scalar(), false, false);
net.setInput(inputBlob);
net.forward(outs, outNames);
其中, pbFile就是pb模型的路径,pbTxt就是pbtxt文件的路径。dnn::blobFromImage的1.0 / 255将数据归一化到0~1。
forward输出的outs是一个1 x 1 x N x 7的矩阵,N是预测结果的数量,可能很多,还需要进行后处理。
float* data = (float*)outs[0].data;
for (int i = 0; i < outs[0].total(); i += 7)
{
float confidence = data[i + 2];
if (confidence > confidenceTH)
{
int left = (int)(data[i + 3] * img.cols);
int top = (int)(data[i + 4] * img.rows);
int right = (int)(data[i + 5] * img.cols);
int bottom = (int)(data[i + 6] * img.rows);
int width = right - left + 1;
int height = bottom - top + 1;
classIds.push_back((int)(data[i + 1])-1);
boxes.push_back(Rect(left, top, width, height));
confidences.push_back(confidence);
}
}
因为是要进行目标检测,所以这里将结果转换成了boundingbox的形式,并且将confidence低于阈值的去掉。
接下来进行非极大值抑制,将重合度高的boundingbox去掉。
vector<int> indices;
dnn::NMSBoxes(boxes, confidences, confidenceTH, nmsTH, indices);
预测已经完成,将结果画出来。
rectangle(img, Point(left, top), Point(right, bottom), Scalar(0, 255, 0));
下载的很多模型效果都不是很好,精度很低。用脚本生成的pbtxt有时候不能使用。
最后成功的是在ssd_mobilenet_v1_coco_11_06_2017下载的,pbtxt文件用的是现成的ssd_mobilenet_v1_coco.pbtxt。有空了再试试其他的模型。