-
-
Notifications
You must be signed in to change notification settings - Fork 863
Expand file tree
/
Copy pathimage_processor.cpp
More file actions
67 lines (52 loc) · 2.65 KB
/
image_processor.cpp
File metadata and controls
67 lines (52 loc) · 2.65 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
#include "image_processor.h"
#include "logger.h"
CLIPImageProcessor::CLIPImageProcessor(const std::string& model_path, const std::string& processor_filename) {
Ort::SessionOptions session_options;
session_options.EnableOrtCustomOps();
auto processor_path = model_path + "/" + processor_filename;
LOG(INFO) << "Loading image processor from " << processor_path;
session_ = std::make_unique<Ort::Session>(env_, processor_path.c_str(), session_options);
}
Option<processed_image_t> CLIPImageProcessor::process_image(const std::string& image_encoded) {
std::unique_lock<std::mutex> lock(mutex_);
// Decode image
auto image = StringUtils::base64_decode(image_encoded);
// convert string to byte array
std::vector<uint8_t> image_bytes(image.begin(), image.end());
// Create input tensor
int64_t input_tensor_size = image_bytes.size();
std::vector<int64_t> input_shape = {input_tensor_size};
std::vector<const char*> input_names = {"image"};
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
auto input_tensor = Ort::Value::CreateTensor<uint8_t>(memory_info, image_bytes.data(), image_bytes.size(), input_shape.data(), input_shape.size());
// Create output tensor
std::vector<const char*> output_names = {"last_hidden_state"};
// Run inference
std::vector<Ort::Value> output_tensors;
// LOG(INFO) << "Running image processor";
try {
output_tensors = session_->Run(Ort::RunOptions{nullptr}, input_names.data(), &input_tensor, 1, output_names.data(), output_names.size());
} catch (...) {
return Option<processed_image_t>(400, "Error while processing image");
}
// Get output tensor
auto output_tensor = output_tensors.front().GetTensorMutableData<float>();
// Convert output tensor to processed_image_t
auto output_shape = output_tensors.front().GetTensorTypeAndShapeInfo().GetShape();
if (output_shape.size() != 4) {
LOG(INFO) << "Output tensor shape is not 4D";
return Option<processed_image_t>(400, "Error while processing image");
}
processed_image_t output;
for (size_t i = 0; i < output_shape[0]; i++) {
for (size_t j = 0; j < output_shape[1]; j++) {
for (size_t k = 0; k < output_shape[2]; k++) {
for (size_t l = 0; l < output_shape[3]; l++) {
output.push_back(output_tensor[i * output_shape[1] * output_shape[2] * output_shape[3] + j * output_shape[2] * output_shape[3] + k * output_shape[3] + l]);
}
}
}
}
// LOG(INFO) << "Image processed";
return Option<processed_image_t>(std::move(output));
}