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