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/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_