Fix classification serialization

Had to set ov::batch(model, 1);

Author: RHeckerIntel

src/cpp/include/utils/config.h

@@ -97,6 +97,32 @@ static inline std::tuple<bool, ov::Layout> makeGuesLayoutFrom4DShape(const ov::P
     return {false, ov::Layout{}};
 }
 
+static inline std::map<std::string, ov::Layout> parseLayoutString(const std::string& layout_string) {
+    // Parse parameter string like "input0:NCHW,input1:NC" or "NCHW" (applied to all
+    // inputs)
+    std::map<std::string, ov::Layout> layouts;
+    std::string searchStr =
+        (layout_string.find_last_of(':') == std::string::npos && !layout_string.empty() ? ":" : "") + layout_string;
+    auto colonPos = searchStr.find_last_of(':');
+    while (colonPos != std::string::npos) {
+        auto startPos = searchStr.find_last_of(',');
+        auto inputName = searchStr.substr(startPos + 1, colonPos - startPos - 1);
+        auto inputLayout = searchStr.substr(colonPos + 1);
+        layouts[inputName] = ov::Layout(inputLayout);
+        searchStr.resize(startPos + 1);
+        if (searchStr.empty() || searchStr.back() != ',') {
+            break;
+        }
+        searchStr.pop_back();
+        colonPos = searchStr.find_last_of(':');
+    }
+    if (!searchStr.empty()) {
+        throw std::invalid_argument("Can't parse input layout string: " + layout_string);
+    }
+    return layouts;
+}
+
+
 static inline ov::Layout getLayoutFromShape(const ov::PartialShape& shape) {
     if (shape.size() == 2) {
         return "NC";
@@ -133,4 +159,21 @@ static inline ov::Layout getLayoutFromShape(const ov::PartialShape& shape) {
     throw std::runtime_error("Usupported " + std::to_string(shape.size()) + "D shape");
 }
 
+static inline ov::Layout getInputLayout(const ov::Output<ov::Node>& input, std::map<std::string, ov::Layout>& inputsLayouts) {
+    ov::Layout layout = ov::layout::get_layout(input);
+    if (layout.empty()) {
+        if (inputsLayouts.empty()) {
+            layout = getLayoutFromShape(input.get_partial_shape());
+            std::cout << "Automatically detected layout '" << layout.to_string() << "' for input '"
+                       << input.get_any_name() << "' will be used." << std::endl;
+        } else if (inputsLayouts.size() == 1) {
+            layout = inputsLayouts.begin()->second;
+        } else {
+            layout = inputsLayouts[input.get_any_name()];
+        }
+    }
+
+    return layout;
+}
+
 }  // namespace utils

src/cpp/src/tasks/classification.cpp

@@ -85,6 +85,11 @@ std::vector<size_t> get_non_xai_output_indices(const std::vector<ov::Output<ov::
 cv::Size Classification::serialize(std::shared_ptr<ov::Model>& ov_model) {
     // --------------------------- Configure input & output -------------------------------------------------
     // --------------------------- Prepare input  ------------------------------------------------------
+    auto config = ov_model->has_rt_info("model_info") ? ov_model->get_rt_info<ov::AnyMap>("model_info") : ov::AnyMap{};
+    std::string layout = "";
+    layout = utils::get_from_any_maps("layout", config, {}, layout);
+    auto inputsLayouts = utils::parseLayoutString(layout);
+
     if (ov_model->inputs().size() != 1) {
         throw std::logic_error("Classification model wrapper supports topologies with only 1 input");
     }
@@ -93,9 +98,7 @@ cv::Size Classification::serialize(std::shared_ptr<ov::Model>& ov_model) {
     auto inputName = input.get_any_name();
 
     const ov::Shape& inputShape = input.get_partial_shape().get_max_shape();
-    const ov::Layout& inputLayout = utils::getLayoutFromShape(inputShape);
-
-    auto config = ov_model->has_rt_info("model_info") ? ov_model->get_rt_info<ov::AnyMap>("model_info") : ov::AnyMap{};
+    const ov::Layout& inputLayout = utils::getInputLayout(input, inputsLayouts);
 
     auto interpolation_mode = cv::INTER_LINEAR;
     utils::RESIZE_MODE resize_mode = utils::RESIZE_FILL;
@@ -121,9 +124,6 @@ cv::Size Classification::serialize(std::shared_ptr<ov::Model>& ov_model) {
                                       mean_values,
                                       scale_values);
 
-    ov::preprocess::PrePostProcessor ppp = ov::preprocess::PrePostProcessor(ov_model);
-    ov_model = ppp.build();
-
     // --------------------------- Prepare output  -----------------------------------------------------
     if (ov_model->outputs().size() > 5) {
         throw std::logic_error("Classification model wrapper supports topologies with up to 4 outputs");

src/cpp/src/utils/preprocessing.cpp

@@ -24,6 +24,7 @@ std::shared_ptr<ov::Model> embedProcessing(std::shared_ptr<ov::Model>& model,
                                            const std::vector<float>& scale,
                                            const std::type_info& dtype) {
     ov::preprocess::PrePostProcessor ppp(model);
+
     // Change the input type to the 8-bit image
     if (dtype == typeid(int)) {
         ppp.input(inputName).tensor().set_element_type(ov::element::u8);
@@ -54,7 +55,9 @@ std::shared_ptr<ov::Model> embedProcessing(std::shared_ptr<ov::Model>& model,
         ppp.input(inputName).preprocess().scale(scale);
     }
 
-    return ppp.build();
+    auto ov_model = ppp.build();
+    ov::set_batch(ov_model, 1);
+    return ov_model;
 }
 
 ov::preprocess::PostProcessSteps::CustomPostprocessOp createResizeGraph(RESIZE_MODE resizeMode,

tests/cpp/test_accuracy.cpp

@@ -103,14 +103,13 @@ TEST_P(ModelParameterizedTest, AccuracyTest) {
     } else if (data.type == "MaskRCNNModel") {
         GTEST_SKIP();
     } else if (data.type == "ClassificationModel") {
-        GTEST_SKIP();
-        //auto model = Classification::load(model_path);
-        //for (auto& test_data : data.test_data) {
-        //    std::string image_path = DATA_DIR + '/' + test_data.image;
-        //    cv::Mat image = cv::imread(image_path);
-        //    auto result = model.infer(image);
-        //    EXPECT_EQ(std::string{result}, test_data.reference[0]);
-        //}
+        auto model = Classification::load(model_path);
+        for (auto& test_data : data.test_data) {
+            std::string image_path = DATA_DIR + '/' + test_data.image;
+            cv::Mat image = cv::imread(image_path);
+            auto result = model.infer(image);
+            EXPECT_EQ(std::string{result}, test_data.reference[0]);
+        }
     } else {
         FAIL() << "No implementation for model type " << data.type;
     }