GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/onnxruntime-extensions

Public

mirrored fromhttps://github.com/microsoft/onnxruntime-extensionsAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
skottmckay/BuildInfra_AndTestImageLibs

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

onnxruntime-extensions/operators/vision/impl

operators/vision/impl/png_encoder_decoder.cc

250lines · modecode

RawDownload
Latest commit unavailable.
1#include "png_encoder_decoder.hpp"
2
3/****************************************************************************************\
4 This part of the file implements PNG codec on base of libpng library,
5 in particular, this code is based on example.c from libpng
6 (see otherlibs/_graphics/readme.txt for copyright notice)
7 and png2bmp sample from libpng distribution (Copyright (C) 1999-2001 MIYASAKA Masaru)
8\****************************************************************************************/
9
10#include "png.h"
11#include "zlib.h"
12
13#include <cassert>
14
15#if defined _MSC_VER && _MSC_VER >= 1200
16// interaction between '_setjmp' and C++ object destruction is non-portable
17#pragma warning(disable : 4611)
18#endif
19
20namespace ort_extensions {
21
22bool PngDecoder::IsPng(const uint8_t* bytes, uint64_t num_bytes) {
23 // '\0x89PGN\r\n<sub>\n'
24 constexpr const uint8_t signature[] = {0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a}; // check start of file for this
25 constexpr int signature_len = 8;
26
27 return num_bytes > signature_len && memcmp(bytes, signature, signature_len) == 0;
28}
29
30PngDecoder::~PngDecoder() {
31 if (png_ptr_) {
32 png_structp png_ptr = (png_structp)png_ptr_;
33 png_infop info_ptr = (png_infop)info_ptr_;
34 png_infop end_info = (png_infop)end_info_;
35 png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
36 }
37}
38
39// read a chunk from bytes_
40void PngDecoder::ReadDataFromBuf(void* _png_ptr, uint8_t* dst, size_t size) {
41 png_structp png_ptr = (png_structp)_png_ptr;
42 PngDecoder* decoder = (PngDecoder*)(png_get_io_ptr(png_ptr));
43 assert(decoder);
44
45 if (decoder->cur_offset_ + size > decoder->NumBytes()) {
46 png_error(png_ptr, "PNG input buffer is incomplete");
47 return;
48 }
49
50 memcpy(dst, decoder->Bytes() + decoder->cur_offset_, size);
51 decoder->cur_offset_ += size;
52}
53
54bool PngDecoder::ReadHeader() {
55 bool result = false;
56
57 png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, 0, 0, 0);
58 assert(png_ptr);
59 png_infop info_ptr = png_create_info_struct(png_ptr);
60 png_infop end_info = png_create_info_struct(png_ptr);
61
62 assert(info_ptr);
63 assert(end_info);
64
65 png_ptr_ = png_ptr;
66 info_ptr_ = info_ptr;
67 end_info_ = end_info;
68
69 if (setjmp(png_jmpbuf(png_ptr)) == 0) {
70 png_set_read_fn(png_ptr, this, (png_rw_ptr)ReadDataFromBuf);
71
72 png_uint_32 wdth, hght;
73 int bit_depth, color_type, num_trans = 0;
74 png_bytep trans;
75 png_color_16p trans_values;
76
77 png_read_info(png_ptr, info_ptr);
78 png_get_IHDR(png_ptr, info_ptr, &wdth, &hght,
79 &bit_depth, &color_type, 0, 0, 0);
80
81 // set the shape based on what Decode will do. doesn't necessarily match the original image though as we're
82 // going to throw away any alpha channel and drop 16 bit depth to 8.
83 SetShape(static_cast<int>(hght), static_cast<int>(wdth), 3);
84
85 color_type_ = color_type;
86 bit_depth_ = bit_depth;
87
88 // TODO: Do we need to handle these combinations?
89 // We want to end up with 3 channel RGB though so they might not be relevant given we (I think) set the png
90 // code to throw away the alpha channel, reduce 16-bit to 8, and convert grayscale to 3 channel.
91 //
92 //
93 // if (bit_depth <= 8 || bit_depth == 16) {
94 // switch (color_type) {
95 // case PNG_COLOR_TYPE_RGB:
96 // case PNG_COLOR_TYPE_PALETTE:
97 // png_get_tRNS(png_ptr, info_ptr, &trans, &num_trans, &trans_values);
98 // if (num_trans > 0)
99 // m_type = CV_8UC4;
100 // else
101 // m_type = CV_8UC3;
102 // break;
103 // case PNG_COLOR_TYPE_GRAY_ALPHA:
104 // case PNG_COLOR_TYPE_RGB_ALPHA:
105 // m_type = CV_8UC4;
106 // break;
107 // default:
108 // m_type = CV_8UC1;
109 // }
110 // if (bit_depth == 16)
111 // m_type = CV_MAKETYPE(CV_16U, CV_MAT_CN(m_type));
112 //
113 //}
114 result = true;
115 }
116
117 return result;
118}
119
120bool PngDecoder::DecodeImpl(uint8_t* output, uint64_t out_bytes) {
121 bool result = false;
122
123 png_structp png_ptr = (png_structp)png_ptr_;
124 png_infop info_ptr = (png_infop)info_ptr_;
125 png_infop end_info = (png_infop)end_info_;
126
127 if (setjmp(png_jmpbuf(png_ptr)) == 0) {
128 if (bit_depth_ == 16) {
129 png_set_strip_16(png_ptr);
130 }
131
132 png_set_strip_alpha(png_ptr);
133
134 if (color_type_ == PNG_COLOR_TYPE_PALETTE) {
135 png_set_palette_to_rgb(png_ptr);
136 }
137
138 if ((color_type_ & PNG_COLOR_MASK_COLOR) == 0 && bit_depth_ < 8) {
139 png_set_expand_gray_1_2_4_to_8(png_ptr);
140 }
141
142 if (!(color_type_ & PNG_COLOR_MASK_COLOR)) {
143 png_set_gray_to_rgb(png_ptr); // Gray->RGB
144 }
145
146 png_set_interlace_handling(png_ptr);
147 png_read_update_info(png_ptr, info_ptr);
148
149 const auto& shape = Shape();
150 auto height = shape[0];
151 auto width = shape[1];
152 auto channels = shape[2];
153
154 std::vector<uint8_t*> row_pointers(height, nullptr);
155 auto row_size = png_get_rowbytes(png_ptr, info_ptr);
156 assert(row_size == width * channels); // check assumption
157
158 for (int row = 0; row < height; ++row) {
159 row_pointers[row] = output + (row * row_size);
160 }
161
162 png_read_image(png_ptr, row_pointers.data());
163 png_read_end(png_ptr, end_info);
164
165 result = true;
166 }
167
168 return result;
169}
170
171/////////////////////// PngEncoder ///////////////////
172
173void PngEncoder::WriteDataToBuf(void* _png_ptr, uint8_t* src, size_t size) {
174 if (size == 0)
175 return;
176
177 png_structp png_ptr = (png_structp)_png_ptr;
178 PngEncoder* encoder = (PngEncoder*)(png_get_io_ptr(png_ptr));
179 assert(encoder);
180
181 auto& buffer = encoder->Buffer();
182
183 if (encoder->cur_offset_ + size > buffer.size()) {
184 assert(false); // unexpected - this means no compression is occurring when writing
185 }
186
187 memcpy(buffer.data() + encoder->cur_offset_, src, size);
188
189 encoder->cur_offset_ += size;
190}
191
192void PngEncoder::FlushBuf(void* png_ptr) {
193 // we're writing to memory so this is a no-op
194}
195
196bool PngEncoder::EncodeImpl() {
197 bool result = false;
198
199 png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, 0, 0, 0);
200 assert(png_ptr);
201
202 png_infop info_ptr = png_create_info_struct(png_ptr);
203 assert(info_ptr);
204
205 const auto& shape = Shape();
206 const int height = static_cast<int>(shape[0]);
207 const int width = static_cast<int>(shape[1]);
208 const int channels = static_cast<int>(shape[2]);
209 const int depth = 8;
210
211 if (setjmp(png_jmpbuf(png_ptr)) == 0) {
212 png_set_write_fn(png_ptr, this, (png_rw_ptr)WriteDataToBuf, (png_flush_ptr)FlushBuf);
213
214 // tune parameters for speed
215 // (see http://wiki.linuxquestions.org/wiki/Libpng)
216 png_set_filter(png_ptr, PNG_FILTER_TYPE_BASE, PNG_FILTER_SUB);
217 png_set_compression_level(png_ptr, Z_BEST_SPEED);
218 png_set_compression_strategy(png_ptr, Z_DEFAULT_STRATEGY);
219
220 png_set_IHDR(png_ptr, info_ptr, width, height, depth, PNG_COLOR_TYPE_RGB,
221 PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT,
222 PNG_FILTER_TYPE_DEFAULT);
223
224 png_write_info(png_ptr, info_ptr);
225
226 std::vector<uint8_t*> row_pointers(height, nullptr);
227 auto row_size = png_get_rowbytes(png_ptr, info_ptr);
228 assert(row_size == width * channels); // check assumption
229
230 // need non-const to make png lib happy
231 uint8_t* orig_image_bytes = const_cast<uint8_t*>(Bytes());
232
233 for (int row = 0; row < height; ++row) {
234 row_pointers[row] = orig_image_bytes + (row * row_size);
235 }
236
237 png_write_image(png_ptr, row_pointers.data());
238 png_write_end(png_ptr, info_ptr);
239
240 Buffer().resize(cur_offset_); // remove unused bytes so size() is correct
241
242 result = true;
243 }
244
245 png_destroy_write_struct(&png_ptr, &info_ptr);
246
247 return result;
248}
249
250} // namespace ort_extensions
251