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
 107         *data = extradata;
 108         *size = extradata_size;
 109
 110         bytestream2_init_writer(&pb_extradata, extradata, extradata_size);
 111         if (s->remove)
 112             bytestream2_init_writer(&pb_filtered_data, filtered_buf->data, filtered_size);
 113
 114         for (i = 0; i < s->av1_pkt.nb_obus; i++) {
 115             AV1OBU *obu = &s->av1_pkt.obus[i];
 116             if (val_in_array(extradata_obu_types, nb_extradata_obu_types,
 117                              obu->type)) {
 118                 bytestream2_put_bufferu(&pb_extradata, obu->raw_data, obu->raw_size);
 119             } else if (s->remove) {
 120                 bytestream2_put_bufferu(&pb_filtered_data, obu->raw_data, obu->raw_size);
 121             }
 122         }
 123
 124         if (s->remove) {
 125             av_buffer_unref(&pkt->buf);
 126             pkt->buf  = filtered_buf;
 127             pkt->data = filtered_buf->data;
 128             pkt->size = filtered_size;
 129         }
 130     }
 131
 132     return 0;
 133 }
 134
 135 static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt,
 136                                    uint8_t **data, int *size)
 137 {
 138     static const int extradata_nal_types_hevc[] = {
 139         HEVC_NAL_VPS, HEVC_NAL_SPS, HEVC_NAL_PPS,
 140     };
 141     static const int extradata_nal_types_h264[] = {
 142         H264_NAL_SPS, H264_NAL_PPS,
 143     };
 144
 145     ExtractExtradataContext *s = ctx->priv_data;
 146
 147     int extradata_size = 0, filtered_size = 0;
 148     const int *extradata_nal_types;
 149     int nb_extradata_nal_types;
 150     int i, has_sps = 0, has_vps = 0, ret = 0;
 151
 152     if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) {
 153         extradata_nal_types    = extradata_nal_types_hevc;
 154         nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_hevc);
 155     } else {
 156         extradata_nal_types    = extradata_nal_types_h264;
 157         nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_h264);
 158     }
 159
 160     ret = ff_h2645_packet_split(&s->h2645_pkt, pkt->data, pkt->size,
 161                                 ctx, 0, 0, ctx->par_in->codec_id, 1, 0);
 162     if (ret < 0)
 163         return ret;
 164
 165     for (i = 0; i < s->h2645_pkt.nb_nals; i++) {
 166         H2645NAL *nal = &s->h2645_pkt.nals[i];
 167         if (val_in_array(extradata_nal_types, nb_extradata_nal_types, nal->type)) {
 168             extradata_size += nal->raw_size + 3;
 169             if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) {
 170                 if (nal->type == HEVC_NAL_SPS) has_sps = 1;
 171                 if (nal->type == HEVC_NAL_VPS) has_vps = 1;
 172             } else {
 173                 if (nal->type == H264_NAL_SPS) has_sps = 1;
 174             }
 175         } else if (s->remove) {
 176             filtered_size += nal->raw_size + 3;
 177         }
 178     }
 179
 180     if (extradata_size &&
 181         ((ctx->par_in->codec_id == AV_CODEC_ID_HEVC && has_sps && has_vps) ||
 182          (ctx->par_in->codec_id == AV_CODEC_ID_H264 && has_sps))) {
 183         AVBufferRef *filtered_buf = NULL;
 184         PutByteContext pb_filtered_data, pb_extradata;
 185         uint8_t *extradata;
 186
 187         if (s->remove) {
 188             filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
 189             if (!filtered_buf) {
 190                 return AVERROR(ENOMEM);
 191             }
 192             memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 193         }
 194
 195         extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
 196         if (!extradata) {
 197             av_buffer_unref(&filtered_buf);
 198             return AVERROR(ENOMEM);
 199         }
 200
 201         *data = extradata;
 202         *size = extradata_size;
 203
 204         bytestream2_init_writer(&pb_extradata, extradata, extradata_size);
 205         if (s->remove)
 206             bytestream2_init_writer(&pb_filtered_data, filtered_buf->data, filtered_size);
 207
 208         for (i = 0; i < s->h2645_pkt.nb_nals; i++) {
 209             H2645NAL *nal = &s->h2645_pkt.nals[i];
 210             if (val_in_array(extradata_nal_types, nb_extradata_nal_types,
 211                              nal->type)) {
 212                 bytestream2_put_be24u(&pb_extradata, 1); //startcode
 213                 bytestream2_put_bufferu(&pb_extradata, nal->raw_data, nal->raw_size);
 214             } else if (s->remove) {
 215                 bytestream2_put_be24u(&pb_filtered_data, 1); // startcode
 216                 bytestream2_put_bufferu(&pb_filtered_data, nal->raw_data, nal->raw_size);
 217             }
 218         }
 219
 220         if (s->remove) {
 221             av_buffer_unref(&pkt->buf);
 222             pkt->buf  = filtered_buf;
 223             pkt->data = filtered_buf->data;
 224             pkt->size = filtered_size;
 225         }
 226     }
 227
 228     return 0;
 229 }
 230
 231 static int extract_extradata_vc1(AVBSFContext *ctx, AVPacket *pkt,
 232                                  uint8_t **data, int *size)
 233 {
 234     ExtractExtradataContext *s = ctx->priv_data;
 235     const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
 236     uint32_t state = UINT32_MAX;
 237     int has_extradata = 0, extradata_size = 0;
 238
 239     while (ptr < end) {
 240         ptr = avpriv_find_start_code(ptr, end, &state);
 241         if (state == VC1_CODE_SEQHDR || state == VC1_CODE_ENTRYPOINT) {
 242             has_extradata = 1;
 243         } else if (has_extradata && IS_MARKER(state)) {
 244             extradata_size = ptr - 4 - pkt->data;
 245             break;
 246         }
 247     }
 248
 249     if (extradata_size) {
 250         *data = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
 251         if (!*data)
 252             return AVERROR(ENOMEM);
 253
 254         memcpy(*data, pkt->data, extradata_size);
 255         *size = extradata_size;
 256
 257         if (s->remove) {
 258             pkt->data += extradata_size;
 259             pkt->size -= extradata_size;
 260         }
 261     }
 262
 263     return 0;
 264 }
 265
 266 static int extract_extradata_mpeg12(AVBSFContext *ctx, AVPacket *pkt,
 267                                      uint8_t **data, int *size)
 268 {
 269     ExtractExtradataContext *s = ctx->priv_data;
 270     uint32_t state = UINT32_MAX;
 271     int i, found = 0;
 272
 273     for (i = 0; i < pkt->size; i++) {
 274         state = (state << 8) | pkt->data[i];
 275         if (state == 0x1B3)
 276             found = 1;
 277         else if (found && state != 0x1B5 && state < 0x200 && state >= 0x100) {
 278             *size = i - 3;
 279             *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
 280             if (!*data)
 281                 return AVERROR(ENOMEM);
 282
 283             memcpy(*data, pkt->data, *size);
 284
 285             if (s->remove) {
 286                 pkt->data += *size;
 287                 pkt->size -= *size;
 288             }
 289             break;
 290         }
 291     }
 292     return 0;
 293 }
 294
 295 static int extract_extradata_mpeg4(AVBSFContext *ctx, AVPacket *pkt,
 296                                    uint8_t **data, int *size)
 297 {
 298     ExtractExtradataContext *s = ctx->priv_data;
 299     const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
 300     uint32_t state = UINT32_MAX;
 301
 302     while (ptr < end) {
 303         ptr = avpriv_find_start_code(ptr, end, &state);
 304         if (state == 0x1B3 || state == 0x1B6) {
 305             if (ptr - pkt->data > 4) {
 306                 *size = ptr - 4 - pkt->data;
 307                 *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
 308                 if (!*data)
 309                     return AVERROR(ENOMEM);
 310
 311                 memcpy(*data, pkt->data, *size);
 312
 313                 if (s->remove) {
 314                     pkt->data += *size;
 315                     pkt->size -= *size;
 316                 }
 317             }
 318             break;
 319         }
 320     }
 321     return 0;
 322 }
 323
 324 static const struct {
 325     enum AVCodecID id;
 326     int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
 327                    uint8_t **data, int *size);
 328 } extract_tab[] = {
 329     { AV_CODEC_ID_AV1,        extract_extradata_av1     },
 330     { AV_CODEC_ID_AVS2,       extract_extradata_mpeg4   },
 331     { AV_CODEC_ID_CAVS,       extract_extradata_mpeg4   },
 332     { AV_CODEC_ID_H264,       extract_extradata_h2645   },
 333     { AV_CODEC_ID_HEVC,       extract_extradata_h2645   },
 334     { AV_CODEC_ID_MPEG1VIDEO, extract_extradata_mpeg12  },
 335     { AV_CODEC_ID_MPEG2VIDEO, extract_extradata_mpeg12  },
 336     { AV_CODEC_ID_MPEG4,      extract_extradata_mpeg4   },
 337     { AV_CODEC_ID_VC1,        extract_extradata_vc1     },
 338 };
 339
 340 static int extract_extradata_init(AVBSFContext *ctx)
 341 {
 342     ExtractExtradataContext *s = ctx->priv_data;
 343     int i;
 344
 345     for (i = 0; i < FF_ARRAY_ELEMS(extract_tab); i++) {
 346         if (extract_tab[i].id == ctx->par_in->codec_id) {
 347             s->extract = extract_tab[i].extract;
 348             break;
 349         }
 350     }
 351     if (!s->extract)
 352         return AVERROR_BUG;
 353
 354     return 0;
 355 }
 356
 357 static int extract_extradata_filter(AVBSFContext *ctx, AVPacket *pkt)
 358 {
 359     ExtractExtradataContext *s = ctx->priv_data;
 360     uint8_t *extradata = NULL;
 361     int extradata_size;
 362     int ret = 0;
 363
 364     ret = ff_bsf_get_packet_ref(ctx, pkt);
 365     if (ret < 0)
 366         return ret;
 367
 368     ret = s->extract(ctx, pkt, &extradata, &extradata_size);
 369     if (ret < 0)
 370         goto fail;
 371
 372     if (extradata) {
 373         memset(extradata + extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 374         ret = av_packet_add_side_data(pkt, AV_PKT_DATA_NEW_EXTRADATA,
 375                                       extradata, extradata_size);
 376         if (ret < 0) {
 377             av_freep(&extradata);
 378             goto fail;
 379         }
 380     }
 381
 382     return 0;
 383
 384 fail:
 385     av_packet_unref(pkt);
 386     return ret;
 387 }
 388
 389 static void extract_extradata_close(AVBSFContext *ctx)
 390 {
 391     ExtractExtradataContext *s = ctx->priv_data;
 392     ff_av1_packet_uninit(&s->av1_pkt);
 393     ff_h2645_packet_uninit(&s->h2645_pkt);
 394 }
 395
 396 static const enum AVCodecID codec_ids[] = {
 397     AV_CODEC_ID_AV1,
 398     AV_CODEC_ID_AVS2,
 399     AV_CODEC_ID_CAVS,
 400     AV_CODEC_ID_H264,
 401     AV_CODEC_ID_HEVC,
 402     AV_CODEC_ID_MPEG1VIDEO,
 403     AV_CODEC_ID_MPEG2VIDEO,
 404     AV_CODEC_ID_MPEG4,
 405     AV_CODEC_ID_VC1,
 406     AV_CODEC_ID_NONE,
 407 };
 408
 409 #define OFFSET(x) offsetof(ExtractExtradataContext, x)
 410 #define FLAGS (AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_BSF_PARAM)
 411 static const AVOption options[] = {
 412     { "remove", "remove the extradata from the bitstream", OFFSET(remove), AV_OPT_TYPE_INT,
 413         { .i64 = 0 }, 0, 1, FLAGS },
 414     { NULL },
 415 };
 416
 417 static const AVClass extract_extradata_class = {
 418     .class_name = "extract_extradata",
 419     .item_name  = av_default_item_name,
 420     .option     = options,
 421     .version    = LIBAVUTIL_VERSION_INT,
 422 };
 423
 424 const AVBitStreamFilter ff_extract_extradata_bsf = {
 425     .name           = "extract_extradata",
 426     .codec_ids      = codec_ids,
 427     .priv_data_size = sizeof(ExtractExtradataContext),
 428     .priv_class     = &extract_extradata_class,
 429     .init           = extract_extradata_init,
 430     .filter         = extract_extradata_filter,
 431     .close          = extract_extradata_close,
 432 };