runtime/contrib/tflite_classify/src/ImageClassifier.cc

   1 /*
   2  * Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *    http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "ImageClassifier.h"
  18
  19 #include <fstream>
  20 #include <queue>
  21 #include <algorithm>
  22
  23 ImageClassifier::ImageClassifier(const std::string &model_file, const std::string &label_file,
  24                                  const int input_size, const int image_mean, const int image_std,
  25                                  const std::string &input_name, const std::string &output_name,
  26                                  const bool use_nnapi)
  27     : _inference(new InferenceInterface(model_file, use_nnapi)), _input_size(input_size),
  28       _image_mean(image_mean), _image_std(image_std), _input_name(input_name),
  29       _output_name(output_name)
  30 {
  31   // Load label
  32   std::ifstream label_stream(label_file.c_str());
  33   assert(label_stream);
  34
  35   std::string line;
  36   while (std::getline(label_stream, line))
  37   {
  38     _labels.push_back(line);
  39   }
  40   _num_classes = _inference->getTensorSize(_output_name);
  41   std::cout << "Output tensor size is " << _num_classes << ", label size is " << _labels.size()
  42             << std::endl;
  43
  44   // Pre-allocate buffers
  45   _fdata.reserve(_input_size * _input_size * 3);
  46   _outputs.reserve(_num_classes);
  47 }
  48
  49 std::vector<Recognition> ImageClassifier::recognizeImage(const cv::Mat &image)
  50 {
  51   // Resize image
  52   cv::Mat cropped;
  53   cv::resize(image, cropped, cv::Size(_input_size, _input_size), 0, 0, cv::INTER_AREA);
  54
  55   // Preprocess the image data from 0~255 int to normalized float based
  56   // on the provided parameters
  57   _fdata.clear();
  58   for (int y = 0; y < cropped.rows; ++y)
  59   {
  60     for (int x = 0; x < cropped.cols; ++x)
  61     {
  62       cv::Vec3b color = cropped.at<cv::Vec3b>(y, x);
  63       color[0] = color[0] - (float)_image_mean / _image_std;
  64       color[1] = color[1] - (float)_image_mean / _image_std;
  65       color[2] = color[2] - (float)_image_mean / _image_std;
  66
  67       _fdata.push_back(color[0]);
  68       _fdata.push_back(color[1]);
  69       _fdata.push_back(color[2]);
  70
  71       cropped.at<cv::Vec3b>(y, x) = color;
  72     }
  73   }
  74
  75   // Copy the input data into model
  76   _inference->feed(_input_name, _fdata, 1, _input_size, _input_size, 3);
  77
  78   // Run the inference call
  79   _inference->run(_output_name);
  80
  81   // Copy the output tensor back into the output array
  82   _inference->fetch(_output_name, _outputs);
  83
  84   // Find the best classifications
  85   auto compare = [](const Recognition &lhs, const Recognition &rhs) {
  86     return lhs.confidence < rhs.confidence;
  87   };
  88
  89   std::priority_queue<Recognition, std::vector<Recognition>, decltype(compare)> pq(compare);
  90   for (int i = 0; i < _num_classes; ++i)
  91   {
  92     if (_outputs[i] > _threshold)
  93     {
  94       pq.push(Recognition(_outputs[i], _labels[i]));
  95     }
  96   }
  97
  98   std::vector<Recognition> results;
  99   int min = std::min(pq.size(), _max_results);
 100   for (int i = 0; i < min; ++i)
 101   {
 102     results.push_back(pq.top());
 103     pq.pop();
 104   }
 105
 106   return results;
 107 }