Add support of holes to contours generation

.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>> child_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.child_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,29 @@ 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;
+            }
+
+            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]);
+                next_child_idx = hierarchy[next_child_idx][0];
+            }
+
             cv::Mat mask = cv::Mat::zeros(imageResult.resultImage.rows,
                                           imageResult.resultImage.cols,
                                           imageResult.resultImage.type());
             cv::drawContours(mask, contours, i, 255, -1);
             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,18 @@ def rotated_rects(self, value):
 
 
 class Contour:
-    def __init__(self, label: str, probability: float, shape: list[tuple[int, int]]):
+    def __init__(self, label: str, probability: float, shape: np.ndarray, child_shapes: list[np.ndarray] | None = None):
         self.shape = shape
         self.label = label
         self.probability = probability
+        self.child_shapes = child_shapes
 
     def __str__(self):
-        return f"{self.label}: {self.probability:.3f}, {len(self.shape)}"
+        num_children = len(self.child_shapes) if self.child_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,13 +207,24 @@ 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
+
+                children = []
+                next_child_idx = hierarchy[i][2]
+                while next_child_idx >= 0:
+                    children.append(contours[next_child_idx])
+                    next_child_idx = hierarchy[next_child_idx][0]
 
-            for contour in contours:
                 mask = np.zeros(prediction.resultImage.shape, dtype=np.uint8)
                 cv2.drawContours(
                     mask,
@@ -223,7 +234,7 @@ def get_contours(
                     thickness=-1,
                 )
                 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

tests/python/accuracy/public_scope.json

@@ -285,7 +285,7 @@
       {
         "image": "coco128/images/train2017/000000000074.jpg",
         "reference": [
-          "0: 1.000, 1: 0.000, [426,640,3], [426,640,3], [1,600,1,1]; backpack: 0.505, 2, "
+          "0: 1.000, 1: 0.000, [426,640,3], [426,640,3], [1,600,1,1]; backpack: 0.505, 2, 0, "
         ]
       }
     ]
@@ -313,7 +313,7 @@
       {
         "image": "coco128/images/train2017/000000000074.jpg",
         "reference": [
-          "1535, 585, 1662, 697, 2 (ellipse): 0.643, 9822, RotatedRect: 1598.500 641.500 111.000 109.000 90.000; 3091, 3097, 3105, 3112, 1 (rectangle): 0.483, 197, RotatedRect: 3097.500 3104.000 14.000 13.000 90.000; 2734, 60, 2867, 324, 1 (rectangle): 0.401, 30622, RotatedRect: 2800.000 188.500 255.000 132.000 90.000; 2; [1,1280,1,1]; ellipse: 0.643, 331; rectangle: 0.483, 48; rectangle: 0.401, 758; "
+          "1535, 585, 1662, 697, 2 (ellipse): 0.643, 9822, RotatedRect: 1598.500 641.500 111.000 109.000 90.000; 3091, 3097, 3105, 3112, 1 (rectangle): 0.483, 197, RotatedRect: 3097.500 3104.000 14.000 13.000 90.000; 2734, 60, 2867, 324, 1 (rectangle): 0.401, 30622, RotatedRect: 2800.000 188.500 255.000 132.000 90.000; 2; [1,1280,1,1]; ellipse: 0.643, 331, 0; rectangle: 0.483, 48, 0; rectangle: 0.401, 758, 0; "
         ]
       }
     ]
@@ -427,7 +427,7 @@
       {
         "image": "coco128/images/train2017/000000000074.jpg",
         "reference": [
-          "0: 0.272, 1: 0.728, [3500,3500,5], [0], [0]; background: 1.404, 311, background: 1.397, 44, background: 1.371, 34, background: 1.377, 12, background: 1.356, 155, background: 1.345, 12, background: 1.183, 219, background: 1.524, 8, background: 1.533, 4, background: 1.519, 2, background: 1.524, 4, background: 1.530, 6, background: 1.537, 2, background: 1.514, 4, background: 1.519, 8, background: 1.529, 6, background: 1.550, 6, background: 1.558, 4, background: 1.520, 2, background: 1.529, 4, background: 1.532, 6, background: 1.535, 6, background: 1.530, 2, background: 1.529, 50, background: 1.528, 22, background: 1.527, 38, background: 1.451, 1476, background: 1.345, 2743, background: 1.609, 2987, background: 1.636, 29909, "
+          "0: 0.272, 1: 0.728, [3500,3500,5], [0], [0]; background: 1.404, 311, 0, background: 1.397, 44, 0, background: 1.371, 34, 0, background: 1.377, 12, 0, background: 1.356, 155, 0, background: 1.345, 12, 0, background: 1.183, 219, 0, background: 1.524, 8, 0, background: 1.533, 4, 0, background: 1.519, 2, 0, background: 1.524, 4, 0, background: 1.530, 6, 0, background: 1.537, 2, 0, background: 1.514, 4, 0, background: 1.519, 8, 0, background: 1.529, 6, 0, background: 1.550, 6, 0, background: 1.558, 4, 0, background: 1.520, 2, 0, background: 1.529, 4, 0, background: 1.532, 6, 0, background: 1.535, 6, 0, background: 1.530, 2, 0, background: 1.529, 50, 0, background: 1.528, 22, 0, background: 1.527, 38, 0, background: 1.451, 1476, 0, background: 1.345, 2743, 9, background: 1.609, 2987, 1, background: 1.398, 20, 0, background: 1.636, 29909, 8, "
         ]
       }
     ]