implement changes

Author: aljazkonec1

bindings/python/src/pipeline/datatype/ImgDetectionsBindings.cpp

@@ -118,19 +118,34 @@ void bind_imgdetections(pybind11::module& m, void* pCallstack) {
             [](ImgDetections& det, std::vector<ImgDetection> val) { det.detections = val; },
             DOC(dai, ImgDetectionsT, detections),
             py::return_value_policy::reference_internal)
+        .def_property(
+            "segmentationMaskWidth",
+            [](ImgDetections& det) { return &det.segmentationMaskWidth; },
+            [](ImgDetections& det, size_t val) { det.segmentationMaskWidth = val; },
+            DOC(dai, ImgDetectionsT, segmentationMaskWidth),
+            py::return_value_policy::reference_internal)
+        .def_property(
+            "segmentationMaskHeight",
+            [](ImgDetections& det) { return &det.segmentationMaskHeight; },
+            [](ImgDetections& det, size_t val) { det.segmentationMaskHeight = val; },
+            DOC(dai, ImgDetectionsT, segmentationMaskHeight),
+            py::return_value_policy::reference_internal)
         .def("getTimestamp", &dai::ImgDetectionsT<dai::ImgDetection>::Buffer::getTimestamp, DOC(dai, Buffer, getTimestamp))
         .def("getTimestampDevice", &dai::ImgDetectionsT<dai::ImgDetection>::Buffer::getTimestampDevice, DOC(dai, Buffer, getTimestampDevice))
         .def("getSequenceNum", &dai::ImgDetectionsT<dai::ImgDetection>::Buffer::getSequenceNum, DOC(dai, Buffer, getSequenceNum))
         .def("getTransformation", [](ImgDetections& msg) { return msg.transformation; })
         .def("setTransformation", [](ImgDetections& msg, const std::optional<ImgTransformation>& transformation) { msg.transformation = transformation; })
         .def("getSegmentationMaskWidth", &ImgDetections::getSegmentationMaskWidth, DOC(dai, ImgDetectionsT, getSegmentationMaskWidth))
         .def("getSegmentationMaskHeight", &ImgDetections::getSegmentationMaskHeight, DOC(dai, ImgDetectionsT, getSegmentationMaskHeight))
-        .def(
-            "setMask", &ImgDetections::setSegmentationMask, py::arg("mask"), py::arg("width"), py::arg("height"), DOC(dai, ImgDetectionsT, setSegmentationMask))
+        .def("setSegmentationMask",
+             py::overload_cast<dai::ImgFrame&>(&ImgDetections::setSegmentationMask),
+             py::arg("frame"),
+             DOC(dai, ImgDetectionsT, setSegmentationMask),
+             py::return_value_policy::reference_internal)
         .def("getMaskData", &ImgDetections::getMaskData, DOC(dai, ImgDetectionsT, getMaskData))
-        .def("getSegmentationMaskAsImgFrame", &ImgDetections::getSegmentationMask, DOC(dai, ImgDetectionsT, getSegmentationMask))
+        .def("getSegmentationMask", &ImgDetections::getSegmentationMask, DOC(dai, ImgDetectionsT, getSegmentationMask))
 #ifdef DEPTHAI_HAVE_OPENCV_SUPPORT
-        .def("setSegmentationMask", &ImgDetections::setCvSegmentationMask, py::arg("mask"), DOC(dai, ImgDetectionsT, setCvSegmentationMask))
+        .def("setCvSegmentationMask", &ImgDetections::setCvSegmentationMask, py::arg("mask"), DOC(dai, ImgDetectionsT, setCvSegmentationMask))
         .def(
             "getCvSegmentationMask",
             [](ImgDetections& self) { return self.getCvSegmentationMask(&g_numpyAllocator); },

include/depthai/pipeline/datatype/ImgDetections.hpp

@@ -139,7 +139,9 @@ struct ImgDetection {
 };
 
 /**
- * ImgDetections message. Carries normalized detection results
+ * ImgDetections message. Carries normalized detections and optional segmentation mask.
+ * The segmentation mask is stored as a single-channel INT8 2-d array, where the value represents the instance index in the list of detections.
+ * The value 255 is treated as a background pixel (no instance).
  */
 class ImgDetections : public ImgDetectionsT<ImgDetection>, public ProtoSerializable {
    public:

include/depthai/pipeline/datatype/ImgDetectionsT.hpp

@@ -45,16 +45,25 @@ class ImgDetectionsT : public Buffer {
     std::size_t getSegmentationMaskHeight() const;
 
     /*
-     * Sets the segmentation mask from a vector of bytes, along with width and height.
+     * Sets the segmentation mask from a vector of bytes.
      * The size of the vector must be equal to width * height.
      */
     void setSegmentationMask(const std::vector<std::uint8_t>& mask, size_t width, size_t height);
 
+    /*
+     * Sets the segmentation mask from an ImgFrame.
+     * @param frame Frame must be of type GRAY8
+     */
+    void setSegmentationMask(dai::ImgFrame& frame);
+
     /*
      * Returns a copy of the segmentation mask data as a vector of bytes. If mask data is not set, returns std::nullopt.
      */
     std::optional<std::vector<std::uint8_t>> getMaskData() const;
 
+    /*
+     * Returns the segmentation mask as an ImgFrame. If mask data is not set, returns std::nullopt.
+     */
     std::optional<dai::ImgFrame> getSegmentationMask() const;
 
 // Optional - OpenCV support
@@ -71,7 +80,7 @@ class ImgDetectionsT : public Buffer {
     void setCvSegmentationMask(cv::Mat mask);
 
     /**
-     * Retrieves data as cv::Mat with specified width and height. If mask data is not set, returns std::nullopt.
+     * Retrieves mask data as a cv::Mat copy with specified width and height. If mask data is not set, returns std::nullopt.
      * @param allocator Allows callers to supply a custom cv::MatAllocator for zero-copy/custom memory management; nullptr uses OpenCV’s default.
      */
     std::optional<cv::Mat> getCvSegmentationMask(cv::MatAllocator* allocator = nullptr);

src/pipeline/datatype/ImgDetectionsT.cpp

@@ -31,6 +31,18 @@ void ImgDetectionsT<DetectionT>::setSegmentationMask(const std::vector<std::uint
     this->segmentationMaskHeight = height;
 }
 
+template <class DetectionT>
+void ImgDetectionsT<DetectionT>::setSegmentationMask(dai::ImgFrame& frame) {
+    if(frame.getType() != dai::ImgFrame::Type::GRAY8) {
+        throw std::runtime_error("SegmentationMask: ImgFrame type must be GRAY8");
+    }
+    auto dataSpan = frame.getData();
+    std::vector<std::uint8_t> vecMask(dataSpan.begin(), dataSpan.end());
+    setData(vecMask);
+    this->segmentationMaskWidth = frame.getWidth();
+    this->segmentationMaskHeight = frame.getHeight();
+}
+
 template <class DetectionT>
 std::optional<std::vector<std::uint8_t>> ImgDetectionsT<DetectionT>::getMaskData() const {
     const auto& d = data->getData();

tests/src/onhost_tests/pipeline/datatype/imgdetections_test.cpp

@@ -274,7 +274,7 @@ TEST_CASE("ImgDetections segmentation mask operations", "[ImgDetections][Segment
     }
 
 #ifdef DEPTHAI_HAVE_OPENCV_SUPPORT
-    SECTION("OpenCV segmentation mask view semantics") {
+    SECTION("OpenCV segmentation mask copy semantics") {
         ImgDetections detections;
         constexpr int rows = 3;
         constexpr int cols = 4;
@@ -303,7 +303,8 @@ TEST_CASE("ImgDetections segmentation mask operations", "[ImgDetections][Segment
         auto shallowData = *optShallowData;
 
         REQUIRE_FALSE(shallowData.empty());
-        REQUIRE(shallowData.front() == shallow.at<uint8_t>(0, 0));
+        REQUIRE(shallowData.front() == mask.at<uint8_t>(0, 0));
+        REQUIRE(shallowData.front() != shallow.at<uint8_t>(0, 0));
 
         cv::Mat constantMask(rows, cols, CV_8UC1, cv::Scalar(7));
         detections.setCvSegmentationMask(constantMask);