145 lines
4.1 KiB
C++
145 lines
4.1 KiB
C++
#include <onnxruntime/cpu_provider_factory.h>
|
|
#include <onnxruntime/onnxruntime_cxx_api.h>
|
|
#include <png.h>
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cassert>
|
|
#include <cmath>
|
|
#include <iostream>
|
|
#include <span>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
// https://onnxruntime.ai/docs/api/c/namespace_ort.html
|
|
// https://github.com/microsoft/onnxruntime-inference-examples/
|
|
|
|
class ImageData;
|
|
|
|
class OnnxMnist
|
|
{
|
|
public:
|
|
OnnxMnist() noexcept = default;
|
|
|
|
auto run(ImageData& image) noexcept -> std::size_t;
|
|
|
|
static constexpr const int kWidth{28};
|
|
static constexpr const int kHeight{28};
|
|
|
|
private:
|
|
// From session we can get the input and output count
|
|
Ort::Session session_{Ort::Env{}, "model.onnx", Ort::SessionOptions{nullptr}};
|
|
};
|
|
|
|
class ImageData
|
|
{
|
|
public:
|
|
explicit ImageData(std::string_view input_file);
|
|
|
|
auto getData() noexcept -> std::span<float> { return data_; }
|
|
auto getWidth() const noexcept { return width_; }
|
|
auto getHeight() const noexcept { return height_; }
|
|
|
|
private:
|
|
std::size_t width_{0};
|
|
std::size_t height_{0};
|
|
|
|
std::vector<float> data_;
|
|
};
|
|
|
|
template<typename T>
|
|
void softmax(std::span<T> input)
|
|
{
|
|
auto row_max = *std::ranges::max_element(input);
|
|
std::vector<float> result(input.size());
|
|
float sum{0.0};
|
|
for (std::size_t i = 0; i != input.size(); ++i) {
|
|
sum += result[i] = std::exp(input[i] - row_max);
|
|
}
|
|
for (std::size_t i = 0; i != input.size(); ++i) {
|
|
input[i] = result[i] / sum;
|
|
}
|
|
}
|
|
|
|
auto main(int argc, char* argv[]) -> int
|
|
{
|
|
if (argc < 2) {
|
|
std::cout << "usage: <input_file>" << std::endl;
|
|
return -1;
|
|
}
|
|
|
|
std::string_view filename{argv[1]}; // NOLINT
|
|
ImageData image{filename};
|
|
|
|
if (image.getWidth() != OnnxMnist::kWidth || image.getHeight() != OnnxMnist::kHeight) {
|
|
std::cout << "Resize to image to 28x28" << std::endl;
|
|
return -1;
|
|
}
|
|
|
|
OnnxMnist model;
|
|
auto result = model.run(image);
|
|
std::cout << "Result : " << result << std::endl;
|
|
return 0;
|
|
}
|
|
|
|
auto OnnxMnist::run(ImageData& image) noexcept -> std::size_t
|
|
{
|
|
const char* kInputNames[]{"Input3"}; // NOLINT
|
|
const char* kOutputNames[]{"Plus214_Output_0"}; // NOLINT
|
|
|
|
constexpr std::array<int64_t, 4> kInputShape{1, 1, kWidth, kHeight};
|
|
constexpr std::array<int64_t, 2> kOutputShape{1, 10};
|
|
|
|
const auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
|
|
|
|
auto image_data = image.getData();
|
|
Ort::Value const input_tensor = Ort::Value::CreateTensor<float>(
|
|
memory_info, image_data.data(), image_data.size(), kInputShape.data(), kInputShape.size());
|
|
|
|
std::array<float, kOutputShape[1]> results_{};
|
|
auto output_tensor = Ort::Value::CreateTensor<float>(
|
|
memory_info, results_.data(), results_.size(), kOutputShape.data(), kOutputShape.size());
|
|
|
|
const Ort::RunOptions run_options;
|
|
session_.Run(run_options, kInputNames, &input_tensor, 1, kOutputNames, &output_tensor, 1);
|
|
|
|
softmax<float>(results_);
|
|
for (std::size_t i = 0; i < results_.size(); i++) {
|
|
std::cout << i << " : " << results_[i] << std::endl;
|
|
}
|
|
|
|
auto result = std::distance(results_.begin(), std::max_element(results_.begin(), results_.end()));
|
|
|
|
return result;
|
|
}
|
|
|
|
ImageData::ImageData(std::string_view input_file)
|
|
{
|
|
// The control structure used by libpng
|
|
// Initialize the 'png_image' structure.
|
|
png_image image;
|
|
memset(&image, 0, (sizeof image));
|
|
image.version = PNG_IMAGE_VERSION;
|
|
if (png_image_begin_read_from_file(&image, input_file.data()) == 0) {
|
|
return;
|
|
}
|
|
image.format = PNG_FORMAT_BGR;
|
|
|
|
const std::size_t input_data_length = PNG_IMAGE_SIZE(image); // NOLINT
|
|
std::vector<png_byte> buffer(input_data_length, 0);
|
|
if (png_image_finish_read(&image, nullptr /*background*/, buffer.data(), 0 /*row_stride*/, nullptr /*colormap*/) ==
|
|
0) {
|
|
return;
|
|
}
|
|
|
|
width_ = image.width;
|
|
height_ = image.height;
|
|
|
|
const auto stride = height_ * width_;
|
|
data_.resize(height_ * width_);
|
|
|
|
for (std::size_t i = 0; i < stride; ++i) {
|
|
data_[i] = buffer[i * 3] == 0 ? 1.0F : 0.0F;
|
|
}
|
|
}
|