Works on zed

umrover · Sep 26, 2024 · 816bed3 · 816bed3
1 parent dc40933
commit 816bed3
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 6 deletions.
diff --git a/perception/object_detector/object_detector.hpp b/perception/object_detector/object_detector.hpp
@@ -18,17 +18,18 @@ namespace mrover {
 
 			std::vector<int64_t> outputTensorSize;
 
+			// Additional space for params that one algorithm might need and not others
 			std::vector<float> buffer;
 
 			// Converts an rgb image to a blob
-			std::function<void(Model const&, cv::Mat&, cv::Mat&)> rbgImageToBlob;
+			std::function<void(Model const&, cv::Mat&, cv::Mat&, cv::Mat&)> rbgImageToBlob;
 
 			// Converts an output tensor into a vector of detections
 			std::function<void(Model const&, cv::Mat&, std::vector<Detection>&)> outputTensorToDetections;
 
 			Model() = default;
 
-			Model(std::string _modelName, std::vector<int> _objectHitCounts, std::vector<std::string> _classes, std::vector<int64_t> _inputTensorSize, std::vector<int64_t> _outputTensorSize, std::vector<float> _buffer, std::function<void(Model const&, cv::Mat&, cv::Mat&)> _rbgImageToBlob, std::function<void(Model const&, cv::Mat&, std::vector<Detection>&)> _outputTensorToDetections) : modelName{std::move(_modelName)}, 	objectHitCounts(std::move(_objectHitCounts)), classes(std::move(_classes)), inputTensorSize(std::move(_inputTensorSize)), outputTensorSize(std::move(_outputTensorSize)), buffer(std::move(_buffer)), rbgImageToBlob{std::move(_rbgImageToBlob)}, outputTensorToDetections{std::move(_outputTensorToDetections)} {}
+			Model(std::string _modelName, std::vector<int> _objectHitCounts, std::vector<std::string> _classes, std::vector<int64_t> _inputTensorSize, std::vector<int64_t> _outputTensorSize, std::vector<float> _buffer, std::function<void(Model const&, cv::Mat&, cv::Mat&, cv::Mat&)> _rbgImageToBlob, std::function<void(Model const&, cv::Mat&, std::vector<Detection>&)> _outputTensorToDetections) : modelName{std::move(_modelName)}, 	objectHitCounts(std::move(_objectHitCounts)), classes(std::move(_classes)), inputTensorSize(std::move(_inputTensorSize)), outputTensorSize(std::move(_outputTensorSize)), buffer(std::move(_buffer)), rbgImageToBlob{std::move(_rbgImageToBlob)}, outputTensorToDetections{std::move(_outputTensorToDetections)} {}
 		};
 
         static constexpr char const* NODE_NAME = "object_detector";
@@ -74,7 +75,7 @@ namespace mrover {
 								      cv::Mat& output,
 							   		  std::vector<Detection>& detections) -> void;
 
-		auto static preprocessYOLOv8Input(Model const& model, cv::Mat& rgbImage, cv::Mat& blob) -> void;
+		auto static preprocessYOLOv8Input(Model const& model, cv::Mat& rgbImage, cv::Mat& blobSizedImage, cv::Mat& blob) -> void;
 
     public:
         explicit ObjectDetectorBase();

diff --git a/perception/object_detector/object_detector.parsing.cpp b/perception/object_detector/object_detector.parsing.cpp
@@ -1,11 +1,14 @@
 #include "object_detector.hpp"
 namespace mrover {
-	auto ObjectDetectorBase::preprocessYOLOv8Input(Model const& model, cv::Mat& rgbImage, cv::Mat& blob) -> void{
-		cv::Mat blobSizedImage;
+	auto ObjectDetectorBase::preprocessYOLOv8Input(Model const& model, cv::Mat& rgbImage, cv::Mat& blobSizedImage, cv::Mat& blob) -> void{
 		if (model.inputTensorSize.size() != 4) {
 			throw std::runtime_error("Expected Blob Size to be of size 4, are you using the correct model type?");
 		}
 
+		if (model.buffer.size() != 2){
+			throw std::runtime_error("Expected 2 additional parameters!");
+		}
+
 		static constexpr double UCHAR_TO_DOUBLE = 1.0 / 255.0;
 
 		cv::Size blobSize{static_cast<int32_t>(model.inputTensorSize[2]), static_cast<int32_t>(model.inputTensorSize[3])};

diff --git a/perception/object_detector/object_detector.processing.cpp b/perception/object_detector/object_detector.processing.cpp
@@ -24,7 +24,7 @@ namespace mrover {
 
         // Convert the RGB Image into the blob Image format
         cv::Mat blobSizedImage;
-		mModel.rbgImageToBlob(mModel, mRgbImage, blobSizedImage);
+		mModel.rbgImageToBlob(mModel, mRgbImage, blobSizedImage, mImageBlob);
 
         mLoopProfiler.measureEvent("Conversion");
 

diff --git a/tensorrt/inference_wrapper.cpp b/tensorrt/inference_wrapper.cpp
@@ -28,3 +28,7 @@ auto InferenceWrapper::getOutputTensor() const -> cv::Mat {
 auto InferenceWrapper::getInputTensorSize() -> std::vector<int64_t>{
 	return mInference->getInputTensorSize();
 }
+
+auto InferenceWrapper::getOutputTensorSize() -> std::vector<int64_t>{
+	return mInference->getInputTensorSize();
+}