Windows Tensorflow C++ Objecct detection API Opencv

Object Detection API C++ Opencv 摄像头

参考TensorFlow在windows上编译https://blog.csdn.net/shakevincent/article/details/80306806

 
 
// Copyright 2015 The TensorFlow Authors. All Rights Reserved.
// Object detection API

#include "tensorflow_test.h"
#include <fstream>
#include <utility>
#include <vector>
#include <iostream>


#include "tensorflow/core/public/session.h"
#include "tensorflow/cc/ops/standard_ops.h"

#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/default_device.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/util/command_line_flags.h"


#include <google/protobuf/arena.h>
#include <google/protobuf/arenastring.h>
#include <google/protobuf/generated_message_util.h>
#include <google/protobuf/metadata.h>
#include <google/protobuf/repeated_field.h>
#include <google/protobuf/extension_set.h>
#include <google/protobuf/generated_enum_reflection.h>

#include <opencv.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include<opencv2/core/core.hpp>
#include<opencv2/highgui/highgui.hpp>

#include <sstream> 
// These are all common classes it's handy to reference with no namespace.
using tensorflow::Flag;
using tensorflow::Tensor;
using tensorflow::Status;
using tensorflow::string;
using tensorflow::int32;


using namespace cv;
using namespace tensorflow;

float const TF_PREDICTION_THRESSHOLD = 0.5;
std::map<int, std::string> TF_LabelMap;
std::unique_ptr<tensorflow::Session> TF_Session;

// Takes a file name, and loads a list of labels from it, one per line into the map object. Expects `CLASSID: CLASSNAME` fmt
static bool TF_loadLabels(const std::string& file_name, std::map<int, std::string>* result, int* found_label_count)
{
	std::ifstream file(file_name);
	if (!file)
	{
		return false;
	}
	result->clear();
	*found_label_count = 0;
	std::string line;
	while (std::getline(file, line))
	{
		std::string::size_type sz;   // alias of size_t
		//std::cout << line << std::endl;
		int i_decimal = std::stoi(line, &sz);
		(*result)[i_decimal] = line.substr(sz + 2); // +2 to account for ':' and following space
		(*found_label_count)++;
	}
	return true;
}
// Reads a model graph definition from disk, and creates a session object you can use to run it.
tensorflow::Status TF_LoadGraph(const std::string& graph_file_name, std::unique_ptr<tensorflow::Session>* session) {
	tensorflow::GraphDef graph_def;
	tensorflow::Status load_graph_status =
		tensorflow::ReadBinaryProto(tensorflow::Env::Default(), graph_file_name, &graph_def);
	if (!load_graph_status.ok()) {
		return tensorflow::errors::NotFound("Failed to load compute graph at '", graph_file_name, "'");
	}
	session->reset(tensorflow::NewSession(tensorflow::SessionOptions()));
	return (*session)->Create(graph_def);
}
static bool TF_init(const std::string& labels_file_name, std::map<int, std::string>* label_map, const std::string& graph_file_name, std::unique_ptr<tensorflow::Session>* tf_session) {
	int argc = 0;
	tensorflow::port::InitMain(NULL, &argc, NULL);

	int label_count;
	if (!TF_loadLabels(labels_file_name, label_map, &label_count)) {
		std::cerr << "TF_loadLabels ERROR" << std::endl;
	}
	else {
		std::cout << "Loaded " << label_count << " dnn class labels" << std::endl;
	}

	tensorflow::Status status = TF_LoadGraph(graph_file_name, tf_session);
	if (!status.ok()) {
		std::cerr << "TF_LoadGraph ERROR: " << status.error_message() << std::endl;
		return false;
	}
	return true;
}

static int PredictionsAndDraw(cv::Mat * cvImg, std::unique_ptr<tensorflow::Session>* tf_session, const float predictionThreshold)
{
	int inputHeight = cvImg->size().height;
	int inputWidth = cvImg->size().width;
	tensorflow::Tensor imgTensorWithSharedData(tensorflow::DT_UINT8, { 1, inputHeight, inputWidth, cvImg->channels() });
	uint8_t *p = imgTensorWithSharedData.flat<uint8_t>().data();
	cv::Mat outputImg(inputHeight, inputWidth, CV_8UC3, p);
	cvImg->convertTo(outputImg, CV_8UC3);

	// Run tensorflow
	// cv::TickMeter tm;
	//tm.start();
	std::vector<tensorflow::Tensor> outputs;
	tensorflow::Status run_status = (*tf_session)->Run({ { "image_tensor:0", imgTensorWithSharedData } },
	{ "detection_boxes:0", "detection_scores:0", "detection_classes:0", "num_detections:0" },
	{},
		&outputs);
	if (!run_status.ok()) {
		std::cerr << "TF_Session->Run Error: " << run_status << std::endl;
	}
	// tm.stop();  
	// std::cout << "Inference time, ms: " << tm.getTimeMilli()  << std::endl;

	tensorflow::TTypes<float>::Flat scores = outputs[1].flat<float>();
	tensorflow::TTypes<float>::Flat classes = outputs[2].flat<float>();
	tensorflow::TTypes<float>::Flat num_detections = outputs[3].flat<float>();
	auto boxes = outputs[0].flat_outer_dims<float, 3>();

	int detectionsCount = (int)(num_detections(0));
	int drawnDetections = 0;
	cv::RNG rng(12345);
	std::cout << "Total detections before threshold: " << detectionsCount << std::endl;
	for (int i = 0; i < detectionsCount && i < 100000; ++i) { // 100000 is infinite loop protection
		if (scores(i) > predictionThreshold) {
			float boxClass = classes(i);

			float x1 = float(outputImg.size().width) * boxes(0, i, 1);
			float y1 = float(outputImg.size().height) * boxes(0, i, 0);

			float x2 = float(outputImg.size().width) * boxes(0, i, 3);
			float y2 = float(outputImg.size().height) * boxes(0, i, 2);

			std::ostringstream label;
			label << TF_LabelMap[boxClass] << ": " << (scores(i) * 100) << "%";
			std::cout << "Detection " << (i + 1) << ": class: " << boxClass << " " << label.str() << ", box: (" << x1 << "," << y1 << "), (" << x2 << "," << y2 << ")" << std::endl;
			cv::Scalar randomColor = cv::Scalar(rng.uniform(0, 255), rng.uniform(0, 255), rng.uniform(0, 255));
			cv::rectangle(*cvImg, cv::Point(x1, y1), cv::Point(x2, y2), randomColor);
			cv::putText(*cvImg, label.str(), cv::Point(x1, y1), 1, 1.0, randomColor);
			drawnDetections++;
		}
	}

	return drawnDetections;
}

int main(int argc, char* argv[]) 
{

	string image = ".\tensorflow_test\\image11.png";
	std::string TF_PB_PATH = ".\\tensorflow_test\\tiny-c3a5.pb";
	std::string TF_LABELLIST_PATH = ".\\tensorflow_test\\coco_label.txt";
	std::string input_layer = "image_tensor:0";
	std::vector<std::string> output_layer = { "detection_boxes:0", "detection_scores:0", "detection_classes:0", "num_detections:0" };

	if (!TF_init(TF_LABELLIST_PATH, &TF_LabelMap, TF_PB_PATH, &TF_Session)) {
		return EXIT_FAILURE;
	}
	VideoCapture cap;
	cap.open(0);
	cv::Mat Image;
	cap >> Image;
	cap >> Image;
	while (1)
	{
		double Time = (double)cvGetTickCount();
		cap >> Image;
		if (!Image.data) {
			std::cerr << "Could not open capture or find the image " << std::endl;
			return EXIT_FAILURE;
		}
		int detectionsCount = PredictionsAndDraw(&Image, &TF_Session, TF_PREDICTION_THRESSHOLD);
		std::cout << "Total detections: " << detectionsCount << std::endl;
		cv::imshow("Display window", Image);
		cv::waitKey(1);
		Time = (double)cvGetTickCount() - Time;
		printf("run time = %gs\n", Time / (cvGetTickFrequency() * 1000000));
	}
	return EXIT_SUCCESS;
}

猜你喜欢

转载自blog.csdn.net/shakevincent/article/details/80307386