示例

#include <iostream>
#include <vector>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/opt.h>
#include <libavutil/time.h>
#include <libavutil/avassert.h>
#include <cstdint>
#include <cstring>
#include <cstdio>  // for file input

// 错误处理宏
#define CHECK_RET(ret, msg) do { \
    if (ret < 0) { \
        std::cerr << msg << ": " << av_err2str(ret) << std::endl; \
        goto end; \
    } \
} while (0)

// 查找 Annex B start code (00 00 00 01 或 00 00 01)
static int find_start_code(const uint8_t* buf, int size, int* start_code_len) {
    if (size >= 4 && buf[0] == 0 && buf[1] == 0 && buf[2] == 0 && buf[3] == 1) {
        *start_code_len = 4;
        return 0;
    }
    if (size >= 3 && buf[0] == 0 && buf[1] == 0 && buf[2] == 1) {
        *start_code_len = 3;
        return 0;
    }
    return -1;
}

// 解析 VPS/SPS/PPS 并存储到 extradata
static int parse_nalu_extradata(AVCodecParserContext* parser, AVCodecContext* codec_ctx,
                                const uint8_t* data, int size,
                                uint8_t** extradata, int* extradata_size) {
    std::vector<uint8_t> extradata_buf;
    int offset = 0;

    while (offset < size) {
        int start_code_len = 0;
        if (find_start_code(data + offset, size - offset, &start_code_len) < 0) {
            offset++;
            continue;
        }

        // 找到下一个 start code 或数据末尾
        int next_offset = offset + start_code_len;
        while (next_offset < size) {
            if (find_start_code(data + next_offset, size - next_offset, &start_code_len) == 0) {
                break;
            }
            next_offset++;
        }
        if (next_offset == size) {
            next_offset = size;  // 最后一个 NALU
        }

        int nalu_size = next_offset - offset;
        uint8_t* nalu_data = (uint8_t*)data + offset;
        int nalu_type = codec_ctx->codec_id == AV_CODEC_ID_H264 ?
                        (nalu_data[start_code_len] & 0x1F) :
                        ((nalu_data[start_code_len] >> 1) & 0x3F);

        // H.264: SPS (7), PPS (8); H.265: VPS (32), SPS (33), PPS (34)
        bool is_extradata_nalu = false;
        if (codec_ctx->codec_id == AV_CODEC_ID_H264 &&
            (nalu_type == 7 || nalu_type == 8)) {
            is_extradata_nalu = true;
        } else if (codec_ctx->codec_id == AV_CODEC_ID_HEVC &&
                   (nalu_type == 32 || nalu_type == 33 || nalu_type == 34)) {
            is_extradata_nalu = true;
        }

        if (is_extradata_nalu) {
            extradata_buf.insert(extradata_buf.end(), nalu_data, nalu_data + nalu_size);
        }

        offset = next_offset;
    }

    if (extradata_buf.empty()) {
        return -1;  // 未找到 extradata
    }

    *extradata_size = extradata_buf.size();
    *extradata = (uint8_t*)av_malloc(*extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
    if (!*extradata) {
        return AVERROR(ENOMEM);
    }
    memcpy(*extradata, extradata_buf.data(), *extradata_size);
    memset(*extradata + *extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
    return 0;
}

int main(int argc, char** argv) {
    if (argc != 4) {
        std::cerr << "Usage: " << argv[0] << " <input_codec> <input_file> <output_fmp4>\n";
        std::cerr << "input_codec: h264 or hevc\n";
        std::cerr << "input_file: path to raw H.264/H.265 stream or 'udp://host:port'\n";
        return -1;
    }

    const char* codec_name = argv[1];  // "h264" or "hevc"
    const char* input_filename = argv[2];
    const char* output_filename = argv[3];

    AVFormatContext* fmt_ctx = nullptr;
    AVStream* stream = nullptr;
    AVCodecParserContext* parser = nullptr;
    AVCodecContext* codec_ctx = nullptr;
    AVPacket pkt;
    int ret = 0;
    AVCodecID codec_id = strcmp(codec_name, "h264") == 0 ? AV_CODEC_ID_H264 : AV_CODEC_ID_HEVC;

    // 1. 初始化 FFmpeg
    avformat_network_init();

    // 2. 分配输出上下文 (mp4 muxer)
    ret = avformat_alloc_output_context2(&fmt_ctx, nullptr, "mp4", output_filename);
    CHECK_RET(ret, "Failed to allocate output context");

    // 3. 创建视频流
    stream = avformat_new_stream(fmt_ctx, nullptr);
    if (!stream) {
        std::cerr << "Failed to create stream" << std::endl;
        ret = AVERROR_UNKNOWN;
        goto end;
    }
    stream->id = fmt_ctx->nb_streams - 1;

    // 4. 初始化 codec 参数和解析器
    const AVCodec* codec = avcodec_find_decoder(codec_id);
    if (!codec) {
        std::cerr << "Decoder not found for " << codec_name << std::endl;
        ret = AVERROR_DEMUXER_NOT_FOUND;
        goto end;
    }

    parser = av_parser_init(codec_id);
    if (!parser) {
        std::cerr << "Failed to init parser" << std::endl;
        ret = AVERROR_UNKNOWN;
        goto end;
    }

    codec_ctx = avcodec_alloc_context3(codec);
    if (!codec_ctx) {
        std::cerr << "Failed to allocate codec context" << std::endl;
        ret = AVERROR(ENOMEM);
        goto end;
    }

    // 设置 codecpar（稍后填充 extradata）
    AVCodecParameters* codecpar = stream->codecpar;
    codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    codecpar->codec_id = codec_id;
    codecpar->format = AV_PIX_FMT_YUV420P;  // 假设
    codecpar->width = 1920;   // 示例；可从 SPS 解析
    codecpar->height = 1080;
    codecpar->bit_rate = 5000000;

    // 设置时间基
    stream->time_base = {1, 30};  // 30fps
    codec_ctx->time_base = stream->time_base;

    // 5. 打开输入（文件或 UDP）
    FILE* input_file = nullptr;
    if (strncmp(input_filename, "udp://", 6) != 0) {
        input_file = fopen(input_filename, "rb");
        if (!input_file) {
            std::cerr << "Failed to open input file: " << input_filename << std::endl;
            ret = AVERROR(EIO);
            goto end;
        }
    } else {
        // UDP 示例：替换为 avformat_open_input
        /*
        AVFormatContext* input_ctx = nullptr;
        ret = avformat_open_input(&input_ctx, input_filename, nullptr, nullptr);
        CHECK_RET(ret, "Failed to open UDP input");
        // 读取逻辑替换 fread
        */
    }

    // 6. 解析 extradata（VPS/SPS/PPS）
    uint8_t buffer[1024 * 1024];  // 1MB 缓冲
    size_t bytes_read = fread(buffer, 1, sizeof(buffer), input_file);
    if (bytes_read == 0) {
        std::cerr << "Failed to read input stream" << std::endl;
        ret = AVERROR(EIO);
        goto end;
    }

    uint8_t* extradata = nullptr;
    int extradata_size = 0;
    ret = parse_nalu_extradata(parser, codec_ctx, buffer, bytes_read, &extradata, &extradata_size);
    CHECK_RET(ret, "Failed to parse extradata");

    codecpar->extradata = extradata;
    codecpar->extradata_size = extradata_size;

    // 7. 打开输出文件
    if (!(fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
        ret = avio_open(&fmt_ctx->pb, output_filename, AVIO_FLAG_WRITE);
        CHECK_RET(ret, "Failed to open output file");
    }

    // 8. fMP4 选项
    AVDictionary* opt = nullptr;
    av_dict_set(&opt, "movflags", "frag_keyframe+empty_moov+default_base_moof", 0);
    av_dict_set(&opt, "fragment_duration", "4000", 0);  // 4s 片段
    av_dict_set(&opt, "flush_packets", "1", 0);
    av_dict_set(&opt, "use_editlist", "0", 0);  // 避免 edit list

    // 9. 写头
    ret = avformat_write_header(fmt_ctx, &opt);
    av_dict_free(&opt);
    CHECK_RET(ret, "Failed to write header");

    // 10. 实时循环：读取裸流，封装 AVPacket
    av_init_packet(&pkt);
    int64_t frame_count = 0;
    int64_t start_time = av_gettime();
    fseek(input_file, 0, SEEK_SET);  // 重置文件指针

    while ((bytes_read = fread(buffer, 1, sizeof(buffer), input_file)) > 0) {
        int offset = 0;
        while (offset < bytes_read) {
            int start_code_len = 0;
            if (find_start_code(buffer + offset, bytes_read - offset, &start_code_len) < 0) {
                offset++;
                continue;
            }

            int next_offset = offset + start_code_len;
            while (next_offset < bytes_read) {
                if (find_start_code(buffer + next_offset, bytes_read - next_offset, &start_code_len) == 0) {
                    break;
                }
                next_offset++;
            }
            if (next_offset == bytes_read) {
                next_offset = bytes_read;
            }

            // 封装 AVPacket
            pkt.data = buffer + offset;
            pkt.size = next_offset - offset;
            pkt.stream_index = stream->id;

            // 时间戳：基于实时时间
            int64_t current_time = av_gettime();
            int64_t time_diff = (current_time - start_time) / 1000000.0 * 30;  // 30fps
            pkt.pts = time_diff;
            pkt.dts = pkt.pts;  // 假设无 B 帧
            pkt.duration = 1;

            // 检查关键帧 (H.264: IDR=5; H.265: IDR=19-21)
            int nalu_type = codec_id == AV_CODEC_ID_H264 ?
                            (buffer[offset + start_code_len] & 0x1F) :
                            ((buffer[offset + start_code_len] >> 1) & 0x3F);
            if ((codec_id == AV_CODEC_ID_H264 && nalu_type == 5) ||
                (codec_id == AV_CODEC_ID_HEVC && (nalu_type >= 19 && nalu_type <= 21))) {
                pkt.flags |= AV_PKT_FLAG_KEY;
            }

            // 写包
            ret = av_interleaved_write_frame(fmt_ctx, &pkt);
            if (ret < 0) {
                std::cerr << "Failed to write frame: " << av_err2str(ret) << std::endl;
                break;
            }

            offset = next_offset;
            frame_count++;
            if (frame_count % 30 == 0) {
                std::cout << "Muxed " << frame_count << " frames" << std::endl;
            }
        }
    }

    // 11. 写尾
    av_write_trailer(fmt_ctx);

end:
    if (extradata) av_freep(&extradata);
    if (input_file) fclose(input_file);
    if (parser) av_parser_close(parser);
    if (codec_ctx) avcodec_free_context(&codec_ctx);
    if (fmt_ctx && !(fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
        avio_closep(&fmt_ctx->pb);
    }
    avformat_free_context(fmt_ctx);
    return ret < 0 ? ret : 0;
}