lavc/cbs: Add JPEG support
[ffmpeg.git] / libavcodec / cbs_jpeg.c
1 /*
2  * This file is part of FFmpeg.
3  *
4  * FFmpeg is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU Lesser General Public
6  * License as published by the Free Software Foundation; either
7  * version 2.1 of the License, or (at your option) any later version.
8  *
9  * FFmpeg is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
12  * Lesser General Public License for more details.
13  *
14  * You should have received a copy of the GNU Lesser General Public
15  * License along with FFmpeg; if not, write to the Free Software
16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
17  */
18
19 #include "cbs.h"
20 #include "cbs_internal.h"
21 #include "cbs_jpeg.h"
22
23
24 #define HEADER(name) do { \
25         ff_cbs_trace_header(ctx, name); \
26     } while (0)
27
28 #define CHECK(call) do { \
29         err = (call); \
30         if (err < 0) \
31             return err; \
32     } while (0)
33
34 #define SUBSCRIPTS(subs, ...) (subs > 0 ? ((int[subs + 1]){ subs, __VA_ARGS__ }) : NULL)
35
36 #define u(width, name, range_min, range_max) \
37     xu(width, name, range_min, range_max, 0)
38 #define us(width, name, sub, range_min, range_max) \
39     xu(width, name, range_min, range_max, 1, sub)
40
41
42 #define READ
43 #define READWRITE read
44 #define RWContext GetBitContext
45 #define FUNC(name) cbs_jpeg_read_ ## name
46
47 #define xu(width, name, range_min, range_max, subs, ...) do { \
48         uint32_t value = range_min; \
49         CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
50                                    SUBSCRIPTS(subs, __VA_ARGS__), \
51                                    &value, range_min, range_max)); \
52         current->name = value; \
53     } while (0)
54
55 #include "cbs_jpeg_syntax_template.c"
56
57 #undef READ
58 #undef READWRITE
59 #undef RWContext
60 #undef FUNC
61 #undef xu
62
63 #define WRITE
64 #define READWRITE write
65 #define RWContext PutBitContext
66 #define FUNC(name) cbs_jpeg_write_ ## name
67
68 #define xu(width, name, range_min, range_max, subs, ...) do { \
69         uint32_t value = current->name; \
70         CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
71                                     SUBSCRIPTS(subs, __VA_ARGS__), \
72                                     value, range_min, range_max)); \
73     } while (0)
74
75
76 #include "cbs_jpeg_syntax_template.c"
77
78 #undef READ
79 #undef READWRITE
80 #undef RWContext
81 #undef FUNC
82 #undef xu
83
84
85 static void cbs_jpeg_free_application_data(void *unit, uint8_t *content)
86 {
87     JPEGRawApplicationData *ad = (JPEGRawApplicationData*)content;
88     av_buffer_unref(&ad->Ap_ref);
89     av_freep(&content);
90 }
91
92 static void cbs_jpeg_free_comment(void *unit, uint8_t *content)
93 {
94     JPEGRawComment *comment = (JPEGRawComment*)content;
95     av_buffer_unref(&comment->Cm_ref);
96     av_freep(&content);
97 }
98
99 static void cbs_jpeg_free_scan(void *unit, uint8_t *content)
100 {
101     JPEGRawScan *scan = (JPEGRawScan*)content;
102     av_buffer_unref(&scan->data_ref);
103     av_freep(&content);
104 }
105
106 static int cbs_jpeg_split_fragment(CodedBitstreamContext *ctx,
107                                    CodedBitstreamFragment *frag,
108                                    int header)
109 {
110     AVBufferRef *data_ref;
111     uint8_t *data;
112     size_t data_size;
113     int unit, start, end, marker, next_start, next_marker;
114     int err, i, j, length;
115
116     if (frag->data_size < 4) {
117         // Definitely too short to be meaningful.
118         return AVERROR_INVALIDDATA;
119     }
120
121     for (i = 0; i + 1 < frag->data_size && frag->data[i] != 0xff; i++);
122     if (i > 0) {
123         av_log(ctx->log_ctx, AV_LOG_WARNING, "Discarding %d bytes at "
124                "beginning of image.\n", i);
125     }
126     for (++i; i + 1 < frag->data_size && frag->data[i] == 0xff; i++);
127     if (i + 1 >= frag->data_size && frag->data[i]) {
128         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
129                "no SOI marker found.\n");
130         return AVERROR_INVALIDDATA;
131     }
132     marker = frag->data[i];
133     if (marker != JPEG_MARKER_SOI) {
134         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: first "
135                "marker is %02x, should be SOI.\n", marker);
136         return AVERROR_INVALIDDATA;
137     }
138     for (++i; i + 1 < frag->data_size && frag->data[i] == 0xff; i++);
139     if (i + 1 >= frag->data_size) {
140         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
141                "no image content found.\n");
142         return AVERROR_INVALIDDATA;
143     }
144     marker = frag->data[i];
145     start  = i + 1;
146
147     for (unit = 0;; unit++) {
148         if (marker == JPEG_MARKER_EOI) {
149             break;
150         } else if (marker == JPEG_MARKER_SOS) {
151             for (i = start; i + 1 < frag->data_size; i++) {
152                 if (frag->data[i] != 0xff)
153                     continue;
154                 end = i;
155                 for (++i; i + 1 < frag->data_size &&
156                           frag->data[i] == 0xff; i++);
157                 if (i + 1 >= frag->data_size) {
158                     next_marker = -1;
159                 } else {
160                     if (frag->data[i] == 0x00)
161                         continue;
162                     next_marker = frag->data[i];
163                     next_start  = i + 1;
164                 }
165                 break;
166             }
167         } else {
168             i = start;
169             if (i + 2 > frag->data_size) {
170                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
171                        "truncated at %02x marker.\n", marker);
172                 return AVERROR_INVALIDDATA;
173             }
174             length = AV_RB16(frag->data + i);
175             if (i + length > frag->data_size) {
176                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
177                        "truncated at %02x marker segment.\n", marker);
178                 return AVERROR_INVALIDDATA;
179             }
180             end = start + length;
181
182             i = end;
183             if (frag->data[i] != 0xff) {
184                 next_marker = -1;
185             } else {
186                 for (++i; i + 1 < frag->data_size &&
187                           frag->data[i] == 0xff; i++);
188                 if (i + 1 >= frag->data_size) {
189                     next_marker = -1;
190                 } else {
191                     next_marker = frag->data[i];
192                     next_start  = i + 1;
193                 }
194             }
195         }
196
197         if (marker == JPEG_MARKER_SOS) {
198             length = AV_RB16(frag->data + start);
199
200             data_ref = NULL;
201             data     = av_malloc(end - start +
202                                  AV_INPUT_BUFFER_PADDING_SIZE);
203             if (!data)
204                 return AVERROR(ENOMEM);
205
206             memcpy(data, frag->data + start, length);
207             for (i = start + length, j = length; i < end; i++, j++) {
208                 if (frag->data[i] == 0xff) {
209                     while (frag->data[i] == 0xff)
210                         ++i;
211                     data[j] = 0xff;
212                 } else {
213                     data[j] = frag->data[i];
214                 }
215             }
216             data_size = j;
217
218             memset(data + data_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
219
220         } else {
221             data      = frag->data + start;
222             data_size = end - start;
223             data_ref  = frag->data_ref;
224         }
225
226         err = ff_cbs_insert_unit_data(ctx, frag, unit, marker,
227                                       data, data_size, data_ref);
228         if (err < 0) {
229             if (!data_ref)
230                 av_freep(&data);
231             return err;
232         }
233
234         if (next_marker == -1)
235             break;
236         marker = next_marker;
237         start  = next_start;
238     }
239
240     return 0;
241 }
242
243 static int cbs_jpeg_read_unit(CodedBitstreamContext *ctx,
244                               CodedBitstreamUnit *unit)
245 {
246     GetBitContext gbc;
247     int err;
248
249     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
250     if (err < 0)
251         return err;
252
253     if (unit->type >= JPEG_MARKER_SOF0 &&
254         unit->type <= JPEG_MARKER_SOF3) {
255         err = ff_cbs_alloc_unit_content(ctx, unit,
256                                         sizeof(JPEGRawFrameHeader),
257                                         NULL);
258         if (err < 0)
259             return err;
260
261         err = cbs_jpeg_read_frame_header(ctx, &gbc, unit->content);
262         if (err < 0)
263             return err;
264
265     } else if (unit->type >= JPEG_MARKER_APPN &&
266                unit->type <= JPEG_MARKER_APPN + 15) {
267         err = ff_cbs_alloc_unit_content(ctx, unit,
268                                         sizeof(JPEGRawApplicationData),
269                                         &cbs_jpeg_free_application_data);
270         if (err < 0)
271             return err;
272
273         err = cbs_jpeg_read_application_data(ctx, &gbc, unit->content);
274         if (err < 0)
275             return err;
276
277     } else if (unit->type == JPEG_MARKER_SOS) {
278         JPEGRawScan *scan;
279         int pos;
280
281         err = ff_cbs_alloc_unit_content(ctx, unit,
282                                         sizeof(JPEGRawScan),
283                                         &cbs_jpeg_free_scan);
284         if (err < 0)
285             return err;
286         scan = unit->content;
287
288         err = cbs_jpeg_read_scan_header(ctx, &gbc, &scan->header);
289         if (err < 0)
290             return err;
291
292         pos = get_bits_count(&gbc);
293         av_assert0(pos % 8 == 0);
294         if (pos > 0) {
295             scan->data_size = unit->data_size - pos / 8;
296             scan->data_ref  = av_buffer_ref(unit->data_ref);
297             if (!scan->data_ref)
298                 return AVERROR(ENOMEM);
299             scan->data = unit->data + pos / 8;
300         }
301
302     } else {
303         switch (unit->type) {
304 #define SEGMENT(marker, type, func, free) \
305         case JPEG_MARKER_ ## marker: \
306             { \
307                 err = ff_cbs_alloc_unit_content(ctx, unit, \
308                                                 sizeof(type), free); \
309                 if (err < 0) \
310                     return err; \
311                 err = cbs_jpeg_read_ ## func(ctx, &gbc, unit->content); \
312                 if (err < 0) \
313                     return err; \
314             } \
315             break
316             SEGMENT(DQT, JPEGRawQuantisationTableSpecification, dqt, NULL);
317             SEGMENT(DHT, JPEGRawHuffmanTableSpecification,      dht, NULL);
318             SEGMENT(COM, JPEGRawComment,  comment, &cbs_jpeg_free_comment);
319 #undef SEGMENT
320         default:
321             return AVERROR(ENOSYS);
322         }
323     }
324
325     return 0;
326 }
327
328 static int cbs_jpeg_write_scan(CodedBitstreamContext *ctx,
329                                CodedBitstreamUnit *unit,
330                                PutBitContext *pbc)
331 {
332     JPEGRawScan *scan = unit->content;
333     int i, err;
334
335     err = cbs_jpeg_write_scan_header(ctx, pbc, &scan->header);
336     if (err < 0)
337         return err;
338
339     if (scan->data) {
340         if (scan->data_size * 8 > put_bits_left(pbc))
341             return AVERROR(ENOSPC);
342
343         for (i = 0; i < scan->data_size; i++)
344             put_bits(pbc, 8, scan->data[i]);
345     }
346
347     return 0;
348 }
349
350 static int cbs_jpeg_write_segment(CodedBitstreamContext *ctx,
351                                   CodedBitstreamUnit *unit,
352                                   PutBitContext *pbc)
353 {
354     int err;
355
356     if (unit->type >= JPEG_MARKER_SOF0 &&
357         unit->type <= JPEG_MARKER_SOF3) {
358         err = cbs_jpeg_write_frame_header(ctx, pbc, unit->content);
359     } else if (unit->type >= JPEG_MARKER_APPN &&
360                unit->type <= JPEG_MARKER_APPN + 15) {
361         err = cbs_jpeg_write_application_data(ctx, pbc, unit->content);
362     } else {
363         switch (unit->type) {
364 #define SEGMENT(marker, func) \
365             case JPEG_MARKER_ ## marker: \
366                 err = cbs_jpeg_write_ ## func(ctx, pbc, unit->content); \
367                 break;
368             SEGMENT(DQT, dqt);
369             SEGMENT(DHT, dht);
370             SEGMENT(COM, comment);
371         default:
372             return AVERROR_PATCHWELCOME;
373         }
374     }
375
376     return err;
377 }
378
379 static int cbs_jpeg_write_unit(CodedBitstreamContext *ctx,
380                                 CodedBitstreamUnit *unit)
381 {
382     CodedBitstreamJPEGContext *priv = ctx->priv_data;
383     PutBitContext pbc;
384     int err;
385
386     if (!priv->write_buffer) {
387         // Initial write buffer size is 1MB.
388         priv->write_buffer_size = 1024 * 1024;
389
390     reallocate_and_try_again:
391         err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
392         if (err < 0) {
393             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
394                    "sufficiently large write buffer (last attempt "
395                    "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
396             return err;
397         }
398     }
399
400     init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
401
402     if (unit->type == JPEG_MARKER_SOS)
403         err = cbs_jpeg_write_scan(ctx, unit, &pbc);
404     else
405         err = cbs_jpeg_write_segment(ctx, unit, &pbc);
406
407     if (err == AVERROR(ENOSPC)) {
408         // Overflow.
409         priv->write_buffer_size *= 2;
410         goto reallocate_and_try_again;
411     }
412     if (err < 0) {
413         // Write failed for some other reason.
414         return err;
415     }
416
417     if (put_bits_count(&pbc) % 8)
418         unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
419     else
420         unit->data_bit_padding = 0;
421
422     unit->data_size = (put_bits_count(&pbc) + 7) / 8;
423     flush_put_bits(&pbc);
424
425     err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
426     if (err < 0)
427         return err;
428
429     memcpy(unit->data, priv->write_buffer, unit->data_size);
430
431     return 0;
432 }
433
434 static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
435                                        CodedBitstreamFragment *frag)
436 {
437     const CodedBitstreamUnit *unit;
438     uint8_t *data;
439     size_t size, dp, sp;
440     int i;
441
442     size = 4; // SOI + EOI.
443     for (i = 0; i < frag->nb_units; i++) {
444         unit = &frag->units[i];
445         size += 2 + unit->data_size;
446         if (unit->type == JPEG_MARKER_SOS) {
447             for (sp = 0; sp < unit->data_size; sp++) {
448                 if (unit->data[sp] == 0xff)
449                     ++size;
450             }
451         }
452     }
453
454     frag->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
455     if (!frag->data_ref)
456         return AVERROR(ENOMEM);
457     data = frag->data_ref->data;
458
459     dp = 0;
460
461     data[dp++] = 0xff;
462     data[dp++] = JPEG_MARKER_SOI;
463
464     for (i = 0; i < frag->nb_units; i++) {
465         unit = &frag->units[i];
466
467         data[dp++] = 0xff;
468         data[dp++] = unit->type;
469
470         if (unit->type != JPEG_MARKER_SOS) {
471             memcpy(data + dp, unit->data, unit->data_size);
472             dp += unit->data_size;
473         } else {
474             sp = AV_RB16(unit->data);
475             av_assert0(sp <= unit->data_size);
476             memcpy(data + dp, unit->data, sp);
477             dp += sp;
478
479             for (; sp < unit->data_size; sp++) {
480                 if (unit->data[sp] == 0xff) {
481                     data[dp++] = 0xff;
482                     data[dp++] = 0x00;
483                 } else {
484                     data[dp++] = unit->data[sp];
485                 }
486             }
487         }
488     }
489
490     data[dp++] = 0xff;
491     data[dp++] = JPEG_MARKER_EOI;
492
493     av_assert0(dp == size);
494
495     memset(data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
496     frag->data      = data;
497     frag->data_size = size;
498
499     return 0;
500 }
501
502 static void cbs_jpeg_close(CodedBitstreamContext *ctx)
503 {
504     CodedBitstreamJPEGContext *priv = ctx->priv_data;
505
506     av_freep(&priv->write_buffer);
507 }
508
509 const CodedBitstreamType ff_cbs_type_jpeg = {
510     .codec_id          = AV_CODEC_ID_MJPEG,
511
512     .priv_data_size    = sizeof(CodedBitstreamJPEGContext),
513
514     .split_fragment    = &cbs_jpeg_split_fragment,
515     .read_unit         = &cbs_jpeg_read_unit,
516     .write_unit        = &cbs_jpeg_write_unit,
517     .assemble_fragment = &cbs_jpeg_assemble_fragment,
518
519     .close             = &cbs_jpeg_close,
520 };