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;
}