git.sesse.net Git - ffmpeg/blob - libavcodec/h264_metadata_bsf.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 "libavutil/avstring.h"
  20 #include "libavutil/common.h"
  21 #include "libavutil/opt.h"
  22
  23 #include "bsf.h"
  24 #include "cbs.h"
  25 #include "cbs_h264.h"
  26 #include "h264.h"
  27 #include "h264_sei.h"
  28
  29 enum {
  30     PASS,
  31     INSERT,
  32     REMOVE,
  33 };
  34
  35 typedef struct H264MetadataContext {
  36     const AVClass *class;
  37
  38     CodedBitstreamContext *cbc;
  39     CodedBitstreamFragment access_unit;
  40
  41     H264RawAUD aud_nal;
  42     H264RawSEI sei_nal;
  43
  44     int aud;
  45
  46     AVRational sample_aspect_ratio;
  47
  48     int video_format;
  49     int video_full_range_flag;
  50     int colour_primaries;
  51     int transfer_characteristics;
  52     int matrix_coefficients;
  53
  54     int chroma_sample_loc_type;
  55
  56     AVRational tick_rate;
  57     int fixed_frame_rate_flag;
  58
  59     int crop_left;
  60     int crop_right;
  61     int crop_top;
  62     int crop_bottom;
  63
  64     const char *sei_user_data;
  65     int sei_first_au;
  66 } H264MetadataContext;
  67
  68
  69 static int h264_metadata_update_sps(AVBSFContext *bsf,
  70                                     H264RawSPS *sps)
  71 {
  72     H264MetadataContext *ctx = bsf->priv_data;
  73     int need_vui = 0;
  74     int crop_unit_x, crop_unit_y;
  75
  76     if (ctx->sample_aspect_ratio.num && ctx->sample_aspect_ratio.den) {
  77         // Table E-1.
  78         static const AVRational sar_idc[] = {
  79             {   0,  0 }, // Unspecified (never written here).
  80             {   1,  1 }, {  12, 11 }, {  10, 11 }, {  16, 11 },
  81             {  40, 33 }, {  24, 11 }, {  20, 11 }, {  32, 11 },
  82             {  80, 33 }, {  18, 11 }, {  15, 11 }, {  64, 33 },
  83             { 160, 99 }, {   4,  3 }, {   3,  2 }, {   2,  1 },
  84         };
  85         int num, den, i;
  86
  87         av_reduce(&num, &den, ctx->sample_aspect_ratio.num,
  88                   ctx->sample_aspect_ratio.den, 65535);
  89
  90         for (i = 1; i < FF_ARRAY_ELEMS(sar_idc); i++) {
  91             if (num == sar_idc[i].num &&
  92                 den == sar_idc[i].den)
  93                 break;
  94         }
  95         if (i == FF_ARRAY_ELEMS(sar_idc)) {
  96             sps->vui.aspect_ratio_idc = 255;
  97             sps->vui.sar_width  = num;
  98             sps->vui.sar_height = den;
  99         } else {
 100             sps->vui.aspect_ratio_idc = i;
 101         }
 102         sps->vui.aspect_ratio_info_present_flag = 1;
 103         need_vui = 1;
 104     }
 105
 106 #define SET_OR_INFER(field, value, present_flag, infer) do { \
 107         if (value >= 0) { \
 108             field = value; \
 109             need_vui = 1; \
 110         } else if (!present_flag) \
 111             field = infer; \
 112     } while (0)
 113
 114     if (ctx->video_format             >= 0 ||
 115         ctx->video_full_range_flag    >= 0 ||
 116         ctx->colour_primaries         >= 0 ||
 117         ctx->transfer_characteristics >= 0 ||
 118         ctx->matrix_coefficients      >= 0) {
 119
 120         SET_OR_INFER(sps->vui.video_format, ctx->video_format,
 121                      sps->vui.video_signal_type_present_flag, 5);
 122
 123         SET_OR_INFER(sps->vui.video_full_range_flag,
 124                      ctx->video_full_range_flag,
 125                      sps->vui.video_signal_type_present_flag, 0);
 126
 127         if (ctx->colour_primaries         >= 0 ||
 128             ctx->transfer_characteristics >= 0 ||
 129             ctx->matrix_coefficients      >= 0) {
 130
 131             SET_OR_INFER(sps->vui.colour_primaries,
 132                          ctx->colour_primaries,
 133                          sps->vui.colour_description_present_flag, 2);
 134
 135             SET_OR_INFER(sps->vui.transfer_characteristics,
 136                          ctx->transfer_characteristics,
 137                          sps->vui.colour_description_present_flag, 2);
 138
 139             SET_OR_INFER(sps->vui.matrix_coefficients,
 140                          ctx->matrix_coefficients,
 141                          sps->vui.colour_description_present_flag, 2);
 142
 143             sps->vui.colour_description_present_flag = 1;
 144         }
 145         sps->vui.video_signal_type_present_flag = 1;
 146         need_vui = 1;
 147     }
 148
 149     if (ctx->chroma_sample_loc_type >= 0) {
 150         sps->vui.chroma_sample_loc_type_top_field =
 151             ctx->chroma_sample_loc_type;
 152         sps->vui.chroma_sample_loc_type_bottom_field =
 153             ctx->chroma_sample_loc_type;
 154         sps->vui.chroma_loc_info_present_flag = 1;
 155         need_vui = 1;
 156     }
 157
 158     if (ctx->tick_rate.num && ctx->tick_rate.den) {
 159         int num, den;
 160
 161         av_reduce(&num, &den, ctx->tick_rate.num, ctx->tick_rate.den,
 162                   UINT32_MAX > INT_MAX ? UINT32_MAX : INT_MAX);
 163
 164         sps->vui.time_scale        = num;
 165         sps->vui.num_units_in_tick = den;
 166
 167         sps->vui.timing_info_present_flag = 1;
 168         need_vui = 1;
 169     }
 170     SET_OR_INFER(sps->vui.fixed_frame_rate_flag,
 171                  ctx->fixed_frame_rate_flag,
 172                  sps->vui.timing_info_present_flag, 0);
 173
 174     if (sps->separate_colour_plane_flag || sps->chroma_format_idc == 0) {
 175         crop_unit_x = 1;
 176         crop_unit_y = 2 - sps->frame_mbs_only_flag;
 177     } else {
 178         crop_unit_x = 1 + (sps->chroma_format_idc < 3);
 179         crop_unit_y = (1 + (sps->chroma_format_idc < 2)) *
 180                        (2 - sps->frame_mbs_only_flag);
 181     }
 182 #define CROP(border, unit) do { \
 183         if (ctx->crop_ ## border >= 0) { \
 184             if (ctx->crop_ ## border % unit != 0) { \
 185                 av_log(bsf, AV_LOG_ERROR, "Invalid value for crop_%s: " \
 186                        "must be a multiple of %d.\n", #border, unit); \
 187                 return AVERROR(EINVAL); \
 188             } \
 189             sps->frame_crop_ ## border ## _offset = \
 190                   ctx->crop_ ## border / unit; \
 191             sps->frame_cropping_flag = 1; \
 192         } \
 193     } while (0)
 194     CROP(left,   crop_unit_x);
 195     CROP(right,  crop_unit_x);
 196     CROP(top,    crop_unit_y);
 197     CROP(bottom, crop_unit_y);
 198 #undef CROP
 199
 200     if (need_vui)
 201         sps->vui_parameters_present_flag = 1;
 202
 203     return 0;
 204 }
 205
 206 static int h264_metadata_filter(AVBSFContext *bsf, AVPacket *out)
 207 {
 208     H264MetadataContext *ctx = bsf->priv_data;
 209     AVPacket *in = NULL;
 210     CodedBitstreamFragment *au = &ctx->access_unit;
 211     int err, i, j, has_sps;
 212     char *sei_udu_string = NULL;
 213
 214     err = ff_bsf_get_packet(bsf, &in);
 215     if (err < 0)
 216         goto fail;
 217
 218     err = ff_cbs_read_packet(ctx->cbc, au, in);
 219     if (err < 0) {
 220         av_log(bsf, AV_LOG_ERROR, "Failed to read packet.\n");
 221         goto fail;
 222     }
 223
 224     if (au->nb_units == 0) {
 225         av_log(bsf, AV_LOG_ERROR, "No NAL units in packet.\n");
 226         err = AVERROR_INVALIDDATA;
 227         goto fail;
 228     }
 229
 230     // If an AUD is present, it must be the first NAL unit.
 231     if (au->units[0].type == H264_NAL_AUD) {
 232         if (ctx->aud == REMOVE)
 233             ff_cbs_delete_unit(ctx->cbc, au, 0);
 234     } else {
 235         if (ctx->aud == INSERT) {
 236             static const int primary_pic_type_table[] = {
 237                 0x084, // 2, 7
 238                 0x0a5, // 0, 2, 5, 7
 239                 0x0e7, // 0, 1, 2, 5, 6, 7
 240                 0x210, // 4, 9
 241                 0x318, // 3, 4, 8, 9
 242                 0x294, // 2, 4, 7, 9
 243                 0x3bd, // 0, 2, 3, 4, 5, 7, 8, 9
 244                 0x3ff, // 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
 245             };
 246             int primary_pic_type_mask = 0xff;
 247             H264RawAUD *aud = &ctx->aud_nal;
 248
 249             for (i = 0; i < au->nb_units; i++) {
 250                 if (au->units[i].type == H264_NAL_SLICE ||
 251                     au->units[i].type == H264_NAL_IDR_SLICE) {
 252                     H264RawSlice *slice = au->units[i].content;
 253                     for (j = 0; j < FF_ARRAY_ELEMS(primary_pic_type_table); j++) {
 254                          if (!(primary_pic_type_table[j] &
 255                                (1 << slice->header.slice_type)))
 256                              primary_pic_type_mask &= ~(1 << j);
 257                     }
 258                 }
 259             }
 260             for (j = 0; j < FF_ARRAY_ELEMS(primary_pic_type_table); j++)
 261                 if (primary_pic_type_mask & (1 << j))
 262                     break;
 263             if (j >= FF_ARRAY_ELEMS(primary_pic_type_table)) {
 264                 av_log(bsf, AV_LOG_ERROR, "No usable primary_pic_type: "
 265                        "invalid slice types?\n");
 266                 err = AVERROR_INVALIDDATA;
 267                 goto fail;
 268             }
 269
 270             aud->nal_unit_header.nal_unit_type = H264_NAL_AUD;
 271             aud->primary_pic_type = j;
 272
 273             err = ff_cbs_insert_unit_content(ctx->cbc, au,
 274                                              0, H264_NAL_AUD, aud);
 275             if (err < 0) {
 276                 av_log(bsf, AV_LOG_ERROR, "Failed to insert AUD.\n");
 277                 goto fail;
 278             }
 279         }
 280     }
 281
 282     has_sps = 0;
 283     for (i = 0; i < au->nb_units; i++) {
 284         if (au->units[i].type == H264_NAL_SPS) {
 285             err = h264_metadata_update_sps(bsf, au->units[i].content);
 286             if (err < 0)
 287                 goto fail;
 288             has_sps = 1;
 289         }
 290     }
 291
 292     // Insert the SEI in access units containing SPSs, and also
 293     // unconditionally in the first access unit we ever see.
 294     if (ctx->sei_user_data && (has_sps || !ctx->sei_first_au)) {
 295         H264RawSEI *sei;
 296         H264RawSEIPayload *payload;
 297         H264RawSEIUserDataUnregistered *udu;
 298         int sei_pos, sei_new;
 299
 300         ctx->sei_first_au = 1;
 301
 302         for (i = 0; i < au->nb_units; i++) {
 303             if (au->units[i].type == H264_NAL_SEI ||
 304                 au->units[i].type == H264_NAL_SLICE ||
 305                 au->units[i].type == H264_NAL_IDR_SLICE)
 306                 break;
 307         }
 308         sei_pos = i;
 309
 310         if (sei_pos < au->nb_units &&
 311             au->units[sei_pos].type == H264_NAL_SEI) {
 312             sei_new = 0;
 313             sei = au->units[sei_pos].content;
 314         } else {
 315             sei_new = 1;
 316             sei = &ctx->sei_nal;
 317             memset(sei, 0, sizeof(*sei));
 318
 319             sei->nal_unit_header.nal_unit_type = H264_NAL_SEI;
 320
 321             err = ff_cbs_insert_unit_content(ctx->cbc, au,
 322                                              sei_pos, H264_NAL_SEI, sei);
 323             if (err < 0) {
 324                 av_log(bsf, AV_LOG_ERROR, "Failed to insert SEI.\n");
 325                 goto fail;
 326             }
 327         }
 328
 329         payload = &sei->payload[sei->payload_count];
 330
 331         payload->payload_type = H264_SEI_TYPE_USER_DATA_UNREGISTERED;
 332         udu = &payload->payload.user_data_unregistered;
 333
 334         for (i = j = 0; j < 32 && ctx->sei_user_data[i]; i++) {
 335             int c, v;
 336             c = ctx->sei_user_data[i];
 337             if (c == '-') {
 338                 continue;
 339             } else if (av_isxdigit(c)) {
 340                 c = av_tolower(c);
 341                 v = (c <= '9' ? c - '0' : c - 'a' + 10);
 342             } else {
 343                 goto invalid_user_data;
 344             }
 345             if (i & 1)
 346                 udu->uuid_iso_iec_11578[j / 2] |= v;
 347             else
 348                 udu->uuid_iso_iec_11578[j / 2] = v << 4;
 349             ++j;
 350         }
 351         if (j == 32 && ctx->sei_user_data[i] == '+') {
 352             sei_udu_string = av_strdup(ctx->sei_user_data + i + 1);
 353             if (!sei_udu_string) {
 354                 err = AVERROR(ENOMEM);
 355                 goto sei_fail;
 356             }
 357
 358             udu->data = sei_udu_string;
 359             udu->data_length = strlen(sei_udu_string);
 360
 361             payload->payload_size = 16 + udu->data_length;
 362
 363             if (!sei_new) {
 364                 // This will be freed by the existing internal
 365                 // reference in fragment_uninit().
 366                 sei_udu_string = NULL;
 367             }
 368
 369         } else {
 370         invalid_user_data:
 371             av_log(bsf, AV_LOG_ERROR, "Invalid user data: "
 372                    "must be \"UUID+string\".\n");
 373             err = AVERROR(EINVAL);
 374         sei_fail:
 375             memset(payload, 0, sizeof(*payload));
 376             goto fail;
 377         }
 378
 379         ++sei->payload_count;
 380     }
 381
 382     err = ff_cbs_write_packet(ctx->cbc, out, au);
 383     if (err < 0) {
 384         av_log(bsf, AV_LOG_ERROR, "Failed to write packet.\n");
 385         goto fail;
 386     }
 387
 388     err = av_packet_copy_props(out, in);
 389     if (err < 0)
 390         goto fail;
 391
 392     err = 0;
 393 fail:
 394     ff_cbs_fragment_uninit(ctx->cbc, au);
 395     av_freep(&sei_udu_string);
 396
 397     av_packet_free(&in);
 398
 399     return err;
 400 }
 401
 402 static int h264_metadata_init(AVBSFContext *bsf)
 403 {
 404     H264MetadataContext *ctx = bsf->priv_data;
 405     CodedBitstreamFragment *au = &ctx->access_unit;
 406     int err, i;
 407
 408     err = ff_cbs_init(&ctx->cbc, AV_CODEC_ID_H264, bsf);
 409     if (err < 0)
 410         return err;
 411
 412     if (bsf->par_in->extradata) {
 413         err = ff_cbs_read_extradata(ctx->cbc, au, bsf->par_in);
 414         if (err < 0) {
 415             av_log(bsf, AV_LOG_ERROR, "Failed to read extradata.\n");
 416             goto fail;
 417         }
 418
 419         for (i = 0; i < au->nb_units; i++) {
 420             if (au->units[i].type == H264_NAL_SPS) {
 421                 err = h264_metadata_update_sps(bsf, au->units[i].content);
 422                 if (err < 0)
 423                     goto fail;
 424             }
 425         }
 426
 427         err = ff_cbs_write_extradata(ctx->cbc, bsf->par_out, au);
 428         if (err < 0) {
 429             av_log(bsf, AV_LOG_ERROR, "Failed to write extradata.\n");
 430             goto fail;
 431         }
 432     }
 433
 434     err = 0;
 435 fail:
 436     ff_cbs_fragment_uninit(ctx->cbc, au);
 437     return err;
 438 }
 439
 440 static void h264_metadata_close(AVBSFContext *bsf)
 441 {
 442     H264MetadataContext *ctx = bsf->priv_data;
 443     ff_cbs_close(&ctx->cbc);
 444 }
 445
 446 #define OFFSET(x) offsetof(H264MetadataContext, x)
 447 static const AVOption h264_metadata_options[] = {
 448     { "aud", "Access Unit Delimiter NAL units",
 449         OFFSET(aud), AV_OPT_TYPE_INT,
 450         { .i64 = PASS }, PASS, REMOVE, 0, "aud" },
 451     { "pass",   NULL, 0, AV_OPT_TYPE_CONST, { .i64 = PASS   }, .unit = "aud" },
 452     { "insert", NULL, 0, AV_OPT_TYPE_CONST, { .i64 = INSERT }, .unit = "aud" },
 453     { "remove", NULL, 0, AV_OPT_TYPE_CONST, { .i64 = REMOVE }, .unit = "aud" },
 454
 455     { "sample_aspect_ratio", "Set sample aspect ratio (table E-1)",
 456         OFFSET(sample_aspect_ratio), AV_OPT_TYPE_RATIONAL,
 457         { .dbl = 0.0 }, 0, 65535 },
 458
 459     { "video_format", "Set video format (table E-2)",
 460         OFFSET(video_format), AV_OPT_TYPE_INT,
 461         { .i64 = -1 }, -1, 7 },
 462     { "video_full_range_flag", "Set video full range flag",
 463         OFFSET(video_full_range_flag), AV_OPT_TYPE_INT,
 464         { .i64 = -1 }, -1, 1 },
 465     { "colour_primaries", "Set colour primaries (table E-3)",
 466         OFFSET(colour_primaries), AV_OPT_TYPE_INT,
 467         { .i64 = -1 }, -1, 255 },
 468     { "transfer_characteristics", "Set transfer characteristics (table E-4)",
 469         OFFSET(transfer_characteristics), AV_OPT_TYPE_INT,
 470         { .i64 = -1 }, -1, 255 },
 471     { "matrix_coefficients", "Set matrix coefficients (table E-5)",
 472         OFFSET(matrix_coefficients), AV_OPT_TYPE_INT,
 473         { .i64 = -1 }, -1, 255 },
 474
 475     { "chroma_sample_loc_type", "Set chroma sample location type (figure E-1)",
 476         OFFSET(chroma_sample_loc_type), AV_OPT_TYPE_INT,
 477         { .i64 = -1 }, -1, 6 },
 478
 479     { "tick_rate", "Set VUI tick rate (num_units_in_tick / time_scale)",
 480         OFFSET(tick_rate), AV_OPT_TYPE_RATIONAL,
 481         { .dbl = 0.0 }, 0, UINT_MAX },
 482     { "fixed_frame_rate_flag", "Set VUI fixed frame rate flag",
 483         OFFSET(fixed_frame_rate_flag), AV_OPT_TYPE_INT,
 484         { .i64 = -1 }, -1, 1 },
 485
 486     { "crop_left", "Set left border crop offset",
 487         OFFSET(crop_left), AV_OPT_TYPE_INT,
 488         { .i64 = -1 }, -1, H264_MAX_WIDTH },
 489     { "crop_right", "Set right border crop offset",
 490         OFFSET(crop_right), AV_OPT_TYPE_INT,
 491         { .i64 = -1 }, -1, H264_MAX_WIDTH },
 492     { "crop_top", "Set top border crop offset",
 493         OFFSET(crop_top), AV_OPT_TYPE_INT,
 494         { .i64 = -1 }, -1, H264_MAX_HEIGHT },
 495     { "crop_bottom", "Set bottom border crop offset",
 496         OFFSET(crop_bottom), AV_OPT_TYPE_INT,
 497         { .i64 = -1 }, -1, H264_MAX_HEIGHT },
 498
 499     { "sei_user_data", "Insert SEI user data (UUID+string)",
 500         OFFSET(sei_user_data), AV_OPT_TYPE_STRING, { .str = NULL } },
 501
 502     { NULL }
 503 };
 504
 505 static const AVClass h264_metadata_class = {
 506     .class_name = "h264_metadata_bsf",
 507     .item_name  = av_default_item_name,
 508     .option     = h264_metadata_options,
 509     .version    = LIBAVCODEC_VERSION_MAJOR,
 510 };
 511
 512 static const enum AVCodecID h264_metadata_codec_ids[] = {
 513     AV_CODEC_ID_H264, AV_CODEC_ID_NONE,
 514 };
 515
 516 const AVBitStreamFilter ff_h264_metadata_bsf = {
 517     .name           = "h264_metadata",
 518     .priv_data_size = sizeof(H264MetadataContext),
 519     .priv_class     = &h264_metadata_class,
 520     .init           = &h264_metadata_init,
 521     .close          = &h264_metadata_close,
 522     .filter         = &h264_metadata_filter,
 523     .codec_ids      = h264_metadata_codec_ids,
 524 };