commit 8c76e5d988560e18c5f73cdca99c1fb9a54ffd0c
author Filip Kujawa <filip.j.kujawa@gmail.com> Sat Apr 08 16:20:27 2023 -0700
committer Filip Kujawa <filip.j.kujawa@gmail.com> Sun Apr 09 19:24:06 2023 -0700
tree 03c1466d5199f69afe10e7600168de23df4088e5
parent 26a2366b566f2aa1df87858c5cc93bd9f90d178b

Actually call yolov5 inference

Signed-off-by: Filip Kujawa <filip.j.kujawa@gmail.com>
Change-Id: I693aa253d09e88f6000ee9ea5a8c68862ac47629

y2023/vision/game_pieces.cc

diff --git a/y2023/vision/game_pieces.cc b/y2023/vision/game_pieces.cc
index 0e2546f..68cc5b2 100644
--- a/y2023/vision/game_pieces.cc
+++ b/y2023/vision/game_pieces.cc
@@ -5,6 +5,8 @@
 #include "aos/events/event_loop.h"
 #include "aos/events/shm_event_loop.h"
 #include "frc971/vision/vision_generated.h"
+#include "y2023/vision/yolov5.h"
+#include <chrono>
 
 // The best_x and best_y are pixel (x, y) cordinates. The 'best'
 // game piece is picked on proximity to the specified cordinates.
@@ -24,18 +26,34 @@
 namespace vision {
 GamePiecesDetector::GamePiecesDetector(aos::EventLoop *event_loop)
     : game_pieces_sender_(event_loop->MakeSender<GamePieces>("/camera")) {
+    LOG(INFO) << "Before load model in constr";
+
+    model = MakeYOLOV5();
+    model->LoadModel("edgetpu_model.tflite");
+
+    LOG(INFO) << "After load model in constr";
   event_loop->MakeWatcher("/camera", [this](const CameraImage &camera_image) {
     this->ProcessImage(camera_image);
   });
 }
 
-// TODO(FILIP): Actually do inference.
-
 void GamePiecesDetector::ProcessImage(const CameraImage &image) {
-  // Param is not used for now.
-  (void)image;
+    auto start = std::chrono::high_resolution_clock::now();
+	LOG(INFO) << reinterpret_cast<const void*>(image.data()->data());
+  cv::Mat image_color_mat(cv::Size(image.cols(), image.rows()), CV_8UC2,
+                          (void *)image.data()->data());
+  std::vector<Detection> detections;
+  cv::Mat image_mat(cv::Size(image.cols(), image.rows()), CV_8UC3);
+  LOG(INFO) << reinterpret_cast<void*>(image_mat.ptr());
+  cv::cvtColor(image_color_mat, image_mat, cv::COLOR_YUV2BGR_YUYV);
+  LOG(INFO) << reinterpret_cast<void*>(image_mat.ptr());
 
-  const int detection_count = 5;
+  detections = model->ProcessImage(image_mat);
+  LOG(INFO) << reinterpret_cast<void*>(image_mat.ptr());
+  LOG(INFO) << reinterpret_cast<void*>(image_color_mat.ptr());
+
+  auto stop = std::chrono::high_resolution_clock::now();
+  LOG(INFO) << "INFERENCE TIME " << std::chrono::duration_cast<std::chrono::milliseconds>(stop - start).count();
 
   auto builder = game_pieces_sender_.MakeBuilder();
 
@@ -45,23 +63,30 @@
   int best_distance_index = 0;
   srand(time(0));
 
-  for (int i = 0; i < detection_count; i++) {
-    int h = rand() % 1000;
-    int w = rand() % 1000;
-    int x = rand() % 250;
-    int y = rand() % 250;
-
+  for (size_t i = 0; i < detections.size(); i++) {
     auto box_builder = builder.MakeBuilder<Box>();
-    box_builder.add_h(h);
-    box_builder.add_w(w);
-    box_builder.add_x(x);
-    box_builder.add_y(y);
+    box_builder.add_h(detections[i].box.height);
+    box_builder.add_w(detections[i].box.width);
+    box_builder.add_x(detections[i].box.x);
+    box_builder.add_y(detections[i].box.y);
     auto box_offset = box_builder.Finish();
 
     auto game_piece_builder = builder.MakeBuilder<GamePiece>();
-    game_piece_builder.add_piece_class(y2023::vision::Class::CONE_DOWN);
+    switch (detections[i].class_id) {
+      case 0:
+        game_piece_builder.add_piece_class(Class::CONE_DOWN);
+        break;
+      case 1:
+        game_piece_builder.add_piece_class(Class::CONE_UP);
+        break;
+      case 2:
+        game_piece_builder.add_piece_class(Class::CUBE);
+        break;
+      default:
+        game_piece_builder.add_piece_class(Class::CONE_DOWN);
+    }
     game_piece_builder.add_box(box_offset);
-    game_piece_builder.add_confidence(0.9);
+    game_piece_builder.add_confidence(detections[i].confidence);
     auto game_piece = game_piece_builder.Finish();
     game_pieces_offsets.push_back(game_piece);
 
@@ -69,8 +94,8 @@
     // Inference returns the top left corner of the bounding box
     // but we want the center of the box for this.
 
-    const int center_x = x + w / 2;
-    const int center_y = y + h / 2;
+    const int center_x = detections[i].box.x + detections[i].box.width / 2;
+    const int center_y = detections[i].box.y + detections[i].box.height / 2;
 
     // Find difference between target x, y and the x, y
     // of the bounding box using Euclidean distance.
@@ -85,8 +110,7 @@
     }
   };
 
-  flatbuffers::FlatBufferBuilder fbb;
-  auto game_pieces_vector = fbb.CreateVector(game_pieces_offsets);
+  auto game_pieces_vector = builder.fbb()->CreateVector(game_pieces_offsets);
 
   auto game_pieces_builder = builder.MakeBuilder<GamePieces>();
   game_pieces_builder.add_game_pieces(game_pieces_vector);
@@ -96,4 +120,4 @@
 }
 
 }  // namespace vision
-}  // namespace y2023
\ No newline at end of file
+}  // namespace y2023