Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/cpp/models/include/models/segmentation_model.h
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,8 @@ class SegmentationModel : public BaseModel {
virtual std::vector<std::unique_ptr<ImageResult>> inferBatch(const std::vector<ImageInputData>& inputImgs);

static std::string ModelType;
std::vector<Contour> getContours(const ImageResultWithSoftPrediction& imageResult);
std::vector<Contour> getContours(const ImageResultWithSoftPrediction& imageResult,
bool simplified_postprocessing = false);

protected:
void prepareInputsOutputs(std::shared_ptr<ov::Model>& model) override;
Expand Down
6 changes: 4 additions & 2 deletions src/cpp/models/src/segmentation_model.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -283,14 +283,16 @@ std::unique_ptr<ResultBase> SegmentationModel::postprocess(InferenceResult& infR
return std::unique_ptr<ResultBase>(result);
}

std::vector<Contour> SegmentationModel::getContours(const ImageResultWithSoftPrediction& imageResult) {
std::vector<Contour> SegmentationModel::getContours(const ImageResultWithSoftPrediction& imageResult,
bool simplified_postprocessing) {
if (imageResult.soft_prediction.channels() == 1) {
throw std::runtime_error{"Cannot get contours from soft prediction with 1 layer"};
}

std::vector<Contour> combined_contours = {};
cv::Mat label_index_map;
cv::Mat current_label_soft_prediction;
int find_contours_mode = simplified_postprocessing ? cv::RETR_EXTERNAL : cv::RETR_CCOMP;
for (int index = 1; index < imageResult.soft_prediction.channels(); index++) {
cv::extractChannel(imageResult.soft_prediction, current_label_soft_prediction, index);
cv::inRange(imageResult.resultImage,
Expand All @@ -299,7 +301,7 @@ std::vector<Contour> SegmentationModel::getContours(const ImageResultWithSoftPre
label_index_map);
std::vector<std::vector<cv::Point>> contours;
std::vector<cv::Vec4i> hierarchy;
cv::findContours(label_index_map, contours, hierarchy, cv::RETR_CCOMP, cv::CHAIN_APPROX_NONE);
cv::findContours(label_index_map, contours, hierarchy, find_contours_mode, cv::CHAIN_APPROX_NONE);

std::string label = getLabelName(index - 1);

Expand Down
14 changes: 13 additions & 1 deletion src/python/model_api/models/segmentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,12 +186,24 @@ def postprocess(self, outputs: dict, meta: dict) -> ImageResultWithSoftPredictio
def get_contours(
self,
prediction: ImageResultWithSoftPrediction,
simplified_postprocessing: bool = False,
) -> list[Contour]:
"""Represents existing masks with contours.

Args:
prediction (ImageResultWithSoftPrediction): Input segmentation prediction.
simplified_postprocessing (bool, optional): Disables searching for holes in masks. Defaults to False.

Returns:
list[Contour]: Contours found.
"""
n_layers = prediction.soft_prediction.shape[2]

if n_layers == 1:
msg = "Cannot get contours from soft prediction with 1 layer"
raise RuntimeError(msg)

find_contours_mode = cv2.RETR_CCOMP if not simplified_postprocessing else cv2.RETR_EXTERNAL
combined_contours = []
for layer_index in range(1, n_layers): # ignoring background
label = self.get_label_name(layer_index - 1)
Expand All @@ -209,7 +221,7 @@ def get_contours(

contours, hierarchy = cv2.findContours(
label_index_map,
cv2.RETR_CCOMP,
find_contours_mode,
cv2.CHAIN_APPROX_NONE,
)
if len(contours):
Expand Down