Add support of holes to contours generation (#335)

* Add support of holes to contours generation

* Fix logic in case of absence of nested contours

* Cleanup python impl

* Add cpp implementation

* Unify contour str reprs

* Fix np warning

* Resolve some of corner-cases

* Update refs

* Update pre-commit config

* Fix numpy issue

* Update ref tests results

* Retain old data type for Contour init args

* Exclude nested contours from contour confidence calculation

* Exclude nested contours from prob computation

* Rename child_shapes for clarity

* Unify calls of drawContours

* Tests debug

* Update ref results

* Add numpy conversion to shape

* Fix list of contours conversion

Author: sovrasov

.github/workflows/test_accuracy.yml

@@ -35,7 +35,7 @@ jobs:
       - name: Run Python Test
         run: |
           source venv/bin/activate
-          pytest --data=./data tests/python/accuracy/test_accuracy.py
+          pytest -v --data=./data tests/python/accuracy/test_accuracy.py
       - name: Install CPP dependencies
         run: |
           sudo bash src/cpp/install_dependencies.sh

.pre-commit-config.yaml

@@ -30,8 +30,8 @@ repos:
       - id: mypy
         additional_dependencies: [types-PyYAML, types-setuptools]
 
-  - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v4.0.0-alpha.8
+  - repo: https://github.com/rbubley/mirrors-prettier
+    rev: v3.6.2
     hooks:
       - id: prettier
 

src/cpp/models/include/models/results.h

@@ -315,10 +315,11 @@ struct Contour {
     std::string label;
     float probability;
     std::vector<cv::Point> shape;
+    std::vector<std::vector<cv::Point>> excluded_shapes;
 
     friend std::ostream& operator<<(std::ostream& os, const Contour& contour) {
         return os << contour.label << ": " << std::fixed << std::setprecision(3) << contour.probability << ", "
-                  << contour.shape.size();
+                  << contour.shape.size() << ", " << contour.excluded_shapes.size();
     }
 };
 
@@ -332,7 +333,7 @@ static inline std::vector<Contour> getContours(const std::vector<SegmentedObject
         if (contours.size() != 1) {
             throw std::runtime_error("findContours() must have returned only one contour");
         }
-        combined_contours.push_back({obj.label, obj.confidence, contours[0]});
+        combined_contours.push_back({obj.label, obj.confidence, contours[0], {}});
     }
     return combined_contours;
 }

src/cpp/models/src/segmentation_model.cpp

@@ -298,17 +298,31 @@ std::vector<Contour> SegmentationModel::getContours(const ImageResultWithSoftPre
                     cv::Scalar(index, index, index),
                     label_index_map);
         std::vector<std::vector<cv::Point>> contours;
-        cv::findContours(label_index_map, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE);
+        std::vector<cv::Vec4i> hierarchy;
+        cv::findContours(label_index_map, contours, hierarchy, cv::RETR_CCOMP, cv::CHAIN_APPROX_NONE);
 
         std::string label = getLabelName(index - 1);
 
-        for (unsigned int i = 0; i < contours.size(); i++) {
+        for (size_t i = 0; i < contours.size(); ++i) {
+            if (hierarchy[i][3] >= 0) {
+                continue;
+            }
+
             cv::Mat mask = cv::Mat::zeros(imageResult.resultImage.rows,
                                           imageResult.resultImage.cols,
                                           imageResult.resultImage.type());
             cv::drawContours(mask, contours, i, 255, -1);
+
+            std::vector<std::vector<cv::Point>> children;
+            int next_child_idx = hierarchy[i][2];
+            while (next_child_idx >= 0) {
+                children.push_back(contours[next_child_idx]);
+                cv::drawContours(mask, contours, next_child_idx, 0, -1);
+                next_child_idx = hierarchy[next_child_idx][0];
+            }
+
             float probability = (float)cv::mean(current_label_soft_prediction, mask)[0];
-            combined_contours.push_back({label, probability, contours[i]});
+            combined_contours.push_back({label, probability, contours[i], children});
         }
     }
 

src/python/model_api/models/result/segmentation.py

@@ -149,13 +149,35 @@ def rotated_rects(self, value):
 
 
 class Contour:
-    def __init__(self, label: str, probability: float, shape: list[tuple[int, int]]):
-        self.shape = shape
+    """Represents a semantic segmentation mask as internals of a contour with "holes".
+    Args:
+        label (str): The label of the contour.
+        probability (float): The probability associated with the contour.
+        shape (np.ndarray | list[tuple[int, int]]): The shape of the contour. Shape is represented as a
+            list of 2d points or an equivalent numpy array (N, 2).
+        excluded_shapes (list[np.ndarray] | list[tuple[int, int]] | None, optional): Shapes of excluded contours.
+            If empty, the main shape is simply connected. Otherwise, excluded_shapes
+            represent "holes". Defaults to None.
+    """
+
+    def __init__(
+        self,
+        label: str,
+        probability: float,
+        shape: np.ndarray | list[tuple[int, int]],
+        excluded_shapes: list[np.ndarray] | list[list[tuple[int, int]]] | None = None,
+    ):
+        self.shape = np.array(shape)
         self.label = label
         self.probability = probability
+        self.excluded_shapes = [np.array(x) for x in excluded_shapes] if excluded_shapes is not None else None
 
     def __str__(self):
-        return f"{self.label}: {self.probability:.3f}, {len(self.shape)}"
+        num_children = len(self.excluded_shapes) if self.excluded_shapes is not None else 0
+        return f"{self.label}: {self.probability:.3f}, {len(self.shape)}, {num_children}"
+
+    def __repr__(self):
+        return self.__str__()
 
 
 class ImageResultWithSoftPrediction(Result):

src/python/model_api/models/segmentation.py

@@ -186,7 +186,7 @@ def postprocess(self, outputs: dict, meta: dict) -> ImageResultWithSoftPredictio
     def get_contours(
         self,
         prediction: ImageResultWithSoftPrediction,
-    ) -> list:
+    ) -> list[Contour]:
         n_layers = prediction.soft_prediction.shape[2]
 
         if n_layers == 1:
@@ -207,23 +207,30 @@ def get_contours(
             obj_group = prediction.resultImage == layer_index
             label_index_map = obj_group.astype(np.uint8) * 255
 
-            contours, _hierarchy = cv2.findContours(
+            contours, hierarchy = cv2.findContours(
                 label_index_map,
-                cv2.RETR_EXTERNAL,
+                cv2.RETR_CCOMP,
                 cv2.CHAIN_APPROX_NONE,
             )
+            if len(contours):
+                hierarchy = hierarchy.squeeze(axis=0)
+
+            for i, contour in enumerate(contours):
+                if hierarchy[i][3] >= 0:
+                    continue
 
-            for contour in contours:
                 mask = np.zeros(prediction.resultImage.shape, dtype=np.uint8)
-                cv2.drawContours(
-                    mask,
-                    np.asarray([contour]),
-                    contourIdx=-1,
-                    color=1,
-                    thickness=-1,
-                )
+                cv2.drawContours(mask, contours, contourIdx=i, color=1, thickness=-1)
+
+                children = []
+                next_child_idx = hierarchy[i][2]
+                while next_child_idx >= 0:
+                    children.append(contours[next_child_idx])
+                    cv2.drawContours(mask, contours, contourIdx=next_child_idx, color=0, thickness=-1)
+                    next_child_idx = hierarchy[next_child_idx][0]
+
                 probability = cv2.mean(current_label_soft_prediction, mask)[0]
-                combined_contours.append(Contour(label, probability, contour))
+                combined_contours.append(Contour(label, probability, contour, children))
 
         return combined_contours
 

src/python/model_api/tilers/instance_segmentation.py

@@ -123,7 +123,7 @@ def _merge_results(self, results, shape) -> InstanceSegmentationResult:
         labels = labels.astype(np.int32)
         resized_masks, label_names = [], []
         for mask, box, label_idx in zip(masks, bboxes, labels):
-            label_names.append(self.model.labels[int(label_idx)])
+            label_names.append(self.model.labels[int(label_idx.squeeze())])
             resized_masks.append(_segm_postprocess(box, mask, *shape[:-1]))
 
         resized_masks = np.stack(resized_masks) if resized_masks else masks