从视频文件中提取YUV数据

extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/imgutils.h>
}

#include <iostream>
#include <string>
#include <vector>
#include <error.h>

// 打开输入文件，返回 AVFormatContext 指针
AVFormatContext* openInputFile(const std::string& filename) {
    AVFormatContext* formatContext = nullptr;
    if (avformat_open_input(&formatContext, filename.c_str(), nullptr, nullptr) != 0) {
        std::cerr << "Failed to open file: " << filename << std::endl;
        return nullptr;
    }

    if (avformat_find_stream_info(formatContext, nullptr) < 0) {
        std::cerr << "Failed to find stream information" << std::endl;
        avformat_close_input(&formatContext);
        return nullptr;
    }

    return formatContext;
}

// 初始化视频解码器，返回 AVCodecContext 指针
AVCodecContext* initVideoDecoder(AVFormatContext* formatContext, int streamIndex) {
    AVCodec* codec = nullptr;
    AVCodecContext* codecContext = nullptr;

    // 查找解码器
    AVCodecParameters* codecParams = formatContext->streams[streamIndex]->codecpar;
    codec = avcodec_find_decoder(codecParams->codec_id);
    if (!codec) {
        std::cerr << "Unsupported codec!" << std::endl;
        return nullptr;
    }

    // 分配解码器上下文
    codecContext = avcodec_alloc_context3(codec);
    if (!codecContext) {
        std::cerr << "Failed to allocate codec context" << std::endl;
        return nullptr;
    }

    // 将流参数复制到解码器上下文
    if (avcodec_parameters_to_context(codecContext, codecParams) < 0) {
        std::cerr << "Failed to copy codec parameters to codec context" << std::endl;
        avcodec_free_context(&codecContext);
        return nullptr;
    }

    // 打开解码器
    if (avcodec_open2(codecContext, codec, nullptr) < 0) {
        std::cerr << "Failed to open codec" << std::endl;
        avcodec_free_context(&codecContext);
        return nullptr;
    }

    return codecContext;
}

void saveYUV(AVFrame *frame, FILE *yuvFile) {
    // 获取 YUV 分量的大小和数据
    uint8_t *y_data = frame->data[0];
    uint8_t *u_data = frame->data[1];
    uint8_t *v_data = frame->data[2];

    // 写入 Y 分量
    for (int i = 0; i < frame->height; ++i) {
        fwrite(y_data + i * frame->linesize[0], 1, frame->width, yuvFile);
    }

//    // 写入 U 分量
//    for (int i = 0; i < frame->height / 2; ++i) {
//        fwrite(u_data + i * frame->linesize[1], 1, frame->width / 2, outputFile);
//    }

//    // 写入 V 分量
//    for (int i = 0; i < frame->height / 2; ++i) {
//        fwrite(v_data + i * frame->linesize[2], 1, frame->width / 2, outputFile);
//    }
}

// 解码视频帧并输出到文件
bool decodeAndOutputFrames(AVFormatContext* formatContext, AVCodecContext* codecContext, const std::string& outputFilename) {
    AVPacket* packet = av_packet_alloc();
    if (!packet) {
        std::cerr << "Failed to allocate AVPacket" << std::endl;
        return false;
    }

    AVFrame* frame = av_frame_alloc();
    if (!frame) {
        std::cerr << "Failed to allocate AVFrame" << std::endl;
        av_packet_free(&packet);
        return false;
    }

    // 打开输出文件
    FILE* outputFile = fopen(outputFilename.c_str(), "wb");
    if (!outputFile) {
        std::cerr << "Failed to open output file: " << outputFilename << std::endl;
        av_frame_free(&frame);
        av_packet_free(&packet);
        return false;
    }

    // 解码和输出循环
    int ret = 0;
    while (av_read_frame(formatContext, packet) >= 0) {
        if (packet->stream_index == 0) {
            ret = avcodec_send_packet(codecContext, packet);
            if (ret < 0) {
                std::cerr << "Error sending a packet for decoding: " << ret << std::endl;
                break;
            }

            while (ret >= 0) {
                ret = avcodec_receive_frame(codecContext, frame);
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                    break;
                else if (ret < 0) {
                    std::cerr << "Error during decoding: " << ret << std::endl;
                    break;
                }

                saveYUV(frame, outputFile);
            }
        }

        av_packet_unref(packet);
    }

    // 清理资源
    fclose(outputFile);
    av_frame_free(&frame);
    av_packet_free(&packet);

    return true;
}

// 关闭输入文件
void closeInputFile(AVFormatContext** formatContext) {
    if (*formatContext) {
        avformat_close_input(formatContext);
        *formatContext = nullptr;
    }
}

using namespace std;


int main() {
    std::string inputFilename = "input.mp4";
    std::string outputFilename = "output.yuv";

    // 打开输入文件
    AVFormatContext* formatContext = openInputFile(inputFilename);
    if (!formatContext) {
        std::cerr << "Failed to open input file" << std::endl;
        return 1;
    }

    // 初始化视频解码器
    AVCodecContext* codecContext = initVideoDecoder(formatContext, 0);  
    if (!codecContext) {
        std::cerr << "Failed to initialize video decoder" << std::endl;
        closeInputFile(&formatContext);
        return 1;
    }

    // 解码和输出帧到文件
    bool success = decodeAndOutputFrames(formatContext, codecContext, outputFilename);
    if (!success) {
        std::cerr << "Failed to decode and output frames" << std::endl;
    }

    // 清理资源
    avcodec_free_context(&codecContext);
    closeInputFile(&formatContext);

    return 0;
}