lavc/libaribb24: protect handled value with parenthesis in RGB_TO_BGR
[ffmpeg.git] / libavcodec / vaapi_encode_mjpeg.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 <va/va.h>
20 #include <va/va_enc_jpeg.h>
21
22 #include "libavutil/avassert.h"
23 #include "libavutil/common.h"
24 #include "libavutil/internal.h"
25 #include "libavutil/opt.h"
26 #include "libavutil/pixdesc.h"
27
28 #include "avcodec.h"
29 #include "bytestream.h"
30 #include "cbs.h"
31 #include "cbs_jpeg.h"
32 #include "internal.h"
33 #include "jpegtables.h"
34 #include "mjpeg.h"
35 #include "put_bits.h"
36 #include "vaapi_encode.h"
37
38
39 // Standard JPEG quantisation tables, in zigzag order.
40 static const unsigned char vaapi_encode_mjpeg_quant_luminance[64] = {
41     16,  11,  12,  14,  12,  10,  16,  14,
42     13,  14,  18,  17,  16,  19,  24,  40,
43     26,  24,  22,  22,  24,  49,  35,  37,
44     29,  40,  58,  51,  61,  60,  57,  51,
45     56,  55,  64,  72,  92,  78,  64,  68,
46     87,  69,  55,  56,  80, 109,  81,  87,
47     95,  98, 103, 104, 103,  62,  77, 113,
48    121, 112, 100, 120,  92, 101, 103,  99,
49 };
50 static const unsigned char vaapi_encode_mjpeg_quant_chrominance[64] = {
51     17,  18,  18,  24,  21,  24,  47,  26,
52     26,  47,  99,  66,  56,  66,  99,  99,
53     99,  99,  99,  99,  99,  99,  99,  99,
54     99,  99,  99,  99,  99,  99,  99,  99,
55     99,  99,  99,  99,  99,  99,  99,  99,
56     99,  99,  99,  99,  99,  99,  99,  99,
57     99,  99,  99,  99,  99,  99,  99,  99,
58     99,  99,  99,  99,  99,  99,  99,  99,
59 };
60
61 typedef struct VAAPIEncodeMJPEGContext {
62     VAAPIEncodeContext common;
63
64     // User options.
65     int jfif;
66     int huffman;
67
68     // Derived settings.
69     int quality;
70     uint8_t jfif_data[14];
71
72     // Writer structures.
73     JPEGRawFrameHeader     frame_header;
74     JPEGRawScan            scan;
75     JPEGRawApplicationData jfif_header;
76     JPEGRawQuantisationTableSpecification quant_tables;
77     JPEGRawHuffmanTableSpecification      huffman_tables;
78
79     CodedBitstreamContext *cbc;
80     CodedBitstreamFragment current_fragment;
81 } VAAPIEncodeMJPEGContext;
82
83 static int vaapi_encode_mjpeg_write_image_header(AVCodecContext *avctx,
84                                                  VAAPIEncodePicture *pic,
85                                                  VAAPIEncodeSlice *slice,
86                                                  char *data, size_t *data_len)
87 {
88     VAAPIEncodeMJPEGContext *priv = avctx->priv_data;
89     CodedBitstreamFragment  *frag = &priv->current_fragment;
90     int err;
91
92     if (priv->jfif) {
93         err = ff_cbs_insert_unit_content(priv->cbc, frag, -1,
94                                          JPEG_MARKER_APPN + 0,
95                                          &priv->jfif_header, NULL);
96         if (err < 0)
97             goto fail;
98     }
99
100     err = ff_cbs_insert_unit_content(priv->cbc, frag, -1,
101                                      JPEG_MARKER_DQT,
102                                      &priv->quant_tables, NULL);
103     if (err < 0)
104         goto fail;
105
106     err = ff_cbs_insert_unit_content(priv->cbc, frag, -1,
107                                      JPEG_MARKER_SOF0,
108                                      &priv->frame_header, NULL);
109     if (err < 0)
110         goto fail;
111
112     if (priv->huffman) {
113         err = ff_cbs_insert_unit_content(priv->cbc, frag, -1,
114                                          JPEG_MARKER_DHT,
115                                          &priv->huffman_tables, NULL);
116         if (err < 0)
117             goto fail;
118     }
119
120     err = ff_cbs_insert_unit_content(priv->cbc, frag, -1,
121                                      JPEG_MARKER_SOS,
122                                      &priv->scan, NULL);
123     if (err < 0)
124         goto fail;
125
126     err = ff_cbs_write_fragment_data(priv->cbc, frag);
127     if (err < 0) {
128         av_log(avctx, AV_LOG_ERROR, "Failed to write image header.\n");
129         goto fail;
130     }
131
132     if (*data_len < 8 * frag->data_size) {
133         av_log(avctx, AV_LOG_ERROR, "Image header too large: "
134                "%zu < %zu.\n", *data_len, 8 * frag->data_size);
135         err = AVERROR(ENOSPC);
136         goto fail;
137     }
138
139     // Remove the EOI at the end of the fragment.
140     memcpy(data, frag->data, frag->data_size - 2);
141     *data_len = 8 * (frag->data_size - 2);
142
143     err = 0;
144 fail:
145     ff_cbs_fragment_uninit(priv->cbc, frag);
146     return err;
147 }
148
149 static int vaapi_encode_mjpeg_write_extra_buffer(AVCodecContext *avctx,
150                                                  VAAPIEncodePicture *pic,
151                                                  int index, int *type,
152                                                  char *data, size_t *data_len)
153 {
154     VAAPIEncodeMJPEGContext *priv = avctx->priv_data;
155     int t, i, k;
156
157     if (index == 0) {
158         // Write quantisation tables.
159         JPEGRawFrameHeader                     *fh = &priv->frame_header;
160         JPEGRawQuantisationTableSpecification *dqt = &priv->quant_tables;
161         VAQMatrixBufferJPEG *quant;
162
163         if (*data_len < sizeof(*quant))
164             return AVERROR(ENOSPC);
165         *type     = VAQMatrixBufferType;
166         *data_len = sizeof(*quant);
167
168         quant = (VAQMatrixBufferJPEG*)data;
169         memset(quant, 0, sizeof(*quant));
170
171         quant->load_lum_quantiser_matrix = 1;
172         for (i = 0; i < 64; i++)
173             quant->lum_quantiser_matrix[i] = dqt->table[fh->Tq[0]].Q[i];
174
175         if (fh->Nf > 1) {
176             quant->load_chroma_quantiser_matrix = 1;
177             for (i = 0; i < 64; i++)
178                 quant->chroma_quantiser_matrix[i] =
179                     dqt->table[fh->Tq[1]].Q[i];
180         }
181
182     } else if (index == 1) {
183         // Write huffman tables.
184         JPEGRawScanHeader                 *sh = &priv->scan.header;
185         JPEGRawHuffmanTableSpecification *dht = &priv->huffman_tables;
186         VAHuffmanTableBufferJPEGBaseline *huff;
187
188         if (*data_len < sizeof(*huff))
189             return AVERROR(ENOSPC);
190         *type     = VAHuffmanTableBufferType;
191         *data_len = sizeof(*huff);
192
193         huff = (VAHuffmanTableBufferJPEGBaseline*)data;
194         memset(huff, 0, sizeof(*huff));
195
196         for (t = 0; t < 1 + (sh->Ns > 1); t++) {
197             const JPEGRawHuffmanTable *ht;
198
199             huff->load_huffman_table[t] = 1;
200
201             ht = &dht->table[2 * t];
202             for (i = k = 0; i < 16; i++)
203                 k += (huff->huffman_table[t].num_dc_codes[i] = ht->L[i]);
204             av_assert0(k <= sizeof(huff->huffman_table[t].dc_values));
205             for (i = 0; i < k; i++)
206                 huff->huffman_table[t].dc_values[i] = ht->V[i];
207
208             ht = &dht->table[2 * t + 1];
209             for (i = k = 0; i < 16; i++)
210                 k += (huff->huffman_table[t].num_ac_codes[i] = ht->L[i]);
211             av_assert0(k <= sizeof(huff->huffman_table[t].ac_values));
212             for (i = 0; i < k; i++)
213                 huff->huffman_table[t].ac_values[i] = ht->V[i];
214         }
215
216     } else {
217         return AVERROR_EOF;
218     }
219     return 0;
220 }
221
222 static int vaapi_encode_mjpeg_init_picture_params(AVCodecContext *avctx,
223                                                   VAAPIEncodePicture *pic)
224 {
225     VAAPIEncodeMJPEGContext         *priv = avctx->priv_data;
226     JPEGRawFrameHeader                *fh = &priv->frame_header;
227     JPEGRawScanHeader                 *sh = &priv->scan.header;
228     VAEncPictureParameterBufferJPEG *vpic = pic->codec_picture_params;
229     const AVPixFmtDescriptor *desc;
230     const uint8_t *components;
231     int t, i, quant_scale, len;
232
233     av_assert0(pic->type == PICTURE_TYPE_IDR);
234
235     desc = av_pix_fmt_desc_get(priv->common.input_frames->sw_format);
236     av_assert0(desc);
237     if (desc->flags & AV_PIX_FMT_FLAG_RGB)
238         components = (uint8_t[3]) { 'R', 'G', 'B' };
239     else
240         components = (uint8_t[3]) {  1,   2,   3  };
241
242     // Frame header.
243
244     fh->P  = 8;
245     fh->Y  = avctx->height;
246     fh->X  = avctx->width;
247     fh->Nf = desc->nb_components;
248
249     for (i = 0; i < fh->Nf; i++) {
250         fh->C[i] = components[i];
251         fh->H[i] = 1 + (i == 0 ? desc->log2_chroma_w : 0);
252         fh->V[i] = 1 + (i == 0 ? desc->log2_chroma_h : 0);
253
254         fh->Tq[i] = !!i;
255     }
256
257     fh->Lf = 8 + 3 * fh->Nf;
258
259     // JFIF header.
260     if (priv->jfif) {
261         JPEGRawApplicationData *app = &priv->jfif_header;
262         AVRational sar = pic->input_image->sample_aspect_ratio;
263         int sar_w, sar_h;
264         PutByteContext pbc;
265
266         bytestream2_init_writer(&pbc, priv->jfif_data,
267                                 sizeof(priv->jfif_data));
268
269         bytestream2_put_buffer(&pbc, "JFIF", 5);
270         bytestream2_put_be16(&pbc, 0x0102);
271         bytestream2_put_byte(&pbc, 0);
272
273         av_reduce(&sar_w, &sar_h, sar.num, sar.den, 65535);
274         if (sar_w && sar_h) {
275             bytestream2_put_be16(&pbc, sar_w);
276             bytestream2_put_be16(&pbc, sar_h);
277         } else {
278             bytestream2_put_be16(&pbc, 1);
279             bytestream2_put_be16(&pbc, 1);
280         }
281
282         bytestream2_put_byte(&pbc, 0);
283         bytestream2_put_byte(&pbc, 0);
284
285         av_assert0(bytestream2_get_bytes_left_p(&pbc) == 0);
286
287         app->Lp     = 2 + sizeof(priv->jfif_data);
288         app->Ap     = priv->jfif_data;
289         app->Ap_ref = NULL;
290     }
291
292     // Quantisation tables.
293
294     if (priv->quality < 50)
295         quant_scale = 5000 / priv->quality;
296     else
297         quant_scale = 200 - 2 * priv->quality;
298
299     len = 2;
300
301     for (t = 0; t < 1 + (fh->Nf > 1); t++) {
302         JPEGRawQuantisationTable *quant = &priv->quant_tables.table[t];
303         const uint8_t *data = t == 0 ?
304             vaapi_encode_mjpeg_quant_luminance :
305             vaapi_encode_mjpeg_quant_chrominance;
306
307         quant->Pq = 0;
308         quant->Tq = t;
309         for (i = 0; i < 64; i++)
310             quant->Q[i] = av_clip(data[i] * quant_scale / 100, 1, 255);
311
312         len += 65;
313     }
314
315     priv->quant_tables.Lq = len;
316
317     // Huffman tables.
318
319     len = 2;
320
321     for (t = 0; t < 2 + 2 * (fh->Nf > 1); t++) {
322         JPEGRawHuffmanTable *huff = &priv->huffman_tables.table[t];
323         const uint8_t *lengths, *values;
324         int k;
325
326         switch (t) {
327         case 0:
328             lengths = avpriv_mjpeg_bits_dc_luminance + 1;
329             values  = avpriv_mjpeg_val_dc;
330             break;
331         case 1:
332             lengths = avpriv_mjpeg_bits_ac_luminance + 1;
333             values  = avpriv_mjpeg_val_ac_luminance;
334             break;
335         case 2:
336             lengths = avpriv_mjpeg_bits_dc_chrominance + 1;
337             values  = avpriv_mjpeg_val_dc;
338             break;
339         case 3:
340             lengths = avpriv_mjpeg_bits_ac_chrominance + 1;
341             values  = avpriv_mjpeg_val_ac_chrominance;
342             break;
343         }
344
345         huff->Tc = t % 2;
346         huff->Th = t / 2;
347
348         for (i = k = 0; i < 16; i++)
349             k += (huff->L[i] = lengths[i]);
350
351         for (i = 0; i < k; i++)
352             huff->V[i] = values[i];
353
354         len += 17 + k;
355     }
356
357     priv->huffman_tables.Lh = len;
358
359     // Scan header.
360
361     sh->Ns = fh->Nf;
362
363     for (i = 0; i < fh->Nf; i++) {
364         sh->Cs[i] = fh->C[i];
365         sh->Td[i] = i > 0;
366         sh->Ta[i] = i > 0;
367     }
368
369     sh->Ss = 0;
370     sh->Se = 63;
371     sh->Ah = 0;
372     sh->Al = 0;
373
374     sh->Ls = 6 + 2 * sh->Ns;
375
376
377     *vpic = (VAEncPictureParameterBufferJPEG) {
378         .reconstructed_picture = pic->recon_surface,
379         .coded_buf             = pic->output_buffer,
380
381         .picture_width  = fh->X,
382         .picture_height = fh->Y,
383
384         .pic_flags.bits = {
385             .profile      = 0,
386             .progressive  = 0,
387             .huffman      = 1,
388             .interleaved  = 0,
389             .differential = 0,
390         },
391
392         .sample_bit_depth = fh->P,
393         .num_scan         = 1,
394         .num_components   = fh->Nf,
395
396         // The driver modifies the provided quantisation tables according
397         // to this quality value; the middle value of 50 makes that the
398         // identity so that they are used unchanged.
399         .quality = 50,
400     };
401
402     for (i = 0; i < fh->Nf; i++) {
403         vpic->component_id[i]             = fh->C[i];
404         vpic->quantiser_table_selector[i] = fh->Tq[i];
405     }
406
407     pic->nb_slices = 1;
408
409     return 0;
410 }
411
412 static int vaapi_encode_mjpeg_init_slice_params(AVCodecContext *avctx,
413                                                 VAAPIEncodePicture *pic,
414                                                 VAAPIEncodeSlice *slice)
415 {
416     VAAPIEncodeMJPEGContext         *priv = avctx->priv_data;
417     JPEGRawScanHeader                 *sh = &priv->scan.header;
418     VAEncSliceParameterBufferJPEG *vslice = slice->codec_slice_params;
419     int i;
420
421     *vslice = (VAEncSliceParameterBufferJPEG) {
422         .restart_interval = 0,
423         .num_components   = sh->Ns,
424     };
425
426     for (i = 0; i < sh->Ns; i++) {
427         vslice->components[i].component_selector = sh->Cs[i];
428         vslice->components[i].dc_table_selector  = sh->Td[i];
429         vslice->components[i].ac_table_selector  = sh->Ta[i];
430     }
431
432     return 0;
433 }
434
435 static av_cold int vaapi_encode_mjpeg_configure(AVCodecContext *avctx)
436 {
437     VAAPIEncodeContext       *ctx = avctx->priv_data;
438     VAAPIEncodeMJPEGContext *priv = avctx->priv_data;
439     int err;
440
441     priv->quality = avctx->global_quality;
442     if (priv->quality < 1 || priv->quality > 100) {
443         av_log(avctx, AV_LOG_ERROR, "Invalid quality value %d "
444                "(must be 1-100).\n", priv->quality);
445         return AVERROR(EINVAL);
446     }
447
448     // Hack: the implementation calls the JPEG image header (which we
449     // will use in the same way as a slice header) generic "raw data".
450     // Therefore, if after the packed header capability check we have
451     // PACKED_HEADER_RAW_DATA available, rewrite it as
452     // PACKED_HEADER_SLICE so that the header-writing code can do the
453     // right thing.
454     if (ctx->va_packed_headers & VA_ENC_PACKED_HEADER_RAW_DATA) {
455         ctx->va_packed_headers &= ~VA_ENC_PACKED_HEADER_RAW_DATA;
456         ctx->va_packed_headers |=  VA_ENC_PACKED_HEADER_SLICE;
457     }
458
459     err = ff_cbs_init(&priv->cbc, AV_CODEC_ID_MJPEG, avctx);
460     if (err < 0)
461         return err;
462
463     return 0;
464 }
465
466 static const VAAPIEncodeProfile vaapi_encode_mjpeg_profiles[] = {
467     { FF_PROFILE_MJPEG_HUFFMAN_BASELINE_DCT,
468             8, 1, 0, 0, VAProfileJPEGBaseline },
469     { FF_PROFILE_MJPEG_HUFFMAN_BASELINE_DCT,
470             8, 3, 1, 1, VAProfileJPEGBaseline },
471     { FF_PROFILE_MJPEG_HUFFMAN_BASELINE_DCT,
472             8, 3, 1, 0, VAProfileJPEGBaseline },
473     { FF_PROFILE_MJPEG_HUFFMAN_BASELINE_DCT,
474             8, 3, 0, 0, VAProfileJPEGBaseline },
475     { FF_PROFILE_UNKNOWN }
476 };
477
478 static const VAAPIEncodeType vaapi_encode_type_mjpeg = {
479     .profiles              = vaapi_encode_mjpeg_profiles,
480
481     .flags                 = FLAG_CONSTANT_QUALITY_ONLY |
482                              FLAG_INTRA_ONLY,
483
484     .configure             = &vaapi_encode_mjpeg_configure,
485
486     .picture_params_size   = sizeof(VAEncPictureParameterBufferJPEG),
487     .init_picture_params   = &vaapi_encode_mjpeg_init_picture_params,
488
489     .slice_params_size     = sizeof(VAEncSliceParameterBufferJPEG),
490     .init_slice_params     = &vaapi_encode_mjpeg_init_slice_params,
491
492     .slice_header_type     = VAEncPackedHeaderRawData,
493     .write_slice_header    = &vaapi_encode_mjpeg_write_image_header,
494
495     .write_extra_buffer    = &vaapi_encode_mjpeg_write_extra_buffer,
496 };
497
498 static av_cold int vaapi_encode_mjpeg_init(AVCodecContext *avctx)
499 {
500     VAAPIEncodeContext *ctx = avctx->priv_data;
501
502     ctx->codec = &vaapi_encode_type_mjpeg;
503
504     // The JPEG image header - see note above.
505     ctx->desired_packed_headers =
506         VA_ENC_PACKED_HEADER_RAW_DATA;
507
508     ctx->surface_width  = FFALIGN(avctx->width,  8);
509     ctx->surface_height = FFALIGN(avctx->height, 8);
510
511     return ff_vaapi_encode_init(avctx);
512 }
513
514 static av_cold int vaapi_encode_mjpeg_close(AVCodecContext *avctx)
515 {
516     VAAPIEncodeMJPEGContext *priv = avctx->priv_data;
517
518     ff_cbs_close(&priv->cbc);
519
520     return ff_vaapi_encode_close(avctx);
521 }
522
523 #define OFFSET(x) offsetof(VAAPIEncodeMJPEGContext, x)
524 #define FLAGS (AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM)
525 static const AVOption vaapi_encode_mjpeg_options[] = {
526     VAAPI_ENCODE_COMMON_OPTIONS,
527
528     { "jfif", "Include JFIF header",
529       OFFSET(jfif), AV_OPT_TYPE_BOOL,
530       { .i64 = 0 }, 0, 1, FLAGS },
531     { "huffman", "Include huffman tables",
532       OFFSET(huffman), AV_OPT_TYPE_BOOL,
533       { .i64 = 1 }, 0, 1, FLAGS },
534
535     { NULL },
536 };
537
538 static const AVCodecDefault vaapi_encode_mjpeg_defaults[] = {
539     { "global_quality", "80" },
540     { "b",              "0"  },
541     { NULL },
542 };
543
544 static const AVClass vaapi_encode_mjpeg_class = {
545     .class_name = "mjpeg_vaapi",
546     .item_name  = av_default_item_name,
547     .option     = vaapi_encode_mjpeg_options,
548     .version    = LIBAVUTIL_VERSION_INT,
549 };
550
551 AVCodec ff_mjpeg_vaapi_encoder = {
552     .name           = "mjpeg_vaapi",
553     .long_name      = NULL_IF_CONFIG_SMALL("MJPEG (VAAPI)"),
554     .type           = AVMEDIA_TYPE_VIDEO,
555     .id             = AV_CODEC_ID_MJPEG,
556     .priv_data_size = sizeof(VAAPIEncodeMJPEGContext),
557     .init           = &vaapi_encode_mjpeg_init,
558     .send_frame     = &ff_vaapi_encode_send_frame,
559     .receive_packet = &ff_vaapi_encode_receive_packet,
560     .close          = &vaapi_encode_mjpeg_close,
561     .priv_class     = &vaapi_encode_mjpeg_class,
562     .capabilities   = AV_CODEC_CAP_HARDWARE |
563                       AV_CODEC_CAP_INTRA_ONLY,
564     .defaults       = vaapi_encode_mjpeg_defaults,
565     .pix_fmts = (const enum AVPixelFormat[]) {
566         AV_PIX_FMT_VAAPI,
567         AV_PIX_FMT_NONE,
568     },
569     .wrapper_name   = "vaapi",
570 };