git.sesse.net Git - ffmpeg/blob - libavcodec/extract_extradata_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 <stdint.h>
  20
  21 #include "libavutil/common.h"
  22 #include "libavutil/intreadwrite.h"
  23 #include "libavutil/log.h"
  24 #include "libavutil/opt.h"
  25
  26 #include "av1.h"
  27 #include "av1_parse.h"
  28 #include "bsf.h"
  29 #include "bsf_internal.h"
  30 #include "bytestream.h"
  31 #include "h2645_parse.h"
  32 #include "h264.h"
  33 #include "hevc.h"
  34 #include "vc1_common.h"
  35
  36 typedef struct ExtractExtradataContext {
  37     const AVClass *class;
  38
  39     int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
  40                    uint8_t **data, int *size);
  41
  42     /* AV1 specific fields */
  43     AV1Packet av1_pkt;
  44
  45     /* H264/HEVC specific fields */
  46     H2645Packet h2645_pkt;
  47
  48     /* AVOptions */
  49     int remove;
  50 } ExtractExtradataContext;
  51
  52 static int val_in_array(const int *arr, int len, int val)
  53 {
  54     int i;
  55     for (i = 0; i < len; i++)
  56         if (arr[i] == val)
  57             return 1;
  58     return 0;
  59 }
  60
  61 static int extract_extradata_av1(AVBSFContext *ctx, AVPacket *pkt,
  62                                  uint8_t **data, int *size)
  63 {
  64     static const int extradata_obu_types[] = {
  65         AV1_OBU_SEQUENCE_HEADER, AV1_OBU_METADATA,
  66     };
  67     ExtractExtradataContext *s = ctx->priv_data;
  68
  69     int extradata_size = 0, filtered_size = 0;
  70     int nb_extradata_obu_types = FF_ARRAY_ELEMS(extradata_obu_types);
  71     int i, has_seq = 0, ret = 0;
  72
  73     ret = ff_av1_packet_split(&s->av1_pkt, pkt->data, pkt->size, ctx);
  74     if (ret < 0)
  75         return ret;
  76
  77     for (i = 0; i < s->av1_pkt.nb_obus; i++) {
  78         AV1OBU *obu = &s->av1_pkt.obus[i];
  79         if (val_in_array(extradata_obu_types, nb_extradata_obu_types, obu->type)) {
  80             extradata_size += obu->raw_size;
  81             if (obu->type == AV1_OBU_SEQUENCE_HEADER)
  82                 has_seq = 1;
  83         } else if (s->remove) {
  84             filtered_size += obu->raw_size;
  85         }
  86     }
  87
  88     if (extradata_size && has_seq) {
  89         AVBufferRef *filtered_buf = NULL;
  90         PutByteContext pb_filtered_data, pb_extradata;
  91         uint8_t *extradata;
  92
  93         if (s->remove) {
  94             filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
  95             if (!filtered_buf) {
  96                 return AVERROR(ENOMEM);
  97             }
  98             memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
  99         }
 100
 101         extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
 102         if (!extradata) {
 103             av_buffer_unref(&filtered_buf);
 104             return AVERROR(ENOMEM);
 105         }
 106         memset(extradata + extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 107
 108         *data = extradata;
 109         *size = extradata_size;
 110
 111         bytestream2_init_writer(&pb_extradata, extradata, extradata_size);
 112         if (s->remove)
 113             bytestream2_init_writer(&pb_filtered_data, filtered_buf->data, filtered_size);
 114
 115         for (i = 0; i < s->av1_pkt.nb_obus; i++) {
 116             AV1OBU *obu = &s->av1_pkt.obus[i];
 117             if (val_in_array(extradata_obu_types, nb_extradata_obu_types,
 118                              obu->type)) {
 119                 bytestream2_put_bufferu(&pb_extradata, obu->raw_data, obu->raw_size);
 120             } else if (s->remove) {
 121                 bytestream2_put_bufferu(&pb_filtered_data, obu->raw_data, obu->raw_size);
 122             }
 123         }
 124
 125         if (s->remove) {
 126             av_buffer_unref(&pkt->buf);
 127             pkt->buf  = filtered_buf;
 128             pkt->data = filtered_buf->data;
 129             pkt->size = filtered_size;
 130         }
 131     }
 132
 133     return 0;
 134 }
 135
 136 static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt,
 137                                    uint8_t **data, int *size)
 138 {
 139     static const int extradata_nal_types_hevc[] = {
 140         HEVC_NAL_VPS, HEVC_NAL_SPS, HEVC_NAL_PPS,
 141     };
 142     static const int extradata_nal_types_h264[] = {
 143         H264_NAL_SPS, H264_NAL_PPS,
 144     };
 145
 146     ExtractExtradataContext *s = ctx->priv_data;
 147
 148     int extradata_size = 0, filtered_size = 0;
 149     const int *extradata_nal_types;
 150     int nb_extradata_nal_types;
 151     int i, has_sps = 0, has_vps = 0, ret = 0;
 152
 153     if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) {
 154         extradata_nal_types    = extradata_nal_types_hevc;
 155         nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_hevc);
 156     } else {
 157         extradata_nal_types    = extradata_nal_types_h264;
 158         nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_h264);
 159     }
 160
 161     ret = ff_h2645_packet_split(&s->h2645_pkt, pkt->data, pkt->size,
 162                                 ctx, 0, 0, ctx->par_in->codec_id, 1, 0);
 163     if (ret < 0)
 164         return ret;
 165
 166     for (i = 0; i < s->h2645_pkt.nb_nals; i++) {
 167         H2645NAL *nal = &s->h2645_pkt.nals[i];
 168         if (val_in_array(extradata_nal_types, nb_extradata_nal_types, nal->type)) {
 169             extradata_size += nal->raw_size + 3;
 170             if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) {
 171                 if (nal->type == HEVC_NAL_SPS) has_sps = 1;
 172                 if (nal->type == HEVC_NAL_VPS) has_vps = 1;
 173             } else {
 174                 if (nal->type == H264_NAL_SPS) has_sps = 1;
 175             }
 176         } else if (s->remove) {
 177             filtered_size += nal->raw_size + 3;
 178         }
 179     }
 180
 181     if (extradata_size &&
 182         ((ctx->par_in->codec_id == AV_CODEC_ID_HEVC && has_sps && has_vps) ||
 183          (ctx->par_in->codec_id == AV_CODEC_ID_H264 && has_sps))) {
 184         AVBufferRef *filtered_buf = NULL;
 185         PutByteContext pb_filtered_data, pb_extradata;
 186         uint8_t *extradata;
 187
 188         if (s->remove) {
 189             filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
 190             if (!filtered_buf) {
 191                 return AVERROR(ENOMEM);
 192             }
 193             memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 194         }
 195
 196         extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
 197         if (!extradata) {
 198             av_buffer_unref(&filtered_buf);
 199             return AVERROR(ENOMEM);
 200         }
 201         memset(extradata + extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 202
 203         *data = extradata;
 204         *size = extradata_size;
 205
 206         bytestream2_init_writer(&pb_extradata, extradata, extradata_size);
 207         if (s->remove)
 208             bytestream2_init_writer(&pb_filtered_data, filtered_buf->data, filtered_size);
 209
 210         for (i = 0; i < s->h2645_pkt.nb_nals; i++) {
 211             H2645NAL *nal = &s->h2645_pkt.nals[i];
 212             if (val_in_array(extradata_nal_types, nb_extradata_nal_types,
 213                              nal->type)) {
 214                 bytestream2_put_be24u(&pb_extradata, 1); //startcode
 215                 bytestream2_put_bufferu(&pb_extradata, nal->raw_data, nal->raw_size);
 216             } else if (s->remove) {
 217                 bytestream2_put_be24u(&pb_filtered_data, 1); // startcode
 218                 bytestream2_put_bufferu(&pb_filtered_data, nal->raw_data, nal->raw_size);
 219             }
 220         }
 221
 222         if (s->remove) {
 223             av_buffer_unref(&pkt->buf);
 224             pkt->buf  = filtered_buf;
 225             pkt->data = filtered_buf->data;
 226             pkt->size = filtered_size;
 227         }
 228     }
 229
 230     return 0;
 231 }
 232
 233 static int extract_extradata_vc1(AVBSFContext *ctx, AVPacket *pkt,
 234                                  uint8_t **data, int *size)
 235 {
 236     ExtractExtradataContext *s = ctx->priv_data;
 237     const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
 238     uint32_t state = UINT32_MAX;
 239     int has_extradata = 0, extradata_size = 0;
 240
 241     while (ptr < end) {
 242         ptr = avpriv_find_start_code(ptr, end, &state);
 243         if (state == VC1_CODE_SEQHDR || state == VC1_CODE_ENTRYPOINT) {
 244             has_extradata = 1;
 245         } else if (has_extradata && IS_MARKER(state)) {
 246             extradata_size = ptr - 4 - pkt->data;
 247             break;
 248         }
 249     }
 250
 251     if (extradata_size) {
 252         *data = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
 253         if (!*data)
 254             return AVERROR(ENOMEM);
 255
 256         memcpy(*data, pkt->data, extradata_size);
 257         memset(*data + extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 258         *size = extradata_size;
 259
 260         if (s->remove) {
 261             pkt->data += extradata_size;
 262             pkt->size -= extradata_size;
 263         }
 264     }
 265
 266     return 0;
 267 }
 268
 269 static int extract_extradata_mpeg12(AVBSFContext *ctx, AVPacket *pkt,
 270                                      uint8_t **data, int *size)
 271 {
 272     ExtractExtradataContext *s = ctx->priv_data;
 273     uint32_t state = UINT32_MAX;
 274     int i, found = 0;
 275
 276     for (i = 0; i < pkt->size; i++) {
 277         state = (state << 8) | pkt->data[i];
 278         if (state == 0x1B3)
 279             found = 1;
 280         else if (found && state != 0x1B5 && state < 0x200 && state >= 0x100) {
 281             if (i > 3) {
 282                 *size = i - 3;
 283                 *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
 284                 if (!*data)
 285                     return AVERROR(ENOMEM);
 286
 287                 memcpy(*data, pkt->data, *size);
 288                 memset(*data + *size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 289
 290                 if (s->remove) {
 291                     pkt->data += *size;
 292                     pkt->size -= *size;
 293                 }
 294             }
 295             break;
 296         }
 297     }
 298     return 0;
 299 }
 300
 301 static int extract_extradata_mpeg4(AVBSFContext *ctx, AVPacket *pkt,
 302                                    uint8_t **data, int *size)
 303 {
 304     ExtractExtradataContext *s = ctx->priv_data;
 305     const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
 306     uint32_t state = UINT32_MAX;
 307
 308     while (ptr < end) {
 309         ptr = avpriv_find_start_code(ptr, end, &state);
 310         if (state == 0x1B3 || state == 0x1B6) {
 311             if (ptr - pkt->data > 4) {
 312                 *size = ptr - 4 - pkt->data;
 313                 *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
 314                 if (!*data)
 315                     return AVERROR(ENOMEM);
 316
 317                 memcpy(*data, pkt->data, *size);
 318                 memset(*data + *size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 319
 320                 if (s->remove) {
 321                     pkt->data += *size;
 322                     pkt->size -= *size;
 323                 }
 324             }
 325             break;
 326         }
 327     }
 328     return 0;
 329 }
 330
 331 static const struct {
 332     enum AVCodecID id;
 333     int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
 334                    uint8_t **data, int *size);
 335 } extract_tab[] = {
 336     { AV_CODEC_ID_AV1,        extract_extradata_av1     },
 337     { AV_CODEC_ID_AVS2,       extract_extradata_mpeg4   },
 338     { AV_CODEC_ID_CAVS,       extract_extradata_mpeg4   },
 339     { AV_CODEC_ID_H264,       extract_extradata_h2645   },
 340     { AV_CODEC_ID_HEVC,       extract_extradata_h2645   },
 341     { AV_CODEC_ID_MPEG1VIDEO, extract_extradata_mpeg12  },
 342     { AV_CODEC_ID_MPEG2VIDEO, extract_extradata_mpeg12  },
 343     { AV_CODEC_ID_MPEG4,      extract_extradata_mpeg4   },
 344     { AV_CODEC_ID_VC1,        extract_extradata_vc1     },
 345 };
 346
 347 static int extract_extradata_init(AVBSFContext *ctx)
 348 {
 349     ExtractExtradataContext *s = ctx->priv_data;
 350     int i;
 351
 352     for (i = 0; i < FF_ARRAY_ELEMS(extract_tab); i++) {
 353         if (extract_tab[i].id == ctx->par_in->codec_id) {
 354             s->extract = extract_tab[i].extract;
 355             break;
 356         }
 357     }
 358     if (!s->extract)
 359         return AVERROR_BUG;
 360
 361     return 0;
 362 }
 363
 364 static int extract_extradata_filter(AVBSFContext *ctx, AVPacket *pkt)
 365 {
 366     ExtractExtradataContext *s = ctx->priv_data;
 367     uint8_t *extradata = NULL;
 368     int extradata_size;
 369     int ret = 0;
 370
 371     ret = ff_bsf_get_packet_ref(ctx, pkt);
 372     if (ret < 0)
 373         return ret;
 374
 375     ret = s->extract(ctx, pkt, &extradata, &extradata_size);
 376     if (ret < 0)
 377         goto fail;
 378
 379     if (extradata) {
 380         ret = av_packet_add_side_data(pkt, AV_PKT_DATA_NEW_EXTRADATA,
 381                                       extradata, extradata_size);
 382         if (ret < 0) {
 383             av_freep(&extradata);
 384             goto fail;
 385         }
 386     }
 387
 388     return 0;
 389
 390 fail:
 391     av_packet_unref(pkt);
 392     return ret;
 393 }
 394
 395 static void extract_extradata_close(AVBSFContext *ctx)
 396 {
 397     ExtractExtradataContext *s = ctx->priv_data;
 398     ff_av1_packet_uninit(&s->av1_pkt);
 399     ff_h2645_packet_uninit(&s->h2645_pkt);
 400 }
 401
 402 static const enum AVCodecID codec_ids[] = {
 403     AV_CODEC_ID_AV1,
 404     AV_CODEC_ID_AVS2,
 405     AV_CODEC_ID_CAVS,
 406     AV_CODEC_ID_H264,
 407     AV_CODEC_ID_HEVC,
 408     AV_CODEC_ID_MPEG1VIDEO,
 409     AV_CODEC_ID_MPEG2VIDEO,
 410     AV_CODEC_ID_MPEG4,
 411     AV_CODEC_ID_VC1,
 412     AV_CODEC_ID_NONE,
 413 };
 414
 415 #define OFFSET(x) offsetof(ExtractExtradataContext, x)
 416 #define FLAGS (AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_BSF_PARAM)
 417 static const AVOption options[] = {
 418     { "remove", "remove the extradata from the bitstream", OFFSET(remove), AV_OPT_TYPE_INT,
 419         { .i64 = 0 }, 0, 1, FLAGS },
 420     { NULL },
 421 };
 422
 423 static const AVClass extract_extradata_class = {
 424     .class_name = "extract_extradata",
 425     .item_name  = av_default_item_name,
 426     .option     = options,
 427     .version    = LIBAVUTIL_VERSION_INT,
 428 };
 429
 430 const AVBitStreamFilter ff_extract_extradata_bsf = {
 431     .name           = "extract_extradata",
 432     .codec_ids      = codec_ids,
 433     .priv_data_size = sizeof(ExtractExtradataContext),
 434     .priv_class     = &extract_extradata_class,
 435     .init           = extract_extradata_init,
 436     .filter         = extract_extradata_filter,
 437     .close          = extract_extradata_close,
 438 };