Remove hardcoded COCO labels size

ajtudela · Dec 13, 2022 · 0af1c42 · 0af1c42
1 parent 4441d8a
commit 0af1c42
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 7 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -6,6 +6,7 @@ Changelog for package object_detection_openvino
 ------------------
 * First ROS2 (Galactic) version.
 * Update to OpenVino 2021.4.
+* Remove hardcoded COCO labels size.
 
 3.2.2 (26-05-2022)
 ------------------

diff --git a/include/object_detection_openvino/objectDetectionVPU.hpp b/include/object_detection_openvino/objectDetectionVPU.hpp
@@ -44,9 +44,6 @@
 // OpenCV
 #include <cv_bridge/cv_bridge.h>
 
-/// The number of the labels in the COCO database.
-#define COCO_CLASSES		80
-
 /// Typedef for easier readability.
 typedef std::chrono::duration<double, std::ratio<1, 1000>> ms;
 

diff --git a/src/objectDetectionVPU.cpp b/src/objectDetectionVPU.cpp
@@ -289,11 +289,11 @@ void ObjectDetectionVPU::color_point_callback(
 			object.ymax = object.ymax > color_height ? color_height : object.ymax;
 
 			// Color of the class
-			int offset = object.classId * 123457 % COCO_CLASSES;
+			int offset = object.classId * 123457 % labels_.size();
 			float color_rgb[3];
-			color_rgb[0] = get_color(2, offset, COCO_CLASSES);
-			color_rgb[1] = get_color(1, offset, COCO_CLASSES);
-			color_rgb[2] = get_color(0, offset, COCO_CLASSES);
+			color_rgb[0] = get_color(2, offset, labels_.size());
+			color_rgb[1] = get_color(1, offset, labels_.size());
+			color_rgb[2] = get_color(0, offset, labels_.size());
 
 			// Create detection2D and push to array
 			vision_msgs::msg::Detection2D detection_2d;