inference-engine/samples/hello_shape_infer_ssd/main.cpp

   1 // Copyright (C) 2018-2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #include <vector>
   6 #include <memory>
   7 #include <string>
   8
   9 #include <inference_engine.hpp>
  10 #include <ext_list.hpp>
  11 #include <samples/ocv_common.hpp>
  12
  13 #include "shape_infer_extension.hpp"
  14
  15 using namespace InferenceEngine;
  16
  17 int main(int argc, char* argv[]) {
  18     try {
  19         // ------------------------------ Parsing and validation of input args ---------------------------------
  20         if (argc != 5) {
  21             std::cout << "Usage : ./hello_shape_infer_ssd <path_to_model> <path_to_image> <device> <batch>"
  22                       << std::endl;
  23             return EXIT_FAILURE;
  24         }
  25         const std::string input_model{argv[1]};
  26         const std::string input_image_path{argv[2]};
  27         const std::string device_name{argv[3]};
  28         const size_t batch_size{std::stoul(argv[4])};
  29         // -----------------------------------------------------------------------------------------------------
  30
  31         // --------------------------- 1. Load Plugin for inference engine -------------------------------------
  32         InferencePlugin plugin = PluginDispatcher().getPluginByDevice(device_name);
  33         IExtensionPtr cpuExtension, inPlaceExtension;
  34         if (device_name == "CPU") {
  35             cpuExtension = std::make_shared<Extensions::Cpu::CpuExtensions>();
  36             inPlaceExtension = std::make_shared<InPlaceExtension>();
  37             plugin.AddExtension(cpuExtension);
  38             // register sample's custom kernel (CustomReLU)
  39             plugin.AddExtension(inPlaceExtension);
  40         }
  41         // -----------------------------------------------------------------------------------------------------
  42
  43         // --------------------------- 2. Read IR Generated by ModelOptimizer (.xml and .bin files) ------------
  44         CNNNetReader network_reader;
  45         network_reader.ReadNetwork(input_model);
  46         network_reader.ReadWeights(input_model.substr(0, input_model.size() - 4) + ".bin");
  47         CNNNetwork network = network_reader.getNetwork();
  48
  49         OutputsDataMap outputs_info(network.getOutputsInfo());
  50         InputsDataMap inputs_info(network.getInputsInfo());
  51         if (inputs_info.size() != 1 && outputs_info.size() != 1)
  52             throw std::logic_error("Sample supports clean SSD network with one input and one output");
  53
  54         // --------------------------- Resize network to match image sizes and given batch----------------------
  55         if (device_name == "CPU") {
  56             network.AddExtension(cpuExtension);
  57             // register sample's custom shape inference (CustomReLU)
  58             network.AddExtension(inPlaceExtension);
  59         }
  60         auto input_shapes = network.getInputShapes();
  61         std::string input_name;
  62         SizeVector input_shape;
  63         std::tie(input_name, input_shape) = *input_shapes.begin();
  64         cv::Mat image = cv::imread(input_image_path);
  65         input_shape[0] = batch_size;
  66         input_shape[2] = image.rows;
  67         input_shape[3] = image.cols;
  68         input_shapes[input_name] = input_shape;
  69         std::cout << "Resizing network to the image size = [" << image.rows << "x" << image.cols << "] "
  70                   << "with batch = " << batch_size << std::endl;
  71         network.reshape(input_shapes);
  72         // -----------------------------------------------------------------------------------------------------
  73
  74         // --------------------------- 3. Configure input & output ---------------------------------------------
  75         // --------------------------- Prepare input blobs -----------------------------------------------------
  76         InputInfo::Ptr input_info;
  77         std::tie(input_name, input_info) = *inputs_info.begin();
  78         input_info->setLayout(Layout::NCHW);
  79         input_info->setPrecision(Precision::U8);
  80         // --------------------------- Prepare output blobs ----------------------------------------------------
  81         DataPtr output_info;
  82         std::string output_name;
  83         std::tie(output_name, output_info) = *outputs_info.begin();
  84         if (output_info->creatorLayer.lock()->type != "DetectionOutput")
  85             throw std::logic_error("Can't find a DetectionOutput layer in the topology");
  86         const SizeVector output_shape = output_info->getTensorDesc().getDims();
  87         const int max_proposal_count = output_shape[2];
  88         const int object_size = output_shape[3];
  89         if (object_size != 7) {
  90             throw std::logic_error("Output item should have 7 as a last dimension");
  91         }
  92         if (output_shape.size() != 4) {
  93             throw std::logic_error("Incorrect output dimensions for SSD model");
  94         }
  95         if (output_info == nullptr) {
  96             THROW_IE_EXCEPTION << "[SAMPLES] internal error - output information is empty";
  97         }
  98
  99         output_info->setPrecision(Precision::FP32);
 100
 101         auto dumpVec = [](const SizeVector& vec) -> std::string {
 102             if (vec.empty()) return "[]";
 103             std::stringstream oss;
 104             oss << "[" << vec[0];
 105             for (size_t i = 1; i < vec.size(); i++) oss << "," << vec[i];
 106             oss << "]";
 107             return oss.str();
 108         };
 109         std::cout << "Resulting input shape = " << dumpVec(input_shape) << std::endl;
 110         std::cout << "Resulting output shape = " << dumpVec(output_shape) << std::endl;
 111         // -----------------------------------------------------------------------------------------------------
 112
 113         // --------------------------- 4. Loading model to the plugin ------------------------------------------
 114         ExecutableNetwork executable_network = plugin.LoadNetwork(network, {});
 115         // -----------------------------------------------------------------------------------------------------
 116
 117         // --------------------------- 5. Create infer request -------------------------------------------------
 118         InferRequest infer_request = executable_network.CreateInferRequest();
 119         // -----------------------------------------------------------------------------------------------------
 120
 121         // --------------------------- 6. Prepare input --------------------------------------------------------
 122         Blob::Ptr input = infer_request.GetBlob(input_name);
 123         for (size_t b = 0; b < batch_size; b++) {
 124             matU8ToBlob<uint8_t>(image, input, b);
 125         }
 126         // -----------------------------------------------------------------------------------------------------
 127
 128         // --------------------------- 7. Do inference --------------------------------------------------------
 129         infer_request.Infer();
 130         // -----------------------------------------------------------------------------------------------------
 131
 132         // --------------------------- 8. Process output ------------------------------------------------------
 133         Blob::Ptr output = infer_request.GetBlob(output_name);
 134         const float* detection = output->buffer().as<PrecisionTrait<Precision::FP32>::value_type*>();
 135
 136         /* Each detection has image_id that denotes processed image */
 137         for (int cur_proposal = 0; cur_proposal < max_proposal_count; cur_proposal++) {
 138             float image_id = detection[cur_proposal * object_size + 0];
 139             float label = detection[cur_proposal * object_size + 1];
 140             float confidence = detection[cur_proposal * object_size + 2];
 141             /* CPU and GPU plugins have difference in DetectionOutput layer, so we need both checks */
 142             if (image_id < 0 || confidence == 0) {
 143                 continue;
 144             }
 145
 146             float xmin = detection[cur_proposal * object_size + 3] * image.cols;
 147             float ymin = detection[cur_proposal * object_size + 4] * image.rows;
 148             float xmax = detection[cur_proposal * object_size + 5] * image.cols;
 149             float ymax = detection[cur_proposal * object_size + 6] * image.rows;
 150
 151             if (confidence > 0.5) {
 152                 /** Drawing only objects with >50% probability **/
 153                 std::ostringstream conf;
 154                 conf << ":" << std::fixed << std::setprecision(3) << confidence;
 155                 cv::rectangle(image, cv::Point2f(xmin, ymin), cv::Point2f(xmax, ymax), cv::Scalar(0, 0, 255));
 156                 std::cout << "[" << cur_proposal << "," << label << "] element, prob = " << confidence <<
 157                           ", bbox = (" << xmin << "," << ymin << ")-(" << xmax << "," << ymax << ")" << ", batch id = "
 158                           << image_id << std::endl;
 159             }
 160         }
 161
 162         cv::imwrite("hello_shape_infer_ssd_output.jpg", image);
 163         std::cout << "The resulting image was saved in the file: hello_shape_infer_ssd_output.jpg" << std::endl;
 164         // -----------------------------------------------------------------------------------------------------
 165     } catch (const std::exception& ex) {
 166         std::cerr << ex.what() << std::endl;
 167         return EXIT_FAILURE;
 168     }
 169     return EXIT_SUCCESS;
 170 }