git.sesse.net Git - ffmpeg/blob - 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_reset(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 = ctx->rc_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     .default_quality       = 80,
 487
 488     .picture_params_size   = sizeof(VAEncPictureParameterBufferJPEG),
 489     .init_picture_params   = &vaapi_encode_mjpeg_init_picture_params,
 490
 491     .slice_params_size     = sizeof(VAEncSliceParameterBufferJPEG),
 492     .init_slice_params     = &vaapi_encode_mjpeg_init_slice_params,
 493
 494     .slice_header_type     = VAEncPackedHeaderRawData,
 495     .write_slice_header    = &vaapi_encode_mjpeg_write_image_header,
 496
 497     .write_extra_buffer    = &vaapi_encode_mjpeg_write_extra_buffer,
 498 };
 499
 500 static av_cold int vaapi_encode_mjpeg_init(AVCodecContext *avctx)
 501 {
 502     VAAPIEncodeContext *ctx = avctx->priv_data;
 503
 504     ctx->codec = &vaapi_encode_type_mjpeg;
 505
 506     // The JPEG image header - see note above.
 507     ctx->desired_packed_headers =
 508         VA_ENC_PACKED_HEADER_RAW_DATA;
 509
 510     ctx->surface_width  = FFALIGN(avctx->width,  8);
 511     ctx->surface_height = FFALIGN(avctx->height, 8);
 512
 513     return ff_vaapi_encode_init(avctx);
 514 }
 515
 516 static av_cold int vaapi_encode_mjpeg_close(AVCodecContext *avctx)
 517 {
 518     VAAPIEncodeMJPEGContext *priv = avctx->priv_data;
 519
 520     ff_cbs_fragment_free(priv->cbc, &priv->current_fragment);
 521     ff_cbs_close(&priv->cbc);
 522
 523     return ff_vaapi_encode_close(avctx);
 524 }
 525
 526 #define OFFSET(x) offsetof(VAAPIEncodeMJPEGContext, x)
 527 #define FLAGS (AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM)
 528 static const AVOption vaapi_encode_mjpeg_options[] = {
 529     VAAPI_ENCODE_COMMON_OPTIONS,
 530
 531     { "jfif", "Include JFIF header",
 532       OFFSET(jfif), AV_OPT_TYPE_BOOL,
 533       { .i64 = 0 }, 0, 1, FLAGS },
 534     { "huffman", "Include huffman tables",
 535       OFFSET(huffman), AV_OPT_TYPE_BOOL,
 536       { .i64 = 1 }, 0, 1, FLAGS },
 537
 538     { NULL },
 539 };
 540
 541 static const AVCodecDefault vaapi_encode_mjpeg_defaults[] = {
 542     { "b",              "0"  },
 543     { NULL },
 544 };
 545
 546 static const AVClass vaapi_encode_mjpeg_class = {
 547     .class_name = "mjpeg_vaapi",
 548     .item_name  = av_default_item_name,
 549     .option     = vaapi_encode_mjpeg_options,
 550     .version    = LIBAVUTIL_VERSION_INT,
 551 };
 552
 553 AVCodec ff_mjpeg_vaapi_encoder = {
 554     .name           = "mjpeg_vaapi",
 555     .long_name      = NULL_IF_CONFIG_SMALL("MJPEG (VAAPI)"),
 556     .type           = AVMEDIA_TYPE_VIDEO,
 557     .id             = AV_CODEC_ID_MJPEG,
 558     .priv_data_size = sizeof(VAAPIEncodeMJPEGContext),
 559     .init           = &vaapi_encode_mjpeg_init,
 560     .send_frame     = &ff_vaapi_encode_send_frame,
 561     .receive_packet = &ff_vaapi_encode_receive_packet,
 562     .close          = &vaapi_encode_mjpeg_close,
 563     .priv_class     = &vaapi_encode_mjpeg_class,
 564     .capabilities   = AV_CODEC_CAP_HARDWARE |
 565                       AV_CODEC_CAP_INTRA_ONLY,
 566     .defaults       = vaapi_encode_mjpeg_defaults,
 567     .pix_fmts = (const enum AVPixelFormat[]) {
 568         AV_PIX_FMT_VAAPI,
 569         AV_PIX_FMT_NONE,
 570     },
 571     .wrapper_name   = "vaapi",
 572 };