【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.pytf_text_graph_mask_rcnn.pytf_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】OpenCV4调用tensorflow pb模型进行目标检测

准备就绪之后,可以开始写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));

【OpenCV】OpenCV4调用tensorflow pb模型进行目标检测

下载的很多模型效果都不是很好,精度很低。用脚本生成的pbtxt有时候不能使用。

最后成功的是在ssd_mobilenet_v1_coco_11_06_2017下载的,pbtxt文件用的是现成的ssd_mobilenet_v1_coco.pbtxt。有空了再试试其他的模型。