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