h264_mp4toannexb_bsf.c 6.53 KB
Newer Older
1
/*
2
 * H.264 MP4 to Annex B byte stream format filter
3
 * Copyright (c) 2007 Benoit Fouet <benoit.fouet@free.fr>
4
 *
5
 * This file is part of Libav.
6
 *
7
 * Libav is free software; you can redistribute it and/or
8 9 10 11
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
12
 * Libav is distributed in the hope that it will be useful,
13 14 15 16 17
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
18
 * License along with Libav; if not, write to the Free Software
19 20 21
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */

22 23
#include <string.h>

24
#include "libavutil/intreadwrite.h"
25
#include "libavutil/mem.h"
26 27 28 29 30
#include "avcodec.h"

typedef struct H264BSFContext {
    uint8_t  length_size;
    uint8_t  first_idr;
31
    int      extradata_parsed;
32 33
} H264BSFContext;

34
static int alloc_and_copy(uint8_t **poutbuf, int *poutbuf_size,
35
                          const uint8_t *sps_pps, uint32_t sps_pps_size,
36 37 38
                          const uint8_t *in, uint32_t in_size)
{
    uint32_t offset         = *poutbuf_size;
39
    uint8_t nal_header_size = offset ? 3 : 4;
40
    int err;
41

42
    *poutbuf_size += sps_pps_size + in_size + nal_header_size;
43 44 45 46 47
    if ((err = av_reallocp(poutbuf,
                           *poutbuf_size + FF_INPUT_BUFFER_PADDING_SIZE)) < 0) {
        *poutbuf_size = 0;
        return err;
    }
48
    if (sps_pps)
49 50
        memcpy(*poutbuf + offset, sps_pps, sps_pps_size);
    memcpy(*poutbuf + sps_pps_size + nal_header_size + offset, in, in_size);
51
    if (!offset) {
52
        AV_WB32(*poutbuf + sps_pps_size, 1);
53
    } else {
54 55 56
        (*poutbuf + offset + sps_pps_size)[0] =
        (*poutbuf + offset + sps_pps_size)[1] = 0;
        (*poutbuf + offset + sps_pps_size)[2] = 1;
57
    }
58 59

    return 0;
60 61
}

62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87
static int h264_extradata_to_annexb(AVCodecContext *avctx, const int padding)
{
    uint16_t unit_size;
    uint64_t total_size                 = 0;
    uint8_t *out                        = NULL, unit_nb, sps_done = 0,
             sps_seen                   = 0, pps_seen = 0;
    const uint8_t *extradata            = avctx->extradata + 4;
    static const uint8_t nalu_header[4] = { 0, 0, 0, 1 };
    int length_size = (*extradata++ & 0x3) + 1; // retrieve length coded size

    if (length_size == 3)
        return AVERROR(EINVAL);

    /* retrieve sps and pps unit(s) */
    unit_nb = *extradata++ & 0x1f; /* number of sps unit(s) */
    if (!unit_nb) {
        unit_nb = *extradata++; /* number of pps unit(s) */
        sps_done++;

        if (unit_nb)
            pps_seen = 1;
    } else {
        sps_seen = 1;
    }

    while (unit_nb--) {
88
        int err;
89 90 91 92 93 94 95 96 97

        unit_size   = AV_RB16(extradata);
        total_size += unit_size + 4;
        if (total_size > INT_MAX - padding ||
            extradata + 2 + unit_size > avctx->extradata +
            avctx->extradata_size) {
            av_free(out);
            return AVERROR(EINVAL);
        }
98 99
        if ((err = av_reallocp(&out, total_size + padding)) < 0)
            return err;
100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130
        memcpy(out + total_size - unit_size - 4, nalu_header, 4);
        memcpy(out + total_size - unit_size, extradata + 2, unit_size);
        extradata += 2 + unit_size;

        if (!unit_nb && !sps_done++) {
            unit_nb = *extradata++; /* number of pps unit(s) */
            if (unit_nb)
                pps_seen = 1;
        }
    }

    if (out)
        memset(out + total_size, 0, FF_INPUT_BUFFER_PADDING_SIZE);

    if (!sps_seen)
        av_log(avctx, AV_LOG_WARNING,
               "Warning: SPS NALU missing or invalid. "
               "The resulting stream may not play.\n");

    if (!pps_seen)
        av_log(avctx, AV_LOG_WARNING,
               "Warning: PPS NALU missing or invalid. "
               "The resulting stream may not play.\n");

    av_free(avctx->extradata);
    avctx->extradata      = out;
    avctx->extradata_size = total_size;

    return length_size;
}

131 132
static int h264_mp4toannexb_filter(AVBitStreamFilterContext *bsfc,
                                   AVCodecContext *avctx, const char *args,
133 134 135 136
                                   uint8_t **poutbuf, int *poutbuf_size,
                                   const uint8_t *buf, int buf_size,
                                   int keyframe)
{
137 138
    H264BSFContext *ctx = bsfc->priv_data;
    uint8_t unit_type;
139
    int32_t nal_size;
140
    uint32_t cumul_size    = 0;
141
    const uint8_t *buf_end = buf + buf_size;
142
    int ret = 0;
143 144 145

    /* nothing to filter */
    if (!avctx->extradata || avctx->extradata_size < 6) {
146
        *poutbuf      = (uint8_t *)buf;
147 148 149 150 151
        *poutbuf_size = buf_size;
        return 0;
    }

    /* retrieve sps and pps NAL units from extradata */
152
    if (!ctx->extradata_parsed) {
153 154 155 156
        ret = h264_extradata_to_annexb(avctx, FF_INPUT_BUFFER_PADDING_SIZE);
        if (ret < 0)
            return ret;
        ctx->length_size      = ret;
157 158
        ctx->first_idr        = 1;
        ctx->extradata_parsed = 1;
159 160 161
    }

    *poutbuf_size = 0;
162
    *poutbuf      = NULL;
163
    do {
164 165 166
        if (buf + ctx->length_size > buf_end)
            goto fail;

167
        if (ctx->length_size == 1) {
168
            nal_size = buf[0];
169
        } else if (ctx->length_size == 2) {
170
            nal_size = AV_RB16(buf);
171
        } else
172 173
            nal_size = AV_RB32(buf);

174
        buf      += ctx->length_size;
175 176
        unit_type = *buf & 0x1f;

177 178 179
        if (buf + nal_size > buf_end || nal_size < 0)
            goto fail;

180 181
        /* prepend only to the first type 5 NAL unit of an IDR picture */
        if (ctx->first_idr && unit_type == 5) {
182
            if (alloc_and_copy(poutbuf, poutbuf_size,
183
                               avctx->extradata, avctx->extradata_size,
184 185
                               buf, nal_size) < 0)
                goto fail;
186
            ctx->first_idr = 0;
187
        } else {
188
            if (alloc_and_copy(poutbuf, poutbuf_size,
189
                               NULL, 0, buf, nal_size) < 0)
190
                goto fail;
191 192 193 194
            if (!ctx->first_idr && unit_type == 1)
                ctx->first_idr = 1;
        }

195
        buf        += nal_size;
196 197 198 199
        cumul_size += nal_size + ctx->length_size;
    } while (cumul_size < buf_size);

    return 1;
200 201 202 203 204

fail:
    av_freep(poutbuf);
    *poutbuf_size = 0;
    return AVERROR(EINVAL);
205 206
}

207
AVBitStreamFilter ff_h264_mp4toannexb_bsf = {
208 209 210 211
    "h264_mp4toannexb",
    sizeof(H264BSFContext),
    h264_mp4toannexb_filter,
};