Skip to content

Commit 1bef809

Browse files
dkalinowskidkalinowski
committed
save to disk
1 parent d459a77 commit 1bef809

File tree

8 files changed

+263
-48
lines changed

8 files changed

+263
-48
lines changed

‎WORKSPACE

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -599,7 +599,7 @@ new_git_repository(
599599
build_file_content = """
600600
cc_library(
601601
name = "image",
602-
hdrs = ["stb_image.h"],
602+
hdrs = ["stb_image.h", "stb_image_write.h"],
603603
visibility = ["//visibility:public"],
604604
local_defines = [
605605
],

‎src/BUILD

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2602,6 +2602,27 @@ cc_library(
26022602
linkopts = LINKOPTS_ADJUSTED,
26032603
)
26042604

2605+
cc_library(
2606+
name = "libimage_conversion",
2607+
hdrs = [
2608+
"image_conversion.hpp",
2609+
],
2610+
srcs = [
2611+
"image_conversion.cpp",
2612+
],
2613+
deps = [
2614+
"@stb//:image",
2615+
"@com_google_absl//absl/strings",
2616+
"//third_party:openvino",
2617+
"//src:libovmslogging",
2618+
"//src:libovmsprofiler",
2619+
],
2620+
visibility = ["//visibility:public",],
2621+
local_defines = COMMON_LOCAL_DEFINES,
2622+
copts = COPTS_ADJUSTED,
2623+
linkopts = LINKOPTS_ADJUSTED,
2624+
)
2625+
26052626
# HTTP Server implementation using net_http of tensorflow
26062627
# To use other library simply create new target and implementation of libhttp_async_writer_interface
26072628
cc_library(

‎src/image_conversion.cpp

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
//****************************************************************************
2+
// Copyright 2025 Intel Corporation
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+
#include "image_conversion.hpp"
17+
18+
#include <iostream>
19+
20+
#define STB_IMAGE_IMPLEMENTATION
21+
#define STB_IMAGE_WRITE_IMPLEMENTATION
22+
#include "logging.hpp"
23+
#include "profiler.hpp"
24+
#pragma warning(push)
25+
#pragma warning(disable : 6262)
26+
#include "stb_image.h" // NOLINT
27+
#include "stb_image_write.h" // NOLINT
28+
#pragma warning(default : 6262)
29+
#pragma warning(disable : 6001 4324 6385 6386)
30+
#include "absl/strings/escaping.h"
31+
#include "absl/strings/str_cat.h"
32+
#pragma warning(pop)
33+
34+
namespace ovms {
35+
36+
void hello() {
37+
std::cout << "Hello, World!" << std::endl;
38+
}
39+
40+
ov::Tensor load_image_stbi(const std::string& imageBytes) {
41+
int x = 0, y = 0, channelsInFile = 0;
42+
constexpr int desiredChannels = 3;
43+
unsigned char* image = stbi_load_from_memory(
44+
(const unsigned char*)imageBytes.data(), imageBytes.size(),
45+
&x, &y, &channelsInFile, desiredChannels);
46+
if (!image) {
47+
std::stringstream errorMessage;
48+
errorMessage << "Failed to load the image";
49+
throw std::runtime_error{errorMessage.str()};
50+
}
51+
struct SharedImageAllocator {
52+
unsigned char* image;
53+
int channels, height, width;
54+
void* allocate(size_t bytes, size_t) const {
55+
if (image && channels * height * width == bytes) {
56+
return image;
57+
}
58+
throw std::runtime_error{"Unexpected number of bytes was requested to allocate."};
59+
}
60+
void deallocate(void*, size_t bytes, size_t) {
61+
if (channels * height * width != bytes) {
62+
throw std::runtime_error{"Unexpected number of bytes was requested to deallocate."};
63+
}
64+
if (image != nullptr) {
65+
stbi_image_free(image);
66+
image = nullptr;
67+
}
68+
}
69+
bool is_equal(const SharedImageAllocator& other) const noexcept { return this == &other; }
70+
};
71+
return ov::Tensor(
72+
ov::element::u8,
73+
ov::Shape{1, size_t(y), size_t(x), size_t(desiredChannels)},
74+
SharedImageAllocator{image, desiredChannels, y, x});
75+
}
76+
77+
std::string save_image_stbi(ov::Tensor tensor) {
78+
// Validate tensor properties
79+
if (tensor.get_element_type() != ov::element::u8) {
80+
throw std::runtime_error{"Only U8 tensor element type is supported for image saving"};
81+
}
82+
83+
if (tensor.get_shape().size() != 4 || tensor.get_shape()[0] != 1) {
84+
throw std::runtime_error{"Tensor must be in NHWC format with batch size 1"};
85+
}
86+
87+
size_t height = tensor.get_shape()[1];
88+
size_t width = tensor.get_shape()[2];
89+
size_t channels = tensor.get_shape()[3];
90+
91+
if (channels != 3 && channels != 1) {
92+
throw std::runtime_error{"Only 1 or 3 channel images are supported for saving"};
93+
}
94+
95+
// Get pointer to image data
96+
unsigned char* image_data = tensor.data<unsigned char>();
97+
98+
// Create a memory buffer to hold the PNG data
99+
std::vector<unsigned char> png_buffer;
100+
101+
// Define the write function that will store data in our buffer
102+
auto write_func = [](void* context, void* data, int size) {
103+
std::vector<unsigned char>* buffer = static_cast<std::vector<unsigned char>*>(context);
104+
unsigned char* bytes = static_cast<unsigned char*>(data);
105+
buffer->insert(buffer->end(), bytes, bytes + size);
106+
};
107+
108+
// Write PNG to memory using our buffer
109+
int success = stbi_write_png_to_func(
110+
write_func, // Our write function
111+
&png_buffer, // Context (our buffer)
112+
width, // Image width
113+
height, // Image height
114+
channels, // Number of channels
115+
image_data, // Image data
116+
width * channels); // Stride (bytes per row)
117+
118+
if (!success) {
119+
throw std::runtime_error{"Failed to encode image to PNG format"};
120+
}
121+
122+
// Convert the buffer to a string
123+
return std::string(png_buffer.begin(), png_buffer.end());
124+
}
125+
126+
} // namespace ovms

‎src/image_conversion.hpp

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
//*****************************************************************************
2+
// Copyright 2025 Intel Corporation
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+
#pragma once
17+
18+
#include <string>
19+
20+
#include <openvino/runtime/tensor.hpp>
21+
22+
namespace ovms {
23+
24+
void hello();
25+
26+
ov::Tensor load_image_stbi(const std::string& imageBytes);
27+
std::string save_image_stbi(ov::Tensor tensor);
28+
29+
} // namespace ovms

‎src/image_gen/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ cc_library(
6262
"//src:libovmslogging",
6363
"image_gen_calculator_cc_proto",
6464
":pipelines",
65+
"//src:libimage_conversion",
6566
]+ select({
6667
"//conditions:default": ["//third_party:genai", ":llm_engine"],
6768
"//:not_genai_bin" : [":llm_engine"],

‎src/image_gen/http_image_gen_calculator.cc

Lines changed: 46 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
// See the License for the specific language governing permissions and
1414
// limitations under the License.
1515
//*****************************************************************************
16+
#include <fstream>
17+
1618
#pragma warning(push)
1719
#pragma warning(disable : 4005 4309 6001 6385 6386 6326 6011 6246 4456 6246)
1820
#pragma GCC diagnostic push
@@ -24,9 +26,27 @@
2426

2527
#include "../http_payload.hpp"
2628
#include "../logging.hpp"
29+
#include "../image_conversion.hpp"
2730

2831
#include "pipelines.hpp"
2932

33+
static void save_png_to_disk(const std::string& png_data, const std::string& filename) {
34+
std::ofstream out_file(filename, std::ios::binary);
35+
if (!out_file) {
36+
throw std::runtime_error("Failed to open file for writing: " + filename);
37+
}
38+
39+
out_file.write(png_data.data(), png_data.size());
40+
41+
if (!out_file.good()) {
42+
out_file.close();
43+
std::remove(filename.c_str()); // Clean up partial file
44+
throw std::runtime_error("Failed to write data to file: " + filename);
45+
}
46+
47+
out_file.close();
48+
}
49+
3050
using namespace ovms;
3151

3252
namespace mediapipe {
@@ -67,15 +87,35 @@ class ImageGenCalculator : public CalculatorBase {
6787
RET_CHECK(it != pipelinesNap.end()) << "Could not find initialized Image Gen node named: " << cc->NodeName();
6888
auto pipe = it->second;
6989

70-
// curl -X POST localhost:11338/v3/endpoint -d '{}'
90+
auto payload = cc->Inputs().Tag(INPUT_TAG_NAME).Get<ovms::HttpPayload>();
91+
if (payload.parsedJson->HasParseError())
92+
return absl::InvalidArgumentError("Failed to parse JSON");
93+
94+
if (!payload.parsedJson->IsObject()) {
95+
return absl::InvalidArgumentError("JSON body must be an object");
96+
}
97+
98+
// get prompt field as string
99+
auto promptIt = payload.parsedJson->FindMember("prompt");
100+
if (promptIt == payload.parsedJson->MemberEnd()) {
101+
return absl::InvalidArgumentError("prompt field is missing in JSON body");
102+
}
103+
if (!promptIt->value.IsString()) {
104+
return absl::InvalidArgumentError("prompt field is not a string");
105+
}
106+
std::string prompt = promptIt->value.GetString();
107+
108+
// curl -X POST localhost:11338/v3/images/generations -H "Content-Type: application/json" -d '{ "model": "endpoint", "prompt": "A cute baby sea otter", "n": 1, "size": "1024x1024" }'
71109
ov::genai::Text2ImagePipeline::GenerationRequest request = pipe->text2ImagePipeline.create_generation_request();
72-
ov::Tensor image = request.generate("a cat", // TODO: get from payload
110+
ov::Tensor image = request.generate(prompt,
73111
ov::AnyMap{
74-
ov::genai::width(512),
75-
ov::genai::height(512),
76-
ov::genai::num_inference_steps(20),
77-
ov::genai::num_images_per_prompt(1)});
112+
ov::genai::width(512), // todo: get from req
113+
ov::genai::height(512), // todo: get from req
114+
ov::genai::num_inference_steps(20), // todo: get from req
115+
ov::genai::num_images_per_prompt(1)}); // todo: get from req
78116

117+
std::string res = save_image_stbi(image);
118+
save_png_to_disk(res, "output.png");
79119
SPDLOG_LOGGER_DEBUG(llm_calculator_logger, "ImageGenCalculator [Node: {}] Process end", cc->NodeName());
80120
return absl::OkStatus();
81121
}

‎src/llm/BUILD

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ cc_library(
7979
"@mediapipe//mediapipe/framework:calculator_framework", # required for absl status
8080
"//src:libovmsprofiler",
8181
"//third_party:opencv",
82-
"@stb//:image",
82+
"//src:libimage_conversion",
8383
] + select({
8484
"//conditions:default": ["//third_party:genai", ":llm_engine"],
8585
"//:not_genai_bin" : [":llm_engine"],

‎src/llm/apis/openai_completions.cpp

Lines changed: 38 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -27,18 +27,16 @@
2727
#include <rapidjson/writer.h>
2828
#pragma warning(pop)
2929

30-
#define STB_IMAGE_IMPLEMENTATION
3130
#include "../../logging.hpp"
3231
#include "../../profiler.hpp"
3332
#pragma warning(push)
34-
#pragma warning(disable : 6262)
35-
#include "stb_image.h" // NOLINT
36-
#pragma warning(default : 6262)
3733
#pragma warning(disable : 6001 4324 6385 6386)
3834
#include "absl/strings/escaping.h"
3935
#include "absl/strings/str_cat.h"
4036
#pragma warning(pop)
4137

38+
#include "../../image_conversion.hpp" // TODO: Rename to stbi_conversions?
39+
4240
using namespace rapidjson;
4341

4442
namespace ovms {
@@ -91,42 +89,42 @@ absl::Status OpenAIChatCompletionsHandler::parseCompletionsPart() {
9189
return absl::OkStatus();
9290
}
9391

94-
ov::Tensor load_image_stbi(const std::string& imageBytes) {
95-
int x = 0, y = 0, channelsInFile = 0;
96-
constexpr int desiredChannels = 3;
97-
unsigned char* image = stbi_load_from_memory(
98-
(const unsigned char*)imageBytes.data(), imageBytes.size(),
99-
&x, &y, &channelsInFile, desiredChannels);
100-
if (!image) {
101-
std::stringstream errorMessage;
102-
errorMessage << "Failed to load the image";
103-
throw std::runtime_error{errorMessage.str()};
104-
}
105-
struct SharedImageAllocator {
106-
unsigned char* image;
107-
int channels, height, width;
108-
void* allocate(size_t bytes, size_t) const {
109-
if (image && channels * height * width == bytes) {
110-
return image;
111-
}
112-
throw std::runtime_error{"Unexpected number of bytes was requested to allocate."};
113-
}
114-
void deallocate(void*, size_t bytes, size_t) {
115-
if (channels * height * width != bytes) {
116-
throw std::runtime_error{"Unexpected number of bytes was requested to deallocate."};
117-
}
118-
if (image != nullptr) {
119-
stbi_image_free(image);
120-
image = nullptr;
121-
}
122-
}
123-
bool is_equal(const SharedImageAllocator& other) const noexcept { return this == &other; }
124-
};
125-
return ov::Tensor(
126-
ov::element::u8,
127-
ov::Shape{1, size_t(y), size_t(x), size_t(desiredChannels)},
128-
SharedImageAllocator{image, desiredChannels, y, x});
129-
}
92+
// ov::Tensor load_image_stbi(const std::string& imageBytes) {
93+
// int x = 0, y = 0, channelsInFile = 0;
94+
// constexpr int desiredChannels = 3;
95+
// unsigned char* image = stbi_load_from_memory(
96+
// (const unsigned char*)imageBytes.data(), imageBytes.size(),
97+
// &x, &y, &channelsInFile, desiredChannels);
98+
// if (!image) {
99+
// std::stringstream errorMessage;
100+
// errorMessage << "Failed to load the image";
101+
// throw std::runtime_error{errorMessage.str()};
102+
// }
103+
// struct SharedImageAllocator {
104+
// unsigned char* image;
105+
// int channels, height, width;
106+
// void* allocate(size_t bytes, size_t) const {
107+
// if (image && channels * height * width == bytes) {
108+
// return image;
109+
// }
110+
// throw std::runtime_error{"Unexpected number of bytes was requested to allocate."};
111+
// }
112+
// void deallocate(void*, size_t bytes, size_t) {
113+
// if (channels * height * width != bytes) {
114+
// throw std::runtime_error{"Unexpected number of bytes was requested to deallocate."};
115+
// }
116+
// if (image != nullptr) {
117+
// stbi_image_free(image);
118+
// image = nullptr;
119+
// }
120+
// }
121+
// bool is_equal(const SharedImageAllocator& other) const noexcept { return this == &other; }
122+
// };
123+
// return ov::Tensor(
124+
// ov::element::u8,
125+
// ov::Shape{1, size_t(y), size_t(x), size_t(desiredChannels)},
126+
// SharedImageAllocator{image, desiredChannels, y, x});
127+
// }
130128

131129
absl::Status OpenAIChatCompletionsHandler::parseMessages() {
132130
auto it = doc.FindMember("messages");

0 commit comments

Comments
 (0)