Add object detection inferencing code
Add a cc_library for inferencing using tflite and edgetpu.
Signed-off-by: Filip Kujawa <filip.j.kujawa@gmail.com>
Change-Id: Ie4cfa1e960a3d7461a75df074ebf12e4d47727e5
diff --git a/y2023/vision/BUILD b/y2023/vision/BUILD
index 0f75aba..0221aad 100644
--- a/y2023/vision/BUILD
+++ b/y2023/vision/BUILD
@@ -1,5 +1,6 @@
load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
load("@com_github_google_flatbuffers//:typescript.bzl", "flatbuffer_ts_library")
+load("//tools/build_rules:select.bzl", "cpu_select")
cc_binary(
name = "camera_reader",
@@ -267,3 +268,23 @@
"@com_github_google_glog//:glog",
],
)
+
+cc_library(
+ name = "yolov5_lib",
+ srcs = ["yolov5.cc"],
+ hdrs = ["yolov5.h"],
+ deps = [
+ "//third_party:opencv",
+ "@com_github_gflags_gflags//:gflags",
+ "@com_github_google_glog//:glog",
+ ] + cpu_select({
+ "amd64": [
+ "@libtensorflowlite//:tensorflow-k8",
+ "@libedgetpu//:libedgetpu-k8",
+ ],
+ "arm": [
+ "@libtensorflowlite//:tensorflow-arm",
+ "@libedgetpu//:libedgetpu-arm",
+ ],
+ }),
+)
diff --git a/y2023/vision/yolov5.cc b/y2023/vision/yolov5.cc
new file mode 100644
index 0000000..df03d12
--- /dev/null
+++ b/y2023/vision/yolov5.cc
@@ -0,0 +1,158 @@
+#include "yolov5.h"
+
+#include <opencv2/core.hpp>
+
+#include "gflags/gflags.h"
+#include "glog/logging.h"
+
+DEFINE_double(conf_threshold, 0.9,
+ "Threshold value for confidence scores. Detections with a "
+ "confidence score below this value will be ignored.");
+
+DEFINE_double(
+ nms_threshold, 0.5,
+ "Threshold value for non-maximum suppression. Detections with an "
+ "intersection-over-union value below this value will be removed.");
+
+DEFINE_int32(nthreads, 6, "Number of threads to use during inference.");
+
+namespace y2023 {
+namespace vision {
+
+void YOLOV5::LoadModel(const std::string path) {
+ model_ = tflite::FlatBufferModel::BuildFromFile(path.c_str());
+ CHECK(model_);
+ size_t num_devices;
+ std::unique_ptr<edgetpu_device, decltype(&edgetpu_free_devices)> devices(
+ edgetpu_list_devices(&num_devices), &edgetpu_free_devices);
+ const auto &device = devices.get()[0];
+ CHECK_EQ(num_devices, 1ul);
+ tflite::ops::builtin::BuiltinOpResolver resolver;
+ CHECK_EQ(tflite::InterpreterBuilder(*model_, resolver)(&interpreter_),
+ kTfLiteOk);
+
+ auto *delegate =
+ edgetpu_create_delegate(device.type, device.path, nullptr, 0);
+ interpreter_->ModifyGraphWithDelegate(delegate);
+
+ TfLiteStatus status = interpreter_->AllocateTensors();
+ CHECK(status == kTfLiteOk);
+
+ input_ = interpreter_->inputs()[0];
+ TfLiteIntArray *dims = interpreter_->tensor(input_)->dims;
+ in_height_ = dims->data[1];
+ in_width_ = dims->data[2];
+ in_channels_ = dims->data[3];
+ in_type_ = interpreter_->tensor(input_)->type;
+ input_8_ = interpreter_->typed_tensor<uint8_t>(input_);
+
+ interpreter_->SetNumThreads(FLAGS_nthreads);
+}
+
+void YOLOV5::Preprocess(cv::Mat image) {
+ cv::resize(image, image, cv::Size(in_height_, in_width_), cv::INTER_CUBIC);
+ cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
+ image.convertTo(image, CV_8U);
+}
+
+void YOLOV5::ConvertCVMatToTensor(const cv::Mat &src, uint8_t *in) {
+ CHECK(src.type() == CV_8UC3);
+ int n = 0, nc = src.channels(), ne = src.elemSize();
+ for (int y = 0; y < src.rows; ++y)
+ for (int x = 0; x < src.cols; ++x)
+ for (int c = 0; c < nc; ++c)
+ in[n++] = src.data[y * src.step + x * ne + c];
+}
+
+std::vector<std::vector<float>> YOLOV5::TensorToVector2D(
+ TfLiteTensor *src_tensor, const int rows, const int columns) {
+ auto scale = src_tensor->params.scale;
+ auto zero_point = src_tensor->params.zero_point;
+ std::vector<std::vector<float>> result_vec;
+ for (int32_t i = 0; i < rows; i++) {
+ std::vector<float> row_values;
+ for (int32_t j = 0; j < columns; j++) {
+ float val_float =
+ ((static_cast<int32_t>(src_tensor->data.uint8[i * columns + j])) -
+ zero_point) *
+ scale;
+ row_values.push_back(val_float);
+ }
+ result_vec.push_back(row_values);
+ }
+ return result_vec;
+}
+
+void YOLOV5::NonMaximumSupression(
+ const std::vector<std::vector<float>> &orig_preds, const int rows,
+ const int columns, std::vector<Detection> *detections,
+ std::vector<int> *indices)
+
+{
+ std::vector<float> scores;
+ double confidence;
+ cv::Point class_id;
+
+ for (int i = 0; i < rows; i++) {
+ if (orig_preds[i][4] > FLAGS_conf_threshold) {
+ int left = (orig_preds[i][0] - orig_preds[i][2] / 2) * img_width_;
+ int top = (orig_preds[i][1] - orig_preds[i][3] / 2) * img_height_;
+ int w = orig_preds[i][2] * img_width_;
+ int h = orig_preds[i][3] * img_height_;
+
+ for (int j = 5; j < columns; j++) {
+ scores.push_back(orig_preds[i][j] * orig_preds[i][4]);
+ }
+
+ cv::minMaxLoc(scores, nullptr, &confidence, nullptr, &class_id);
+ if (confidence > FLAGS_conf_threshold) {
+ Detection detection{cv::Rect(left, top, w, h), confidence,
+ class_id.x - kClassIdOffset};
+ detections->push_back(detection);
+ }
+ }
+ }
+
+ std::vector<cv::Rect> boxes;
+ std::vector<float> confidences;
+
+ for (const Detection &d : *detections) {
+ boxes.push_back(d.box);
+ confidences.push_back(d.confidence);
+ }
+
+ cv::dnn::NMSBoxes(boxes, confidences, FLAGS_conf_threshold,
+ FLAGS_nms_threshold, *indices);
+}
+
+std::vector<Detection> YOLOV5::ProcessImage(cv::Mat frame) {
+ img_height_ = frame.rows;
+ img_width_ = frame.cols;
+
+ Preprocess(frame);
+ ConvertCVMatToTensor(frame, input_8_);
+
+ // Inference
+ TfLiteStatus status = interpreter_->Invoke();
+ CHECK_EQ(status, kTfLiteOk);
+
+ int output_tensor_index = interpreter_->outputs()[0];
+ TfLiteIntArray *out_dims = interpreter_->tensor(output_tensor_index)->dims;
+ int num_rows = out_dims->data[1];
+ int num_columns = out_dims->data[2];
+
+ TfLiteTensor *src_tensor = interpreter_->tensor(interpreter_->outputs()[0]);
+ std::vector<std::vector<float>> orig_preds =
+ TensorToVector2D(src_tensor, num_rows, num_columns);
+
+ std::vector<int> indices;
+ std::vector<Detection> detections;
+
+ NonMaximumSupression(orig_preds, num_rows, num_columns, &detections,
+ &indices);
+
+ return detections;
+};
+
+} // namespace vision
+} // namespace y2023
diff --git a/y2023/vision/yolov5.h b/y2023/vision/yolov5.h
new file mode 100644
index 0000000..23cefd5
--- /dev/null
+++ b/y2023/vision/yolov5.h
@@ -0,0 +1,84 @@
+#ifndef Y2023_VISION_YOLOV5_H_
+#define Y2023_VISION_YOLOV5_H_
+
+#include <tensorflow/lite/interpreter.h>
+#include <tensorflow/lite/kernels/register.h>
+#include <tensorflow/lite/model.h>
+#include <tflite/public/edgetpu_c.h>
+
+#include <chrono>
+#include <cmath>
+#include <cstdint>
+#include <fstream>
+#include <iostream>
+#include <opencv2/core.hpp>
+#include <opencv2/dnn.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgcodecs.hpp>
+#include <opencv2/imgproc.hpp>
+#include <vector>
+
+namespace y2023 {
+namespace vision {
+
+struct Detection {
+ cv::Rect box;
+ double confidence;
+ int class_id;
+};
+
+class YOLOV5 {
+ public:
+ // Takes a model path as string and loads a pre-trained
+ // YOLOv5 model from the specified path.
+ void LoadModel(const std::string path);
+
+ // Takes an image and returns a Detection.
+ std::vector<Detection> ProcessImage(cv::Mat image);
+
+ private:
+ // Convert an OpenCV Mat object to a tensor input
+ // that can be fed to the TensorFlow Lite model.
+ void ConvertCVMatToTensor(const cv::Mat &src, uint8_t *in);
+
+ // Resizes, converts color space, and converts
+ // image data type before inference.
+ void Preprocess(cv::Mat image);
+
+ // Converts a TensorFlow Lite tensor to a 2D vector.
+ std::vector<std::vector<float>> TensorToVector2D(TfLiteTensor *src_tensor,
+ const int rows,
+ const int columns);
+
+ // Performs non-maximum suppression to remove overlapping bounding boxes.
+ void NonMaximumSupression(const std::vector<std::vector<float>> &orig_preds,
+ const int rows, const int columns,
+ std::vector<Detection> *detections,
+ std::vector<int> *indices);
+ // Models
+ std::unique_ptr<tflite::FlatBufferModel> model_;
+ std::unique_ptr<tflite::Interpreter> interpreter_;
+ tflite::StderrReporter error_reporter_;
+
+ // Parameters of interpreter's input
+ int input_;
+ int in_height_;
+ int in_width_;
+ int in_channels_;
+ int in_type_;
+
+ // Parameters of original image
+ int img_height_;
+ int img_width_;
+
+ // Input of the interpreter
+ uint8_t *input_8_;
+
+ // Subtract this offset from class labels to get the actual label.
+ static constexpr int kClassIdOffset = 5;
+};
+
+} // namespace vision
+} // namespace y2023
+
+#endif // Y2023_VISION_YOLOV5_H_