libjpeg + FFmpeg API を使用してビデオのフレーム画像を保存する

はじめに

前回の記事では FFmpeg API のエンコーダーを使い、あらかじめ作成したフレームをビデオに書き込むためのコードを書きました。
今回の記事では、デコーダーから受け取った映像ファイルのフレームを画像として保存するキャプチャーのコードを作ってみたいと思います。

コード

というわけで作成したコードがこちら。
手順としては、デコーダーから受け取ったフレームをlibjpegを使って保存するということをやっています。
(2021/09/19 追記：libjpeg関連の箇所を変更しました)

#include <iostream>
#include <string>
#include <filesystem>
extern "C"{
    #include <libavformat/avformat.h>
    #include <libavcodec/avcodec.h>
    #include <libavutil/imgutils.h>
    #include <libavutil/opt.h>
    #include <libswscale/swscale.h>
    #include <libswresample/swresample.h>
    #include <jpeglib.h>
}

int main(int argc, char *argv[]){
    std::string dir_sep = "/";
    #if defined(_WIN32)
        dir_sep = "\\";
    #endif
    //使用するビデオ
    const char *input = argv[1];
    //フレーム画像の画質
    int quality = 75;
    if (argv[2]){
        quality = atoi(argv[2]);
    }
    //保存先のディレクトリ
    std::string dir = "output";
    std::filesystem::create_directory(dir);
    AVFormatContext *inputFmtContxt = NULL;
    const AVCodec *decoder = NULL;
    AVCodecContext *decoderContxt = NULL;
    int ret = 0, video_stream_index = 0;
    ret = avformat_open_input(&inputFmtContxt, input, NULL, NULL);
    if (ret < 0){
        std::cout << "Could not open input video" << std::endl;
    }
    ret = avformat_find_stream_info(inputFmtContxt, NULL);
    if (ret < 0){
        std::cout << "Could not find the stream info" << std::endl;
    }
    //デコーダーの設定
    for (int i=0; i<(int)inputFmtContxt->nb_streams; ++i){
        AVStream *in_stream = inputFmtContxt->streams[i];
        AVCodecParameters *in_par = in_stream->codecpar;
        if (in_par->codec_type == AVMEDIA_TYPE_VIDEO){
            video_stream_index = i;
            decoder = avcodec_find_decoder(in_par->codec_id);
            decoderContxt = avcodec_alloc_context3(decoder);
            avcodec_parameters_to_context(decoderContxt, in_par);
            decoderContxt->framerate = in_stream->r_frame_rate;
            decoderContxt->time_base = in_stream->time_base;
            avcodec_open2(decoderContxt, decoder, NULL);
        }
    }
    //YUVからRGBへの変換用
    enum AVPixelFormat pix_fmt = AV_PIX_FMT_RGB24;
    int HEIGHT = decoderContxt->height;
    int WIDTH = decoderContxt->width;
    SwsContext *scaler = sws_getContext(WIDTH, HEIGHT, decoderContxt->pix_fmt, 
                                        WIDTH, HEIGHT, pix_fmt, SWS_BICUBIC, NULL, NULL, NULL);

    //パケットとフレームの準備
    int res = 0;
    AVPacket *packet = av_packet_alloc();
    //デコーダーから受け取るフレーム
    AVFrame *frame = av_frame_alloc();
    //RGBへの変換先のフレーム
    AVFrame *rgbframe = av_frame_alloc();
    rgbframe->width = decoderContxt->width;
    rgbframe->height = decoderContxt->height;
    rgbframe->format = pix_fmt;
    ret = av_frame_get_buffer(rgbframe, 0);
    uint8_t *buf = (uint8_t*) av_malloc(av_image_get_buffer_size(pix_fmt, decoderContxt->width, decoderContxt->height, 1));
    ret = av_image_fill_arrays(rgbframe->data, rgbframe->linesize, buf, pix_fmt, decoderContxt->width, decoderContxt->height, 1);
    //デコードとキャプチャの開始
    int count = 0;
    while (true){
        ret = av_read_frame(inputFmtContxt, packet);
        if (ret < 0){
            break;
        }
        AVStream *input_stream = inputFmtContxt->streams[packet->stream_index];
        if (input_stream->codecpar->codec_type == video_stream_index){
            res = avcodec_send_packet(decoderContxt, packet);
            while (res >= 0){
                res = avcodec_receive_frame(decoderContxt, frame);
                if (res == AVERROR(EAGAIN) || res == AVERROR_EOF){
                    break;
                }
                if (res >= 0){
                    //YUVフレームをRGBに変換
                    sws_scale(scaler, frame->data, frame->linesize, 0, frame->height, rgbframe->data, rgbframe->linesize);
                    ++count;
                    std::string filename = dir + dir_sep + "frame_" + std::to_string(count) + ".png";
                    //JPEG画像として保存
                    struct jpeg_compress_struct cinfo;
                    struct jpeg_error_mgr jerr;
                    cinfo.err = jpeg_std_error(&jerr);
                    jpeg_create_compress(&cinfo);
                    FILE *f = fopen(filename.c_str(), "wb");
                    int stride = rgbframe->linesize[0];
                    jpeg_stdio_dest(&cinfo, f);
                    cinfo.image_width = rgbframe->width;
                    cinfo.image_height = rgbframe->height;
                    cinfo.input_components = 3;
                    cinfo.in_color_space = JCS_RGB;
                    jpeg_set_defaults(&cinfo);
                    jpeg_set_quality(&cinfo, quality, TRUE);
                    jpeg_start_compress(&cinfo, TRUE);
                    uint8_t *row = rgbframe->data[0];
                    for (int i=0; i<rgbframe->height; ++i){
                        jpeg_write_scanlines(&cinfo, &row, 1);
                        row += stride;
                    }
                    jpeg_finish_compress(&cinfo);
                    jpeg_destroy_compress(&cinfo);
                    fclose(f);
                }
            }
            av_frame_unref(frame);
        }
        av_packet_unref(packet);
    }
    //各メモリの解放
    av_packet_free(&packet);
    av_frame_free(&frame);
    av_frame_free(&rgbframe);
    avformat_free_context(inputFmtContxt);
    avcodec_free_context(&decoderContxt);
    av_freep(&buf);
    sws_freeContext(scaler);
    return 0;
}