git.sesse.net Git - ffmpeg/blob - libavcodec/mpeg2_metadata_bsf.c

   1 /*
   2  * This file is part of Libav.
   3  *
   4  * Libav 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  * Libav 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 Libav; 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_mpeg2.h"
  26
  27 typedef struct MPEG2MetadataContext {
  28     const AVClass *class;
  29
  30     CodedBitstreamContext cbc;
  31     CodedBitstreamFragment fragment;
  32
  33     MPEG2RawExtensionData sequence_display_extension;
  34
  35     AVRational display_aspect_ratio;
  36
  37     AVRational frame_rate;
  38
  39     int video_format;
  40     int colour_primaries;
  41     int transfer_characteristics;
  42     int matrix_coefficients;
  43
  44     int mpeg1_warned;
  45 } MPEG2MetadataContext;
  46
  47
  48 static int mpeg2_metadata_update_fragment(AVBSFContext *bsf,
  49                                           CodedBitstreamFragment *frag)
  50 {
  51     MPEG2MetadataContext             *ctx = bsf->priv_data;
  52     MPEG2RawSequenceHeader            *sh = NULL;
  53     MPEG2RawSequenceExtension         *se = NULL;
  54     MPEG2RawSequenceDisplayExtension *sde = NULL;
  55     int i, se_pos, add_sde = 0;
  56
  57     for (i = 0; i < frag->nb_units; i++) {
  58         if (frag->units[i].type == MPEG2_START_SEQUENCE_HEADER) {
  59             sh = frag->units[i].content;
  60         } else if (frag->units[i].type == MPEG2_START_EXTENSION) {
  61             MPEG2RawExtensionData *ext = frag->units[i].content;
  62             if (ext->extension_start_code_identifier ==
  63                 MPEG2_EXTENSION_SEQUENCE) {
  64                 se = &ext->data.sequence;
  65                 se_pos = i;
  66             } else if (ext->extension_start_code_identifier ==
  67                 MPEG2_EXTENSION_SEQUENCE_DISPLAY) {
  68                 sde = &ext->data.sequence_display;
  69             }
  70         }
  71     }
  72
  73     if (!sh || !se) {
  74         // No sequence header and sequence extension: not an MPEG-2 video
  75         // sequence.
  76         if (sh && !ctx->mpeg1_warned) {
  77             av_log(bsf, AV_LOG_WARNING, "Stream contains a sequence "
  78                    "header but not a sequence extension: maybe it's "
  79                    "actually MPEG-1?\n");
  80             ctx->mpeg1_warned = 1;
  81         }
  82         return 0;
  83     }
  84
  85     if (ctx->display_aspect_ratio.num && ctx->display_aspect_ratio.den) {
  86         int num, den;
  87
  88         av_reduce(&num, &den, ctx->display_aspect_ratio.num,
  89                   ctx->display_aspect_ratio.den, 65535);
  90
  91         if (num == 4 && den == 3)
  92             sh->aspect_ratio_information = 2;
  93         else if (num == 16 && den == 9)
  94             sh->aspect_ratio_information = 3;
  95         else if (num == 221 && den == 100)
  96             sh->aspect_ratio_information = 4;
  97         else
  98             sh->aspect_ratio_information = 1;
  99     }
 100
 101     if (ctx->frame_rate.num && ctx->frame_rate.den) {
 102         // Table 6-4.
 103         static AVRational frame_rate_table[] = {
 104             { 0, 0 },
 105             { 24000, 1001 },
 106             { 24,    1    },
 107             { 25,    1    },
 108             { 30000, 1001 },
 109             { 30,    1    },
 110             { 50,    1    },
 111             { 60000, 1001 },
 112             { 60,    1    },
 113         };
 114         int code, ext_n, ext_d;
 115         AVRational best_error = { INT_MAX, 1 };
 116
 117         for (i = 1; i < FF_ARRAY_ELEMS(frame_rate_table); i++) {
 118             if (av_cmp_q(ctx->frame_rate, frame_rate_table[i]) == 0) {
 119                 code  = i;
 120                 ext_n = 1;
 121                 ext_d = 1;
 122                 goto found_frame_rate;
 123             }
 124         }
 125
 126         for (i = 1; i < FF_ARRAY_ELEMS(frame_rate_table); i++) {
 127             AVRational fr, error;
 128             int n, d, cmp;
 129             for (n = 1; n <= 4; n++) {
 130                 for (d = 1; d <= 32; d++) {
 131                     fr = av_mul_q(frame_rate_table[i],
 132                                   (AVRational) { n, d });
 133                     cmp = av_cmp_q(fr, ctx->frame_rate);
 134                     if (cmp == 0) {
 135                         code  = i;
 136                         ext_n = n;
 137                         ext_d = d;
 138                         goto found_frame_rate;
 139                     }
 140                     if (cmp < 0)
 141                         error = av_div_q(ctx->frame_rate, fr);
 142                     else
 143                         error = av_div_q(fr, ctx->frame_rate);
 144                     cmp = av_cmp_q(error, best_error);
 145                     if (cmp < 0 || (cmp == 0 && n == 1 && d == 1)) {
 146                         code  = i;
 147                         ext_n = n;
 148                         ext_d = d;
 149                         best_error = error;
 150                     }
 151                 }
 152             }
 153         }
 154
 155     found_frame_rate:
 156         sh->frame_rate_code        = code;
 157         se->frame_rate_extension_n = ext_n - 1;
 158         se->frame_rate_extension_d = ext_d - 1;
 159     }
 160
 161     if (ctx->video_format             >= 0 ||
 162         ctx->colour_primaries         >= 0 ||
 163         ctx->transfer_characteristics >= 0 ||
 164         ctx->matrix_coefficients      >= 0) {
 165         if (!sde) {
 166             add_sde = 1;
 167             ctx->sequence_display_extension.extension_start_code =
 168                 MPEG2_START_EXTENSION;
 169             ctx->sequence_display_extension.extension_start_code_identifier =
 170                 MPEG2_EXTENSION_SEQUENCE_DISPLAY;
 171             sde = &ctx->sequence_display_extension.data.sequence_display;
 172
 173             *sde = (MPEG2RawSequenceDisplayExtension) {
 174                 .video_format = 5,
 175
 176                 .colour_description       = 0,
 177                 .colour_primaries         = 2,
 178                 .transfer_characteristics = 2,
 179                 .matrix_coefficients      = 2,
 180
 181                 .display_horizontal_size =
 182                     se->horizontal_size_extension << 12 | sh->horizontal_size_value,
 183                 .display_vertical_size =
 184                     se->vertical_size_extension << 12 | sh->vertical_size_value,
 185             };
 186         }
 187
 188         if (ctx->video_format >= 0)
 189             sde->video_format = ctx->video_format;
 190
 191         if (ctx->colour_primaries         >= 0 ||
 192             ctx->transfer_characteristics >= 0 ||
 193             ctx->matrix_coefficients      >= 0) {
 194             sde->colour_description = 1;
 195
 196             if (ctx->colour_primaries >= 0)
 197                 sde->colour_primaries = ctx->colour_primaries;
 198             else if (add_sde)
 199                 sde->colour_primaries = 2;
 200
 201             if (ctx->transfer_characteristics >= 0)
 202                 sde->transfer_characteristics = ctx->transfer_characteristics;
 203             else if (add_sde)
 204                 sde->transfer_characteristics = 2;
 205
 206             if (ctx->matrix_coefficients >= 0)
 207                 sde->matrix_coefficients = ctx->matrix_coefficients;
 208             else if (add_sde)
 209                 sde->matrix_coefficients = 2;
 210         }
 211     }
 212
 213     if (add_sde) {
 214         int err;
 215
 216         err = ff_cbs_insert_unit_content(&ctx->cbc, frag, se_pos + 1,
 217                                          MPEG2_START_EXTENSION,
 218                                          &ctx->sequence_display_extension);
 219         if (err < 0) {
 220             av_log(bsf, AV_LOG_ERROR, "Failed to insert new sequence "
 221                    "display extension.\n");
 222             return err;
 223         }
 224     }
 225
 226     return 0;
 227 }
 228
 229 static int mpeg2_metadata_filter(AVBSFContext *bsf, AVPacket *out)
 230 {
 231     MPEG2MetadataContext *ctx = bsf->priv_data;
 232     AVPacket *in = NULL;
 233     CodedBitstreamFragment *frag = &ctx->fragment;
 234     int err;
 235
 236     err = ff_bsf_get_packet(bsf, &in);
 237     if (err < 0)
 238         goto fail;
 239
 240     err = ff_cbs_read_packet(&ctx->cbc, frag, in);
 241     if (err < 0) {
 242         av_log(bsf, AV_LOG_ERROR, "Failed to read packet.\n");
 243         goto fail;
 244     }
 245
 246     err = mpeg2_metadata_update_fragment(bsf, frag);
 247     if (err < 0) {
 248         av_log(bsf, AV_LOG_ERROR, "Failed to update frame fragment.\n");
 249         goto fail;
 250     }
 251
 252     err = ff_cbs_write_packet(&ctx->cbc, out, frag);
 253     if (err < 0) {
 254         av_log(bsf, AV_LOG_ERROR, "Failed to write packet.\n");
 255         goto fail;
 256     }
 257
 258     err = av_packet_copy_props(out, in);
 259     if (err < 0) {
 260         av_packet_unref(out);
 261         goto fail;
 262     }
 263
 264     err = 0;
 265 fail:
 266     ff_cbs_fragment_uninit(&ctx->cbc, frag);
 267
 268     av_packet_free(&in);
 269
 270     return err;
 271 }
 272
 273 static int mpeg2_metadata_init(AVBSFContext *bsf)
 274 {
 275     MPEG2MetadataContext *ctx = bsf->priv_data;
 276     CodedBitstreamFragment *frag = &ctx->fragment;
 277     int err;
 278
 279     err = ff_cbs_init(&ctx->cbc, AV_CODEC_ID_MPEG2VIDEO, bsf);
 280     if (err < 0)
 281         return err;
 282
 283     if (bsf->par_in->extradata) {
 284         err = ff_cbs_read_extradata(&ctx->cbc, frag, bsf->par_in);
 285         if (err < 0) {
 286             av_log(bsf, AV_LOG_ERROR, "Failed to read extradata.\n");
 287             goto fail;
 288         }
 289
 290         err = mpeg2_metadata_update_fragment(bsf, frag);
 291         if (err < 0) {
 292             av_log(bsf, AV_LOG_ERROR, "Failed to update metadata fragment.\n");
 293             goto fail;
 294         }
 295
 296         err = ff_cbs_write_extradata(&ctx->cbc, bsf->par_out, frag);
 297         if (err < 0) {
 298             av_log(bsf, AV_LOG_ERROR, "Failed to write extradata.\n");
 299             goto fail;
 300         }
 301     }
 302
 303     err = 0;
 304 fail:
 305     ff_cbs_fragment_uninit(&ctx->cbc, frag);
 306     return err;
 307 }
 308
 309 static void mpeg2_metadata_close(AVBSFContext *bsf)
 310 {
 311     MPEG2MetadataContext *ctx = bsf->priv_data;
 312     ff_cbs_close(&ctx->cbc);
 313 }
 314
 315 #define OFFSET(x) offsetof(MPEG2MetadataContext, x)
 316 static const AVOption mpeg2_metadata_options[] = {
 317     { "display_aspect_ratio", "Set display aspect ratio (table 6-3)",
 318         OFFSET(display_aspect_ratio), AV_OPT_TYPE_RATIONAL,
 319         { .i64 = 0 }, 0, 65535 },
 320
 321     { "frame_rate", "Set frame rate",
 322         OFFSET(frame_rate), AV_OPT_TYPE_RATIONAL,
 323         { .i64 = 0 }, 0, UINT_MAX },
 324
 325     { "video_format", "Set video format (table 6-6)",
 326         OFFSET(video_format), AV_OPT_TYPE_INT,
 327         { .i64 = -1 }, -1, 7 },
 328     { "colour_primaries", "Set colour primaries (table 6-7)",
 329         OFFSET(colour_primaries), AV_OPT_TYPE_INT,
 330         { .i64 = -1 }, -1, 255 },
 331     { "transfer_characteristics", "Set transfer characteristics (table 6-8)",
 332         OFFSET(transfer_characteristics), AV_OPT_TYPE_INT,
 333         { .i64 = -1 }, -1, 255 },
 334     { "matrix_coefficients", "Set matrix coefficients (table 6-9)",
 335         OFFSET(matrix_coefficients), AV_OPT_TYPE_INT,
 336         { .i64 = -1 }, -1, 255 },
 337
 338     { NULL }
 339 };
 340
 341 static const AVClass mpeg2_metadata_class = {
 342     .class_name = "mpeg2_metadata_bsf",
 343     .item_name  = av_default_item_name,
 344     .option     = mpeg2_metadata_options,
 345     .version    = LIBAVCODEC_VERSION_MAJOR,
 346 };
 347
 348 static const enum AVCodecID mpeg2_metadata_codec_ids[] = {
 349     AV_CODEC_ID_MPEG2VIDEO, AV_CODEC_ID_NONE,
 350 };
 351
 352 const AVBitStreamFilter ff_mpeg2_metadata_bsf = {
 353     .name           = "mpeg2_metadata",
 354     .priv_data_size = sizeof(MPEG2MetadataContext),
 355     .priv_class     = &mpeg2_metadata_class,
 356     .init           = &mpeg2_metadata_init,
 357     .close          = &mpeg2_metadata_close,
 358     .filter         = &mpeg2_metadata_filter,
 359     .codec_ids      = mpeg2_metadata_codec_ids,
 360 };