From 7d3940c293d9397e73ccc293896f4fafcd209f89 Mon Sep 17 00:00:00 2001
From: Vladislav Sovrasov <vladislav.sovrasov@intel.com>
Date: Thu, 26 Sep 2024 06:25:10 +0900
Subject: [PATCH] Add checks of the input entities in cpp sseg tiler

---
 .../cpp/tilers/src/semantic_segmentation.cpp  | 20 +++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/model_api/cpp/tilers/src/semantic_segmentation.cpp b/model_api/cpp/tilers/src/semantic_segmentation.cpp
index a514c8c5..c9bed98c 100644
--- a/model_api/cpp/tilers/src/semantic_segmentation.cpp
+++ b/model_api/cpp/tilers/src/semantic_segmentation.cpp
@@ -65,16 +65,17 @@ std::unique_ptr<ImageResultWithSoftPrediction> SemanticSegmentationTiler::run(co
 }
 
 std::unique_ptr<ResultBase> SemanticSegmentationTiler::postprocess_tile(std::unique_ptr<ResultBase> tile_result, const cv::Rect&) {
+    ImageResultWithSoftPrediction* soft = dynamic_cast<ImageResultWithSoftPrediction*>(tile_result.get());
+    if (!soft) {
+        throw std::runtime_error("SemanticSegmentationTiler requires the underlying model to return ImageResultWithSoftPrediction");
+    }
     return tile_result;
 }
 
 std::unique_ptr<ResultBase> SemanticSegmentationTiler::merge_results(const std::vector<std::unique_ptr<ResultBase>>& tiles_results,
                                                                      const cv::Size& image_size, const std::vector<cv::Rect>& tile_coords) {
-    auto* result = new ImageResultWithSoftPrediction();
-    auto retVal = std::unique_ptr<ResultBase>(result);
-
     if (tiles_results.empty()) {
-        return retVal;
+        return std::unique_ptr<ResultBase>(new ImageResultWithSoftPrediction());
     }
 
     cv::Mat voting_mask(cv::Size(image_size.width, image_size.height), CV_32SC1, cv::Scalar(0));
@@ -90,11 +91,18 @@ std::unique_ptr<ResultBase> SemanticSegmentationTiler::merge_results(const std::
     normalize_soft_prediction(merged_soft_prediction, voting_mask);
 
     cv::Mat hard_prediction = create_hard_prediction_from_soft_prediction(merged_soft_prediction, soft_threshold, blur_strength);
-    result->resultImage = hard_prediction;
 
+    std::unique_ptr<ResultBase> retVal;
     if (return_soft_prediction) {
+        auto* result = new ImageResultWithSoftPrediction();
+        retVal = std::unique_ptr<ResultBase>(result);
         result->soft_prediction = merged_soft_prediction;
+        result->resultImage = hard_prediction;
+    }
+    else {
+        auto* result = new ImageResult();
+        retVal = std::unique_ptr<ResultBase>(result);
+        result->resultImage = hard_prediction;
     }
-
     return retVal;
 }