Add support of holes to contours generation

Author: sovrasov

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

@@ -149,13 +149,20 @@ 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)}"
+        repr = f"{self.label}: {self.probability:.3f}, {len(self.shape)}"
+        if self.child_shapes is not None:
+            repr += f", {len(self.child_shapes)}"
+        return repr
+
+    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,25 @@ 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,
             )
+            hierarchy = hierarchy.squeeze()
 
-            for contour in contours:
+            for i, contour in enumerate(contours):
+                children = None
+
+                if hierarchy[i][3] >= 0:
+                    continue
+
+                children = []
+                if hierarchy[i][2] >= 0:
+                    child_next_idx = hierarchy[i][2]
+                    while child_next_idx >= 0:
+                        children.append(contours[child_next_idx])
+                        child_next_idx = hierarchy[child_next_idx][0]
                 mask = np.zeros(prediction.resultImage.shape, dtype=np.uint8)
                 cv2.drawContours(
                     mask,
@@ -223,7 +235,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