LCOV - code coverage report
Current view: top level - libavcodec - extract_extradata_bsf.c (source / functions) Hit Total Coverage
Test: coverage.info Lines: 120 149 80.5 %
Date: 2017-12-18 06:23:41 Functions: 7 7 100.0 %

          Line data    Source code
       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 "avcodec.h"
      27             : #include "bsf.h"
      28             : #include "h2645_parse.h"
      29             : #include "h264.h"
      30             : #include "hevc.h"
      31             : #include "vc1_common.h"
      32             : 
      33             : typedef struct ExtractExtradataContext {
      34             :     const AVClass *class;
      35             : 
      36             :     int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
      37             :                    uint8_t **data, int *size);
      38             : 
      39             :     /* AVOptions */
      40             :     int remove;
      41             : } ExtractExtradataContext;
      42             : 
      43        4569 : static int val_in_array(const int *arr, int len, int val)
      44             : {
      45             :     int i;
      46       12616 :     for (i = 0; i < len; i++)
      47       10205 :         if (arr[i] == val)
      48        2158 :             return 1;
      49        2411 :     return 0;
      50             : }
      51             : 
      52         435 : static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt,
      53             :                                    uint8_t **data, int *size)
      54             : {
      55             :     static const int extradata_nal_types_hevc[] = {
      56             :         HEVC_NAL_VPS, HEVC_NAL_SPS, HEVC_NAL_PPS,
      57             :     };
      58             :     static const int extradata_nal_types_h264[] = {
      59             :         H264_NAL_SPS, H264_NAL_PPS,
      60             :     };
      61             : 
      62         435 :     ExtractExtradataContext *s = ctx->priv_data;
      63             : 
      64         435 :     H2645Packet h2645_pkt = { 0 };
      65         435 :     int extradata_size = 0;
      66             :     const int *extradata_nal_types;
      67             :     int nb_extradata_nal_types;
      68         435 :     int i, has_sps = 0, has_vps = 0, ret = 0;
      69             : 
      70         435 :     if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) {
      71         195 :         extradata_nal_types    = extradata_nal_types_hevc;
      72         195 :         nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_hevc);
      73             :     } else {
      74         240 :         extradata_nal_types    = extradata_nal_types_h264;
      75         240 :         nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_h264);
      76             :     }
      77             : 
      78         435 :     ret = ff_h2645_packet_split(&h2645_pkt, pkt->data, pkt->size,
      79         435 :                                 ctx, 0, 0, ctx->par_in->codec_id, 1);
      80         435 :     if (ret < 0)
      81           0 :         goto fail;
      82             : 
      83        2801 :     for (i = 0; i < h2645_pkt.nb_nals; i++) {
      84        2366 :         H2645NAL *nal = &h2645_pkt.nals[i];
      85        2366 :         if (val_in_array(extradata_nal_types, nb_extradata_nal_types, nal->type)) {
      86        1079 :             extradata_size += nal->raw_size + 3;
      87        1079 :             if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) {
      88         639 :                 if (nal->type == HEVC_NAL_SPS) has_sps = 1;
      89         639 :                 if (nal->type == HEVC_NAL_VPS) has_vps = 1;
      90             :             } else {
      91         440 :                 if (nal->type == H264_NAL_SPS) has_sps = 1;
      92             :             }
      93             :         }
      94             :     }
      95             : 
      96         835 :     if (extradata_size &&
      97         800 :         ((ctx->par_in->codec_id == AV_CODEC_ID_HEVC && has_sps && has_vps) ||
      98         432 :          (ctx->par_in->codec_id == AV_CODEC_ID_H264 && has_sps))) {
      99             :         AVBufferRef *filtered_buf;
     100             :         uint8_t *extradata, *filtered_data;
     101             : 
     102         400 :         if (s->remove) {
     103           0 :             filtered_buf = av_buffer_alloc(pkt->size + AV_INPUT_BUFFER_PADDING_SIZE);
     104           0 :             if (!filtered_buf) {
     105           0 :                 ret = AVERROR(ENOMEM);
     106           0 :                 goto fail;
     107             :             }
     108           0 :             filtered_data = filtered_buf->data;
     109             :         }
     110             : 
     111         400 :         extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
     112         400 :         if (!extradata) {
     113           0 :             av_buffer_unref(&filtered_buf);
     114           0 :             ret = AVERROR(ENOMEM);
     115           0 :             goto fail;
     116             :         }
     117             : 
     118         400 :         *data = extradata;
     119         400 :         *size = extradata_size;
     120             : 
     121        2603 :         for (i = 0; i < h2645_pkt.nb_nals; i++) {
     122        2203 :             H2645NAL *nal = &h2645_pkt.nals[i];
     123        2203 :             if (val_in_array(extradata_nal_types, nb_extradata_nal_types,
     124             :                              nal->type)) {
     125        1079 :                 AV_WB24(extradata, 1); // startcode
     126        1079 :                 memcpy(extradata + 3, nal->raw_data, nal->raw_size);
     127        1079 :                 extradata += 3 + nal->raw_size;
     128        1124 :             } else if (s->remove) {
     129           0 :                 AV_WB24(filtered_data, 1); // startcode
     130           0 :                 memcpy(filtered_data + 3, nal->raw_data, nal->raw_size);
     131           0 :                 filtered_data += 3 + nal->raw_size;
     132             :             }
     133             :         }
     134             : 
     135         400 :         if (s->remove) {
     136           0 :             av_buffer_unref(&pkt->buf);
     137           0 :             pkt->buf  = filtered_buf;
     138           0 :             pkt->data = filtered_buf->data;
     139           0 :             pkt->size = filtered_data - filtered_buf->data;
     140             :         }
     141             :     }
     142             : 
     143         470 : fail:
     144         435 :     ff_h2645_packet_uninit(&h2645_pkt);
     145         435 :     return ret;
     146             : }
     147             : 
     148           6 : static int extract_extradata_vc1(AVBSFContext *ctx, AVPacket *pkt,
     149             :                                  uint8_t **data, int *size)
     150             : {
     151           6 :     ExtractExtradataContext *s = ctx->priv_data;
     152           6 :     const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
     153           6 :     uint32_t state = UINT32_MAX;
     154           6 :     int has_extradata = 0, extradata_size = 0;
     155             : 
     156          29 :     while (ptr < end) {
     157          23 :         ptr = avpriv_find_start_code(ptr, end, &state);
     158          23 :         if (state == VC1_CODE_SEQHDR || state == VC1_CODE_ENTRYPOINT) {
     159          17 :             has_extradata = 1;
     160           6 :         } else if (has_extradata && IS_MARKER(state)) {
     161           6 :             extradata_size = ptr - 4 - pkt->data;
     162           6 :             break;
     163             :         }
     164             :     }
     165             : 
     166           6 :     if (extradata_size) {
     167           6 :         *data = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
     168           6 :         if (!*data)
     169           0 :             return AVERROR(ENOMEM);
     170             : 
     171           6 :         memcpy(*data, pkt->data, extradata_size);
     172           6 :         *size = extradata_size;
     173             : 
     174           6 :         if (s->remove) {
     175           0 :             pkt->data += extradata_size;
     176           0 :             pkt->size -= extradata_size;
     177             :         }
     178             :     }
     179             : 
     180           6 :     return 0;
     181             : }
     182             : 
     183         368 : static int extract_extradata_mpeg12(AVBSFContext *ctx, AVPacket *pkt,
     184             :                                      uint8_t **data, int *size)
     185             : {
     186         368 :     ExtractExtradataContext *s = ctx->priv_data;
     187         368 :     uint32_t state = UINT32_MAX;
     188         368 :     int i, found = 0;
     189             : 
     190     2397393 :     for (i = 0; i < pkt->size; i++) {
     191     2397191 :         state = (state << 8) | pkt->data[i];
     192     2397191 :         if (state == 0x1B3)
     193         166 :             found = 1;
     194     2397025 :         else if (found && state != 0x1B5 && state < 0x200 && state >= 0x100) {
     195         166 :             if (i > 3) {
     196         166 :                 *size = i - 3;
     197         166 :                 *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
     198         166 :                 if (!*data)
     199           0 :                     return AVERROR(ENOMEM);
     200             : 
     201         166 :                 memcpy(*data, pkt->data, *size);
     202             : 
     203         166 :                 if (s->remove) {
     204           0 :                     pkt->data += *size;
     205           0 :                     pkt->size -= *size;
     206             :                 }
     207             :             }
     208         166 :             break;
     209             :         }
     210             :     }
     211         368 :     return 0;
     212             : }
     213             : 
     214         115 : static int extract_extradata_mpeg4(AVBSFContext *ctx, AVPacket *pkt,
     215             :                                    uint8_t **data, int *size)
     216             : {
     217         115 :     ExtractExtradataContext *s = ctx->priv_data;
     218         115 :     const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
     219         115 :     uint32_t state = UINT32_MAX;
     220             : 
     221         489 :     while (ptr < end) {
     222         374 :         ptr = avpriv_find_start_code(ptr, end, &state);
     223         374 :         if (state == 0x1B3 || state == 0x1B6) {
     224         115 :             if (ptr - pkt->data > 4) {
     225          65 :                 *size = ptr - 4 - pkt->data;
     226          65 :                 *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
     227          65 :                 if (!*data)
     228           0 :                     return AVERROR(ENOMEM);
     229             : 
     230          65 :                 memcpy(*data, pkt->data, *size);
     231             : 
     232          65 :                 if (s->remove) {
     233           0 :                     pkt->data += *size;
     234           0 :                     pkt->size -= *size;
     235             :                 }
     236             :             }
     237         115 :             break;
     238             :         }
     239             :     }
     240         115 :     return 0;
     241             : }
     242             : 
     243             : static const struct {
     244             :     enum AVCodecID id;
     245             :     int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
     246             :                    uint8_t **data, int *size);
     247             : } extract_tab[] = {
     248             :     { AV_CODEC_ID_CAVS,       extract_extradata_mpeg4   },
     249             :     { AV_CODEC_ID_H264,       extract_extradata_h2645   },
     250             :     { AV_CODEC_ID_HEVC,       extract_extradata_h2645   },
     251             :     { AV_CODEC_ID_MPEG1VIDEO, extract_extradata_mpeg12  },
     252             :     { AV_CODEC_ID_MPEG2VIDEO, extract_extradata_mpeg12  },
     253             :     { AV_CODEC_ID_MPEG4,      extract_extradata_mpeg4   },
     254             :     { AV_CODEC_ID_VC1,        extract_extradata_vc1     },
     255             : };
     256             : 
     257         640 : static int extract_extradata_init(AVBSFContext *ctx)
     258             : {
     259         640 :     ExtractExtradataContext *s = ctx->priv_data;
     260             :     int i;
     261             : 
     262        2236 :     for (i = 0; i < FF_ARRAY_ELEMS(extract_tab); i++) {
     263        2236 :         if (extract_tab[i].id == ctx->par_in->codec_id) {
     264         640 :             s->extract = extract_tab[i].extract;
     265         640 :             break;
     266             :         }
     267             :     }
     268         640 :     if (!s->extract)
     269           0 :         return AVERROR_BUG;
     270             : 
     271         640 :     return 0;
     272             : }
     273             : 
     274        1211 : static int extract_extradata_filter(AVBSFContext *ctx, AVPacket *out)
     275             : {
     276        1211 :     ExtractExtradataContext *s = ctx->priv_data;
     277             :     AVPacket *in;
     278        1211 :     uint8_t *extradata = NULL;
     279             :     int extradata_size;
     280        1211 :     int ret = 0;
     281             : 
     282        1211 :     ret = ff_bsf_get_packet(ctx, &in);
     283        1211 :     if (ret < 0)
     284         287 :         return ret;
     285             : 
     286         924 :     ret = s->extract(ctx, in, &extradata, &extradata_size);
     287         924 :     if (ret < 0)
     288           0 :         goto fail;
     289             : 
     290         924 :     if (extradata) {
     291         637 :         ret = av_packet_add_side_data(in, AV_PKT_DATA_NEW_EXTRADATA,
     292             :                                       extradata, extradata_size);
     293         637 :         if (ret < 0) {
     294           0 :             av_freep(&extradata);
     295           0 :             goto fail;
     296             :         }
     297             :     }
     298             : 
     299         924 :     av_packet_move_ref(out, in);
     300             : 
     301         924 : fail:
     302         924 :     av_packet_free(&in);
     303         924 :     return ret;
     304             : }
     305             : 
     306             : static const enum AVCodecID codec_ids[] = {
     307             :     AV_CODEC_ID_CAVS,
     308             :     AV_CODEC_ID_H264,
     309             :     AV_CODEC_ID_HEVC,
     310             :     AV_CODEC_ID_MPEG1VIDEO,
     311             :     AV_CODEC_ID_MPEG2VIDEO,
     312             :     AV_CODEC_ID_MPEG4,
     313             :     AV_CODEC_ID_VC1,
     314             :     AV_CODEC_ID_NONE,
     315             : };
     316             : 
     317             : #define OFFSET(x) offsetof(ExtractExtradataContext, x)
     318             : static const AVOption options[] = {
     319             :     { "remove", "remove the extradata from the bitstream", OFFSET(remove), AV_OPT_TYPE_INT,
     320             :         { .i64 = 0 }, 0, 1 },
     321             :     { NULL },
     322             : };
     323             : 
     324             : static const AVClass extract_extradata_class = {
     325             :     .class_name = "extract_extradata",
     326             :     .item_name  = av_default_item_name,
     327             :     .option     = options,
     328             :     .version    = LIBAVUTIL_VERSION_INT,
     329             : };
     330             : 
     331             : const AVBitStreamFilter ff_extract_extradata_bsf = {
     332             :     .name           = "extract_extradata",
     333             :     .codec_ids      = codec_ids,
     334             :     .priv_data_size = sizeof(ExtractExtradataContext),
     335             :     .priv_class     = &extract_extradata_class,
     336             :     .init           = extract_extradata_init,
     337             :     .filter         = extract_extradata_filter,
     338             : };

Generated by: LCOV version 1.13