1 files changed, 157 insertions, 0 deletions
diff --git a/python/openvino/runtime/common/models/src/segmentation_model.cpp b/python/openvino/runtime/common/models/src/segmentation_model.cpp
new file mode 100644
index 0000000..82a153b
--- /dev/null
+++ b/python/openvino/runtime/common/models/src/segmentation_model.cpp
@@ -0,0 +1,157 @@
+/*
+// Copyright (C) 2020-2022 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+*/
+
+#include "models/segmentation_model.h"
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <fstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include <opencv2/core.hpp>
+#include <opencv2/imgproc.hpp>
+#include <openvino/openvino.hpp>
+
+#include "models/internal_model_data.h"
+#include "models/results.h"
+
+SegmentationModel::SegmentationModel(const std::string& modelFileName, bool useAutoResize, const std::string& layout)
+    : ImageModel(modelFileName, useAutoResize, layout) {}
+
+std::vector<std::string> SegmentationModel::loadLabels(const std::string& labelFilename) {
+    std::vector<std::string> labelsList;
+
+    /* Read labels (if any) */
+    if (!labelFilename.empty()) {
+        std::ifstream inputFile(labelFilename);
+        if (!inputFile.is_open())
+            throw std::runtime_error("Can't open the labels file: " + labelFilename);
+        std::string label;
+        while (std::getline(inputFile, label)) {
+            labelsList.push_back(label);
+        }
+        if (labelsList.empty())
+            throw std::logic_error("File is empty: " + labelFilename);
+    }
+
+    return labelsList;
+}
+
+void SegmentationModel::prepareInputsOutputs(std::shared_ptr<ov::Model>& model) {
+    // --------------------------- Configure input & output ---------------------------------------------
+    // --------------------------- Prepare input  -----------------------------------------------------
+    if (model->inputs().size() != 1) {
+        throw std::logic_error("Segmentation model wrapper supports topologies with only 1 input");
+    }
+    const auto& input = model->input();
+    inputsNames.push_back(input.get_any_name());
+
+    const ov::Layout& inputLayout = getInputLayout(input);
+    const ov::Shape& inputShape = input.get_shape();
+    if (inputShape.size() != 4 || inputShape[ov::layout::channels_idx(inputLayout)] != 3) {
+        throw std::logic_error("3-channel 4-dimensional model's input is expected");
+    }
+
+    ov::preprocess::PrePostProcessor ppp(model);
+    ppp.input().tensor().set_element_type(ov::element::u8).set_layout({"NHWC"});
+
+    if (useAutoResize) {
+        ppp.input().tensor().set_spatial_dynamic_shape();
+
+        ppp.input()
+            .preprocess()
+            .convert_element_type(ov::element::f32)
+            .resize(ov::preprocess::ResizeAlgorithm::RESIZE_LINEAR);
+    }
+
+    ppp.input().model().set_layout(inputLayout);
+    model = ppp.build();
+    // --------------------------- Prepare output  -----------------------------------------------------
+    if (model->outputs().size() != 1) {
+        throw std::logic_error("Segmentation model wrapper supports topologies with only 1 output");
+    }
+
+    const auto& output = model->output();
+    outputsNames.push_back(output.get_any_name());
+
+    const ov::Shape& outputShape = output.get_shape();
+    ov::Layout outputLayout("");
+    switch (outputShape.size()) {
+        case 3:
+            outputLayout = "CHW";
+            outChannels = 1;
+            outHeight = static_cast<int>(outputShape[ov::layout::height_idx(outputLayout)]);
+            outWidth = static_cast<int>(outputShape[ov::layout::width_idx(outputLayout)]);
+            break;
+        case 4:
+            outputLayout = "NCHW";
+            outChannels = static_cast<int>(outputShape[ov::layout::channels_idx(outputLayout)]);
+            outHeight = static_cast<int>(outputShape[ov::layout::height_idx(outputLayout)]);
+            outWidth = static_cast<int>(outputShape[ov::layout::width_idx(outputLayout)]);
+            break;
+        default:
+            throw std::logic_error("Unexpected output tensor shape. Only 4D and 3D outputs are supported.");
+    }
+}
+
+std::unique_ptr<ResultBase> SegmentationModel::postprocess(InferenceResult& infResult) {
+    ImageResult* result = new ImageResult(infResult.frameId, infResult.metaData);
+    const auto& inputImgSize = infResult.internalModelData->asRef<InternalImageModelData>();
+    const auto& outTensor = infResult.getFirstOutputTensor();
+
+    result->resultImage = cv::Mat(outHeight, outWidth, CV_8UC1);
+
+    if (outChannels == 1 && outTensor.get_element_type() == ov::element::i32) {
+        cv::Mat predictions(outHeight, outWidth, CV_32SC1, outTensor.data<int32_t>());
+        predictions.convertTo(result->resultImage, CV_8UC1);
+    } else if (outChannels == 1 && outTensor.get_element_type() == ov::element::i64) {
+        cv::Mat predictions(outHeight, outWidth, CV_32SC1);
+        const auto data = outTensor.data<int64_t>();
+        for (size_t i = 0; i < predictions.total(); ++i) {
+            reinterpret_cast<int32_t*>(predictions.data)[i] = int32_t(data[i]);
+        }
+        predictions.convertTo(result->resultImage, CV_8UC1);
+    } else if (outTensor.get_element_type() == ov::element::f32) {
+        const float* data = outTensor.data<float>();
+        for (int rowId = 0; rowId < outHeight; ++rowId) {
+            for (int colId = 0; colId < outWidth; ++colId) {
+                int classId = 0;
+                float maxProb = -1.0f;
+                for (int chId = 0; chId < outChannels; ++chId) {
+                    float prob = data[chId * outHeight * outWidth + rowId * outWidth + colId];
+                    if (prob > maxProb) {
+                        classId = chId;
+                        maxProb = prob;
+                    }
+                }  // nChannels
+
+                result->resultImage.at<uint8_t>(rowId, colId) = classId;
+            }  // width
+        }  // height
+    }
+
+    cv::resize(result->resultImage,
+               result->resultImage,
+               cv::Size(inputImgSize.inputImgWidth, inputImgSize.inputImgHeight),
+               0,
+               0,
+               cv::INTER_NEAREST);
+
+    return std::unique_ptr<ResultBase>(result);
+}