@@ -298,17 +298,31 @@ std::vector<Contour> SegmentationModel::getContours(const ImageResultWithSoftPre
298298 cv::Scalar (index, index, index),
299299 label_index_map);
300300 std::vector<std::vector<cv::Point>> contours;
301- cv::findContours (label_index_map, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE);
301+ std::vector<cv::Vec4i> hierarchy;
302+ cv::findContours (label_index_map, contours, hierarchy, cv::RETR_CCOMP, cv::CHAIN_APPROX_NONE);
302303
303304 std::string label = getLabelName (index - 1 );
304305
305- for (unsigned int i = 0 ; i < contours.size (); i++) {
306+ for (size_t i = 0 ; i < contours.size (); ++i) {
307+ if (hierarchy[i][3 ] >= 0 ) {
308+ continue ;
309+ }
310+
306311 cv::Mat mask = cv::Mat::zeros (imageResult.resultImage .rows ,
307312 imageResult.resultImage .cols ,
308313 imageResult.resultImage .type ());
309314 cv::drawContours (mask, contours, i, 255 , -1 );
315+
316+ std::vector<std::vector<cv::Point>> children;
317+ int next_child_idx = hierarchy[i][2 ];
318+ while (next_child_idx >= 0 ) {
319+ children.push_back (contours[next_child_idx]);
320+ cv::drawContours (mask, contours, next_child_idx, 0 , -1 );
321+ next_child_idx = hierarchy[next_child_idx][0 ];
322+ }
323+
310324 float probability = (float )cv::mean (current_label_soft_prediction, mask)[0 ];
311- combined_contours.push_back ({label, probability, contours[i]});
325+ combined_contours.push_back ({label, probability, contours[i], children });
312326 }
313327 }
314328
0 commit comments