SourceXtractorPlusPlus  0.12
Please provide a description of the project.
OnnxSourceTask.cpp
Go to the documentation of this file.
1 
22 #include <NdArray/NdArray.h>
24 #include <onnxruntime_cxx_api.h>
25 
26 namespace NdArray = Euclid::NdArray;
27 
28 namespace SourceXtractor {
29 
30 
31 template<typename T>
32 static void fillCutout(const Image<T>& image, int center_x, int center_y, int width, int height, std::vector<T>& out) {
33  int x_start = center_x - width / 2;
34  int y_start = center_y - height / 2;
35  int x_end = x_start + width;
36  int y_end = y_start + height;
37 
38  int index = 0;
39  for (int iy = y_start; iy < y_end; iy++) {
40  for (int ix = x_start; ix < x_end; ix++, index++) {
41  if (ix >= 0 && iy >= 0 && ix < image.getWidth() && iy < image.getHeight()) {
42  out[index] = image.getValue(ix, iy);
43  }
44  }
45  }
46 }
47 
48 OnnxSourceTask::OnnxSourceTask(const std::vector<OnnxModel>& models) : m_models(models) {}
49 
57 template<typename O>
59 computePropertiesSpecialized(const OnnxModel& model, const DetectionFrameImages& detection_frame_images,
60  const PixelCentroid& centroid) {
61  Ort::RunOptions run_options;
62  auto mem_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
63 
64  const int center_x = static_cast<int>(centroid.getCentroidX() + 0.5);
65  const int center_y = static_cast<int>(centroid.getCentroidY() + 0.5);
66 
67  // Allocate memory
68  std::vector<int64_t> input_shape(model.m_input_shape.begin(), model.m_input_shape.end());
69  input_shape[0] = 1;
70  size_t input_size = std::accumulate(input_shape.begin(), input_shape.end(), 1u, std::multiplies<size_t>());
71  std::vector<float> input_data(input_size);
72 
73  std::vector<int64_t> output_shape(model.m_output_shape.begin(), model.m_output_shape.end());
74  output_shape[0] = 1;
75  size_t output_size = std::accumulate(output_shape.begin(), output_shape.end(), 1u, std::multiplies<size_t>());
76  std::vector<O> output_data(output_size);
77 
78  // Cut the needed area
79  {
80  const auto& image = detection_frame_images.getLockedImage(LayerSubtractedImage);
81  fillCutout(*image, center_x, center_y, input_shape[2], input_shape[3], input_data);
82  }
83 
84  // Setup input/output tensors
85  auto input_tensor = Ort::Value::CreateTensor<float>(
86  mem_info, input_data.data(), input_data.size(), input_shape.data(), input_shape.size());
87  auto output_tensor = Ort::Value::CreateTensor<O>(
88  mem_info, output_data.data(), output_data.size(), output_shape.data(), output_shape.size());
89 
90  // Run the model
91  const char *input_name = model.m_input_name.c_str();
92  const char *output_name = model.m_output_name.c_str();
93  model.m_session->Run(run_options,
94  &input_name, &input_tensor, 1,
95  &output_name, &output_tensor, 1);
96 
97  // Set the output
98  std::vector<size_t> catalog_shape{model.m_output_shape.begin() + 1, model.m_output_shape.end()};
99  return Euclid::make_unique<OnnxProperty::NdWrapper<O>>(catalog_shape, output_data);
100 }
101 
103  const auto& detection_frame_images = source.getProperty<DetectionFrameImages>();
104  const auto& centroid = source.getProperty<PixelCentroid>();
105 
107 
108  for (const auto& model : m_models) {
110 
111  switch (model.m_output_type) {
112  case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
113  result = computePropertiesSpecialized<float>(model, detection_frame_images, centroid);
114  break;
115  case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
116  result = computePropertiesSpecialized<int32_t>(model, detection_frame_images, centroid);
117  break;
118  default:
119  throw Elements::Exception() << "This should have not happened!" << model.m_output_type;
120  }
121 
122  output_dict.emplace(model.m_prop_name, std::move(result));
123  }
124 
125  source.setProperty<OnnxProperty>(std::move(output_dict));
126 }
127 
128 } // end of namespace SourceXtractor
OnnxSourceTask(const std::vector< OnnxModel > &models)
Euclid::NdArray::NdArray< T > NdArray
std::vector< std::int64_t > m_input_shape
Input tensor shape.
Definition: OnnxModel.h:38
T end(T... args)
std::shared_ptr< Image< SeFloat > > getLockedImage(FrameImageLayer layer) const
The centroid of all the pixels in the source, weighted by their DetectionImage pixel values...
Definition: PixelCentroid.h:37
static std::unique_ptr< OnnxProperty::NdWrapperBase > computePropertiesSpecialized(const OnnxModel &model, const DetectionFrameImages &detection_frame_images, const PixelCentroid &centroid)
virtual int getHeight() const =0
Returns the height of the image in pixels.
std::string m_input_name
Input tensor name.
Definition: OnnxModel.h:34
STL class.
const std::vector< OnnxModel > & m_models
std::string m_output_name
Output tensor name.
Definition: OnnxModel.h:35
SeFloat getCentroidX() const
X coordinate of centroid.
Definition: PixelCentroid.h:48
T data(T... args)
const PropertyType & getProperty(unsigned int index=0) const
Convenience template method to call getProperty() with a more user-friendly syntax.
virtual T getValue(int x, int y) const =0
Returns the value of the pixel with the coordinates (x,y)
void computeProperties(SourceInterface &source) const override
Computes one or more properties for the Source.
T move(T... args)
std::vector< std::int64_t > m_output_shape
Output tensor shape.
Definition: OnnxModel.h:39
T size(T... args)
STL class.
STL class.
T begin(T... args)
SeFloat getCentroidY() const
Y coordinate of centroid.
Definition: PixelCentroid.h:53
T c_str(T... args)
Interface representing an image.
Definition: Image.h:43
static void fillCutout(const Image< T > &image, int center_x, int center_y, int width, int height, std::vector< T > &out)
T accumulate(T... args)
The SourceInterface is an abstract "source" that has properties attached to it.
virtual int getWidth() const =0
Returns the width of the image in pixels.
std::unique_ptr< Ort::Session > m_session
Session, one per model. In theory, it is thread-safe.
Definition: OnnxModel.h:41