1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.h"
16
17 #include <cstdint>
18 #include <memory>
19 #include <string>
20
21 #include "tensorflow/lite/c/c_api_types.h"
22 #include "tensorflow/lite/c/common.h"
23 #include "tensorflow/lite/kernels/internal/compatibility.h"
24 #include "tensorflow/lite/minimal_logging.h"
25 #include "tensorflow/lite/string_util.h"
26
27 namespace tflite {
28 namespace acceleration {
29 namespace decode_jpeg_kernel {
30
31 /*
32
33 JPEG file overall file structure
34
35 SOI Marker FFD8
36 Marker XX size=SSSS FFXX SSSS DDDD......
37 Marker YY size=TTTT FFYY TTTT DDDD......
38 SOFn marker with the info we want
39 SOS Marker size=UUUU FFDA UUUU DDDD....
40 Image stream I I I I....
41 EOI Marker FFD9
42
43 The first marker is either APP0 (JFIF format) or APP1 (EXIF format)
44
45 We support only JFIF images
46 */
47
48 namespace {
49
50 using MarkerId = uint16_t;
51
AsWord(int value,char * msb,char * lsb)52 void AsWord(int value, char* msb, char* lsb) {
53 *msb = static_cast<char>(value >> 8);
54 *lsb = static_cast<char>(value);
55 }
56
57 // JFIF spec at
58 // https://www.ecma-international.org/publications-and-standards/technical-reports/ecma-tr-98/
59 // Marker definition summary at
60 // http://lad.dsc.ufcg.edu.br/multimidia/jpegmarker.pdf
61 // Overall JPEG File structure with discussion of the supported number of
62 // channels per format
63 // https://docs.oracle.com/javase/8/docs/api/javax/imageio/metadata/doc-files/jpeg_metadata.html
64 //
65
66 class JfifHeaderParser {
67 public:
JfifHeaderParser(const tflite::StringRef & jpeg_image_data)68 explicit JfifHeaderParser(const tflite::StringRef& jpeg_image_data)
69 : jpeg_image_data_(jpeg_image_data), offset_(0) {
70 if (!IsJpegImage(jpeg_image_data_)) {
71 is_valid_image_buffer_ = false;
72 validation_error_message_ = "Not a valid JPEG image.";
73 } else if (!IsJfifImage(jpeg_image_data_)) {
74 is_valid_image_buffer_ = false;
75 validation_error_message_ = "Image is not in JFIF format.";
76 return;
77 } else {
78 is_valid_image_buffer_ = true;
79 }
80 }
81
82 #define ENSURE_READ_STATUS(a) \
83 do { \
84 const TfLiteStatus s = (a); \
85 if (s != kTfLiteOk) { \
86 return {s, "Error trying to parse JPEG header."}; \
87 } \
88 } while (0)
89
ReadJpegHeader(JpegHeader * result)90 Status ReadJpegHeader(JpegHeader* result) {
91 if (!is_valid_image_buffer_) {
92 return {kTfLiteError, validation_error_message_};
93 }
94
95 Status move_to_sof_status = MoveToStartOfFrameMarker();
96 if (move_to_sof_status.code != kTfLiteOk) {
97 return move_to_sof_status;
98 }
99
100 ENSURE_READ_STATUS(SkipBytes(2)); // skipping marker length
101 char precision;
102 ENSURE_READ_STATUS(ReadByte(&precision));
103 uint16_t height;
104 ENSURE_READ_STATUS(ReadWord(&height));
105 uint16_t width;
106 ENSURE_READ_STATUS(ReadWord(&width));
107 char num_of_components;
108 ENSURE_READ_STATUS(ReadByte(&num_of_components));
109
110 if (num_of_components != 1 && num_of_components != 3) {
111 return {kTfLiteError,
112 "A JFIF image without App14 marker doesn't support a number of "
113 "components = " +
114 std::to_string(static_cast<int>(num_of_components))};
115 }
116
117 result->width = width;
118 result->height = height;
119 result->channels = num_of_components;
120 result->bits_per_sample = precision;
121
122 return {kTfLiteOk, ""};
123 }
124
ApplyHeaderToImage(const JpegHeader & new_header,std::string & write_to)125 Status ApplyHeaderToImage(const JpegHeader& new_header,
126 std::string& write_to) {
127 if (!is_valid_image_buffer_) {
128 return {kTfLiteError, validation_error_message_};
129 }
130
131 Status move_to_sof_status = MoveToStartOfFrameMarker();
132 if (move_to_sof_status.code != kTfLiteOk) {
133 return move_to_sof_status;
134 }
135 ENSURE_READ_STATUS(SkipBytes(2)); // skipping marker length
136
137 if (!HasData(6)) {
138 return {kTfLiteError,
139 "Invalid SOF marker, image buffer ends before end of marker"};
140 }
141
142 char header[6];
143 header[0] = static_cast<char>(new_header.bits_per_sample);
144 AsWord(new_header.height, header + 1, header + 2);
145 AsWord(new_header.width, header + 3, header + 4);
146 header[5] = static_cast<char>(new_header.channels);
147
148 write_to.clear();
149 write_to.append(jpeg_image_data_.str, offset_);
150 write_to.append(header, 6);
151
152 ENSURE_READ_STATUS(SkipBytes(6));
153 if (HasData()) {
154 write_to.append(jpeg_image_data_.str + offset_,
155 jpeg_image_data_.len - offset_);
156 }
157
158 return {kTfLiteOk, ""};
159 }
160
161 private:
162 const tflite::StringRef jpeg_image_data_;
163 // Using int for consistency with the size in StringRef
164 int offset_;
165 bool is_valid_image_buffer_;
166 std::string validation_error_message_;
167
168 // Moves to the begin of the first SOF marker
MoveToStartOfFrameMarker()169 Status MoveToStartOfFrameMarker() {
170 const MarkerId kStartOfStreamMarkerId = 0xFFDA; // Start of image data
171
172 offset_ = 0;
173 ENSURE_READ_STATUS(SkipBytes(4)); // skipping SOI and APP0 marker IDs
174 ENSURE_READ_STATUS(SkipCurrentMarker()); // skipping APP0
175 MarkerId curr_marker_id;
176 // We need at least 2 bytes for the marker ID and 2 for the length
177 while (HasData(/*min_data_size=*/4)) {
178 ENSURE_READ_STATUS(ReadWord(&curr_marker_id));
179 // We are breaking at the first met SOF marker. This won't generate
180 // results inconsistent with LibJPEG because only
181 // image with a single SOF marker are successfully parsed by it.
182 // LibJPEG fails if more than one marker is found in the header (see
183 // https://github.com/libjpeg-turbo/libjpeg-turbo/blob/main/jerror.h#L121
184 // and
185 // https://github.com/libjpeg-turbo/libjpeg-turbo/blob/main/jdmarker.c#L264-L265
186 if (IsStartOfFrameMarkerId(curr_marker_id)) {
187 break;
188 }
189 if (curr_marker_id == kStartOfStreamMarkerId) {
190 return {kTfLiteError, "Error trying to parse JPEG header."};
191 }
192 ENSURE_READ_STATUS(SkipCurrentMarker());
193 }
194
195 return {kTfLiteOk, ""};
196 }
197
198 #undef ENSURE_READ_STATUS
199
HasData(int min_data_size=1)200 bool HasData(int min_data_size = 1) {
201 return offset_ <= jpeg_image_data_.len - min_data_size;
202 }
203
SkipBytes(int bytes)204 TfLiteStatus SkipBytes(int bytes) {
205 if (!HasData(bytes)) {
206 TFLITE_LOG(TFLITE_LOG_WARNING,
207 "Trying to move out of image boundaries from offset %d, "
208 "skipping %d bytes",
209 offset_, bytes);
210 return kTfLiteError;
211 }
212
213 offset_ += bytes;
214
215 return kTfLiteOk;
216 }
217
ReadByte(char * result)218 TfLiteStatus ReadByte(char* result) {
219 if (!HasData()) {
220 return kTfLiteError;
221 }
222
223 *result = jpeg_image_data_.str[offset_];
224
225 return SkipBytes(1);
226 }
227
ReadWord(uint16_t * result)228 TfLiteStatus ReadWord(uint16_t* result) {
229 TF_LITE_ENSURE_STATUS(ReadWordAt(jpeg_image_data_, offset_, result));
230 return SkipBytes(2);
231 }
232
SkipCurrentMarker()233 TfLiteStatus SkipCurrentMarker() {
234 // We just read the marker ID so we are on top of the marker len
235 uint16_t full_marker_len;
236 TF_LITE_ENSURE_STATUS(ReadWord(&full_marker_len));
237 if (full_marker_len <= 2) {
238 TFLITE_LOG(TFLITE_LOG_WARNING,
239 "Invalid marker length %d read at offset %X", full_marker_len,
240 offset_);
241 return kTfLiteError;
242 }
243
244 // The marker len includes the 2 bytes of marker length
245 return SkipBytes(full_marker_len - 2);
246 }
247
ReadWordAt(const tflite::StringRef & jpeg_image_data,int read_offset,uint16_t * result)248 static TfLiteStatus ReadWordAt(const tflite::StringRef& jpeg_image_data,
249 int read_offset, uint16_t* result) {
250 if (read_offset < 0 || read_offset > jpeg_image_data.len - 2) {
251 return kTfLiteError;
252 }
253 // Cast to unsigned since char can be signed.
254 const unsigned char* buf =
255 reinterpret_cast<const unsigned char*>(jpeg_image_data.str);
256
257 *result = (buf[read_offset] << 8) + buf[read_offset + 1];
258
259 return kTfLiteOk;
260 }
261
IsJpegImage(const tflite::StringRef & jpeg_image_data)262 static bool IsJpegImage(const tflite::StringRef& jpeg_image_data) {
263 const MarkerId kStartOfImageMarkerId = 0xFFD8;
264 const MarkerId kEndOfImageMarkerId = 0xFFD9;
265
266 MarkerId soi_marker_id;
267 MarkerId eoi_marker_id;
268 if (ReadWordAt(jpeg_image_data, 0, &soi_marker_id) != kTfLiteOk) {
269 return false;
270 }
271 if (ReadWordAt(jpeg_image_data, jpeg_image_data.len - 2, &eoi_marker_id) !=
272 kTfLiteOk) {
273 return false;
274 }
275
276 return (soi_marker_id == kStartOfImageMarkerId) &&
277 (eoi_marker_id == kEndOfImageMarkerId);
278 }
279
IsJfifImage(const tflite::StringRef & jpeg_image_data)280 static bool IsJfifImage(const tflite::StringRef& jpeg_image_data) {
281 const MarkerId kApp0MarkerId = 0xFFE0; // First marker in JIFF image
282
283 MarkerId app_marker_id;
284 if ((ReadWordAt(jpeg_image_data, 2, &app_marker_id) != kTfLiteOk) ||
285 (app_marker_id != kApp0MarkerId)) {
286 return false;
287 }
288
289 // Checking Jfif identifier string "JFIF\0" in APP0 Marker
290 const std::string kJfifIdString{"JFIF\0", 5};
291
292 // The ID starts after SOI (2 bytes), APP0 marker IDs (2 bytes) and 2 other
293 // bytes with APP0 marker length
294 const int KJfifIdStringStartOffset = 6;
295
296 if (KJfifIdStringStartOffset + kJfifIdString.size() >=
297 jpeg_image_data.len) {
298 TFLITE_LOG(TFLITE_LOG_WARNING,
299 "Invalid image, reached end of data at offset while "
300 "parsing APP0 header");
301 return false;
302 }
303
304 const std::string actualImgId(
305 jpeg_image_data.str + KJfifIdStringStartOffset, kJfifIdString.size());
306 if (kJfifIdString != actualImgId) {
307 TFLITE_LOG(TFLITE_LOG_WARNING, "Invalid image, invalid APP0 header");
308
309 return false;
310 }
311
312 return true;
313 }
314
IsStartOfFrameMarkerId(MarkerId marker_id)315 static bool IsStartOfFrameMarkerId(MarkerId marker_id) {
316 return 0xFFC0 <= marker_id && marker_id < 0xFFCF;
317 }
318 };
319
320 } // namespace
ReadJpegHeader(const tflite::StringRef & jpeg_image_data,JpegHeader * header)321 Status ReadJpegHeader(const tflite::StringRef& jpeg_image_data,
322 JpegHeader* header) {
323 JfifHeaderParser parser(jpeg_image_data);
324
325 return parser.ReadJpegHeader(header);
326 }
327
BuildImageWithNewHeader(const tflite::StringRef & orig_jpeg_image_data,const JpegHeader & new_header,std::string & new_image_data)328 Status BuildImageWithNewHeader(const tflite::StringRef& orig_jpeg_image_data,
329 const JpegHeader& new_header,
330 std::string& new_image_data) {
331 JfifHeaderParser parser(orig_jpeg_image_data);
332
333 return parser.ApplyHeaderToImage(new_header, new_image_data);
334 }
335
336 } // namespace decode_jpeg_kernel
337 } // namespace acceleration
338 } // namespace tflite
339