Final Refactoring

umrover · Sep 18, 2024 · 755f093 · 755f093
1 parent 60be043
commit 755f093
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 14 deletions.
diff --git a/perception/object_detector/object_detector.cpp b/perception/object_detector/object_detector.cpp
@@ -23,7 +23,7 @@ namespace mrover {
 
         RCLCPP_INFO_STREAM(get_logger(), "Opening Model " << mModelName);
 
-        mLearning = Learning{mModelName, packagePath};
+        mTensorRT = TensortRT{mModelName, packagePath};
 
         mDebugImgPub = create_publisher<sensor_msgs::msg::Image>("object_detector/debug_img", 1);
 

diff --git a/perception/object_detector/object_detector.hpp b/perception/object_detector/object_detector.hpp
@@ -20,7 +20,7 @@ namespace mrover {
 
         LoopProfiler mLoopProfiler;
 
-        Learning mLearning;
+        TensortRT mTensorRT;
 
         cv::Mat mRgbImage, mImageBlob;
         sensor_msgs::msg::Image mDetectionsImageMessage;

diff --git a/perception/object_detector/object_detector.processing.cpp b/perception/object_detector/object_detector.processing.cpp
@@ -26,7 +26,7 @@ namespace mrover {
 
         // Run the blob through the model
         std::vector<Detection> detections{};
-        mLearning.modelForwardPass(mImageBlob, detections, mModelScoreThreshold, mModelNmsThreshold);
+        mTensorRT.modelForwardPass(mImageBlob, detections, mModelScoreThreshold, mModelNmsThreshold);
 
         mLoopProfiler.measureEvent("Execution");
 

diff --git a/perception/object_detector/pch.hpp b/perception/object_detector/pch.hpp
@@ -49,4 +49,4 @@
 
 #include "parameter.hpp"
 #include "point.hpp"
-#include <learning.hpp>
+#include <tensorrt.hpp>
diff --git a/tensorrt/learning.cpp → tensorrt/tensorrt.cpp b/tensorrt/learning.cpp → tensorrt/tensorrt.cpp
@@ -1,10 +1,10 @@
-#include "learning.hpp"
+#include "tensorrt.hpp"
 
 using namespace std;
 
-Learning::Learning() = default;
+TensortRT::TensortRT() = default;
 
-Learning::Learning(string modelName, string& packagePathString) : mModelName{std::move(modelName)} {
+TensortRT::TensortRT(string modelName, string& packagePathString) : mModelName{std::move(modelName)} {
 
     std::filesystem::path packagePath = packagePathString;
     std::filesystem::path modelFileName = mModelName.append(".onnx");
@@ -13,15 +13,15 @@ Learning::Learning(string modelName, string& packagePathString) : mModelName{std
     mInferenceWrapper = InferenceWrapper{modelPath, mModelName, packagePath};
 }
 
-Learning::~Learning() = default;
+TensortRT::~TensortRT() = default;
 
-auto Learning::modelForwardPass(cv::Mat const& blob, std::vector<Detection>& detections, float modelScoreThreshold, float modelNMSThreshold) const -> void {
+auto TensortRT::modelForwardPass(cv::Mat const& blob, std::vector<Detection>& detections, float modelScoreThreshold, float modelNMSThreshold) const -> void {
     mInferenceWrapper.doDetections(blob);
     cv::Mat output = mInferenceWrapper.getOutputTensor();
     parseModelOutput(output, detections, modelScoreThreshold, modelNMSThreshold);
 }
 
-auto Learning::parseModelOutput(cv::Mat& output, std::vector<Detection>& detections, float modelScoreThreshold, float modelNMSThreshold) const -> void {
+auto TensortRT::parseModelOutput(cv::Mat& output, std::vector<Detection>& detections, float modelScoreThreshold, float modelNMSThreshold) const -> void {
     // Parse model specific dimensioning from the output
 
     // The input to this function is expecting a YOLOv8 style model, thus the dimensions should be > rows

diff --git a/tensorrt/learning.hpp → tensorrt/tensorrt.hpp b/tensorrt/learning.hpp → tensorrt/tensorrt.hpp
@@ -9,7 +9,7 @@ struct Detection {
     cv::Rect box;
 };
 
-class Learning {
+class TensortRT {
     std::string mModelName;
 
     std::vector<std::string> classes{"bottle", "hammer"};
@@ -22,11 +22,11 @@ class Learning {
                           float modelNMSThreshold = 0.5) const -> void;
 
 public:
-    Learning();
+    TensortRT();
 
-    explicit Learning(std::string modelName, std::string& packagePathString);
+    explicit TensortRT(std::string modelName, std::string& packagePathString);
 
-    ~Learning();
+    ~TensortRT();
 
     auto modelForwardPass(cv::Mat const& blob,
                           std::vector<Detection>& detections,