Video FFMpeg使用指南 (v2.0)

#include <iostream>
extern "C"
{
	#include <libavformat/avformat.h>
	#include <libavcodec/avcodec.h>
	#include <libswscale/swscale.h>
}

static enum AVPixelFormat get_hw_format(AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts) {
    for (const enum AVPixelFormat *p = pix_fmts; *p != -1; p++) {
        if (*p == AV_PIX_FMT_CUDA)
		{
			std::cout<<"receive cuda fmt"<<std::endl;
            return *p;
		}
    }
    return AV_PIX_FMT_NONE;
}

static const char* get_cuvid_decoder_name(AVCodecID codec_id) {
    switch(codec_id) {
        case AV_CODEC_ID_H264:
            return "h264_cuvid";
        case AV_CODEC_ID_HEVC:
            return "hevc_cuvid";
        case AV_CODEC_ID_VP9:
            return "vp9_cuvid";
        case AV_CODEC_ID_AV1:
            return "av1_cuvid";
        default:
            return "";
    }
}

int main(int argc, char** argv) {
    if (argc < 2) {
        std::cout << "Usage: " << argv[0] << " <input.hevc>" << std::endl;
        return -1;
    }


    AVFormatContext* fmt_ctx = nullptr;
    if (avformat_open_input(&fmt_ctx, argv[1], nullptr, nullptr) != 0) {
        std::cout << "Could not open input file" << std::endl;
        return -1;
    }

    if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
        std::cout << "Could not find stream information" << std::endl;
        avformat_close_input(&fmt_ctx);
        return -1;
    }


    int video_stream_idx = -1;
    AVCodecParameters* codec_params = nullptr;
    for (unsigned int i = 0; i < fmt_ctx->nb_streams; i++) {
        if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_idx = i;
            codec_params = fmt_ctx->streams[i]->codecpar;
            break;
        }
    }

    if (video_stream_idx == -1) {
        std::cout << "Could not find video stream" << std::endl;
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    AVHWDeviceType type = av_hwdevice_find_type_by_name("cuda");
    if (type == AV_HWDEVICE_TYPE_NONE) {
        std::cout << "Could not support cuda hw" << std::endl;
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    AVBufferRef *hw_device_ctx = NULL;
    int err = av_hwdevice_ctx_create(&hw_device_ctx, type, NULL, NULL, 0);
    if (err < 0) {
        std::cout << "Could not create hw context"<<std::endl;
		avformat_close_input(&fmt_ctx);
		return -1;
    }
	


    const char* cuvid_decoder = get_cuvid_decoder_name(codec_params->codec_id);
    const AVCodec* codec = NULL;
    if (!cuvid_decoder) {
        codec = avcodec_find_decoder(codec_params->codec_id);
    } else {
        codec = avcodec_find_decoder_by_name(cuvid_decoder);
    }
    if (!codec) {
        std::cout << "Unsupported codec!" << std::endl;
        avformat_close_input(&fmt_ctx);
        return -1;
    }


    AVCodecContext* codec_ctx = avcodec_alloc_context3(codec);
    if (avcodec_parameters_to_context(codec_ctx, codec_params) < 0) {
        std::cout << "Could not initialize codec context" << std::endl;
        avcodec_free_context(&codec_ctx);
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    //set buf to hwbuffer
    codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
	
	if (!codec_ctx->hw_device_ctx) {
        std::cout << "Failed to set hardware device context." << std::endl;
        avcodec_free_context(&codec_ctx);
        av_buffer_unref(&hw_device_ctx);
		avformat_close_input(&fmt_ctx);
        return -1;
    }
	codec_ctx->get_format = get_hw_format;


    if (avcodec_open2(codec_ctx, codec, nullptr) < 0) {
        std::cout << "Could not open codec" << std::endl;
        avcodec_free_context(&codec_ctx);
        avformat_close_input(&fmt_ctx);
		av_buffer_unref(&hw_device_ctx);
        return -1;
    }


    AVFrame* frame = av_frame_alloc();
    AVPacket* pkt = av_packet_alloc();

	
	int frame_count = 0;
    while (av_read_frame(fmt_ctx, pkt) >= 0) {
        if (pkt->stream_index == video_stream_idx) {

            int ret = avcodec_send_packet(codec_ctx, pkt);
            if (ret < 0) {
                std::cout << "Error sending packet to decoder" << std::endl;
                break;
            }


            while (ret >= 0) {
                ret = avcodec_receive_frame(codec_ctx, frame);
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
                    break;
                } else if (ret < 0) {
                    std::cout << "Error during decoding" << std::endl;
                    break;
                }
				frame_count++;
                std::cout << "Decoded frame " << frame_count << std::endl;
            }
        }
        av_packet_unref(pkt);
    }

    av_frame_free(&frame);
    avcodec_free_context(&codec_ctx);
    avformat_close_input(&fmt_ctx);

    return 0;
}

diff --git a/src/torchcodec/_core/SingleStreamDecoder.cpp b/src/torchcodec/_core/SingleStreamDecoder.cpp
index 865179a..b5788e7 100644
--- a/src/torchcodec/_core/SingleStreamDecoder.cpp
+++ b/src/torchcodec/_core/SingleStreamDecoder.cpp
@@ -27,6 +27,17 @@ int64_t secondsToClosestPts(double seconds, const AVRational& timeBase) {
       std::round(seconds * timeBase.den / timeBase.num));
 }
 
+inline char* find_codec(const char* input) {
+    const char* codecs[] = {"h264", "hevc", "av1", "vp9"};
+    size_t codec_len = sizeof(codecs) / sizeof(codecs[0]);
+    for (size_t i = 0; i < codec_len; ++i) {
+        if (strstr(input, codecs[i])) {
+            return (char*)codecs[i];
+        }
+    }
+    return NULL;
+}
+
 // Some videos aren't properly encoded and do not specify pts values for
 // packets, and thus for frames. Unset values correspond to INT64_MIN. When that
 // happens, we fallback to the dts value which hopefully exists and is correct.
@@ -425,9 +436,22 @@ void SingleStreamDecoder::addStream(
   // addStream() which is supposed to be generic
   if (mediaType == AVMEDIA_TYPE_VIDEO) {
     if (deviceInterface_) {
-      avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
-          deviceInterface_->findCodec(streamInfo.stream->codecpar->codec_id)
-              .value_or(avCodec));
+      if (device.type() != torch::kCUDA) {
+        avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
+            deviceInterface_->findCodec(streamInfo.stream->codecpar->codec_id)
+                .value_or(avCodec));
+      }
+      else {
+        const char* cuvid_suffix = "_cuvid";
+        char* codec_name = find_codec(avCodec->name);
+        size_t cuvid_length = std::strlen(codec_name) + std::strlen(cuvid_suffix) + 1;
+        char* cuvid_name = new char[cuvid_length];
+        std::strcpy(cuvid_name, codec_name);
+        std::strcat(cuvid_name, cuvid_suffix);
+        avCodec = avcodec_find_decoder_by_name(cuvid_name);
+        delete[] cuvid_name;
+        TORCH_CHECK(avCodec != nullptr);
+      }
     }
   }
 

diff --git a/src/torchcodec/_core/SingleStreamDecoder.cpp b/src/torchcodec/_core/SingleStreamDecoder.cpp
index b73f703..0222087 100644
--- a/src/torchcodec/_core/SingleStreamDecoder.cpp
+++ b/src/torchcodec/_core/SingleStreamDecoder.cpp
@@ -26,6 +26,17 @@ int64_t secondsToClosestPts(double seconds, const AVRational& timeBase) {
       std::round(seconds * timeBase.den / timeBase.num));
 }
 
+inline char* find_codec(const char* input) {
+    const char* codecs[] = {"h264", "hevc", "av1", "vp9"};
+    size_t codec_len = sizeof(codecs) / sizeof(codecs[0]);
+    for (size_t i = 0; i < codec_len; ++i) {
+        if (strstr(input, codecs[i])) {
+            return (char*)codecs[i];
+        }
+    }
+    return NULL;
+}
+
 // Some videos aren't properly encoded and do not specify pts values for
 // packets, and thus for frames. Unset values correspond to INT64_MIN. When that
 // happens, we fallback to the dts value which hopefully exists and is correct.
@@ -388,9 +399,22 @@ void SingleStreamDecoder::addStream(
   // addStream() which is supposed to be generic
   if (mediaType == AVMEDIA_TYPE_VIDEO) {
     if (deviceInterface_) {
-      avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
-          deviceInterface_->findCodec(streamInfo.stream->codecpar->codec_id)
-              .value_or(avCodec));
+      if (device.type() != torch::kCUDA) {
+          avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
+              deviceInterface_->findCodec(streamInfo.stream->codecpar->codec_id)
+                  .value_or(avCodec));
+      }
+      else {
+        const char* cuvid_suffix = "_cuvid";
+        char* codec_name = find_codec(avCodec->name);
+        size_t cuvid_length = std::strlen(codec_name) + std::strlen(cuvid_suffix) + 1;
+        char* cuvid_name = new char[cuvid_length];
+        std::strcpy(cuvid_name, codec_name);
+        std::strcat(cuvid_name, cuvid_suffix);
+        avCodec = avcodec_find_decoder_by_name(cuvid_name);
+        delete[] cuvid_name;
+        TORCH_CHECK(avCodec != nullptr);
+      }
     }
   }
 
@@ -417,6 +441,12 @@ void SingleStreamDecoder::addStream(
     throw std::invalid_argument(getFFMPEGErrorStringFromErrorCode(retVal));
   }
 
+  if (device.type() == torch::kCUDA) {
+    codecContext->hw_frames_ctx = av_hwframe_ctx_alloc(codecContext->hw_device_ctx);
+    AVHWFramesContext* hwframe_ctx = (AVHWFramesContext*)codecContext->hw_frames_ctx->data;
+    // in avcodec_open2, cuvid will not set sw_format until sequence callback in avcodec_send_packet()
+    hwframe_ctx->sw_format = codecContext->sw_pix_fmt;
+  }
   codecContext->time_base = streamInfo.stream->time_base;
   containerMetadata_.allStreamMetadata[activeStreamIndex_].codecName =
       std::string(avcodec_get_name(codecContext->codec_id));

diff --git a/src/torchcodec/decoders/_core/VideoDecoder.cpp b/src/torchcodec/decoders/_core/VideoDecoder.cpp
index 97214ce..c88c581 100644
--- a/src/torchcodec/decoders/_core/VideoDecoder.cpp
+++ b/src/torchcodec/decoders/_core/VideoDecoder.cpp
@@ -40,6 +40,17 @@ int64_t secondsToClosestPts(double seconds, const AVRational& timeBase) {
   return static_cast<int64_t>(std::round(seconds * timeBase.den));
 }
 
+inline char* find_codec(const char* input) {
+    const char* codecs[] = {"h264", "hevc", "av1", "vp9"};
+    size_t codec_len = sizeof(codecs) / sizeof(codecs[0]);
+    for (size_t i = 0; i < codec_len; ++i) {
+        if (strstr(input, codecs[i])) {
+            return (char*)codecs[i];
+        }
+    }
+    return NULL;
+}
+
 std::vector<std::string> splitStringWithDelimiters(
     const std::string& str,
     const std::string& delims) {
@@ -449,9 +460,15 @@ void VideoDecoder::addStream(
   // TODO_CODE_QUALITY it's pretty meh to have a video-specific logic within
   // addStream() which is supposed to be generic
   if (mediaType == AVMEDIA_TYPE_VIDEO && device.type() == torch::kCUDA) {
-    avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
-        findCudaCodec(device, streamInfo.stream->codecpar->codec_id)
-            .value_or(avCodec));
+    const char* cuvid_suffix = "_cuvid";
+    char* codec_name = find_codec(avCodec->name);
+    size_t cuvid_length = std::strlen(codec_name) + std::strlen(cuvid_suffix) + 1;
+    char* cuvid_name = new char[cuvid_length];
+    std::strcpy(cuvid_name, codec_name);
+    std::strcat(cuvid_name, cuvid_suffix);
+    avCodec = avcodec_find_decoder_by_name(cuvid_name);
+    delete[] cuvid_name;
+    TORCH_CHECK(avCodec != nullptr);
   }
 
   AVCodecContext* codecContext = avcodec_alloc_context3(avCodec);
@@ -474,6 +491,12 @@ void VideoDecoder::addStream(
     throw std::invalid_argument(getFFMPEGErrorStringFromErrorCode(retVal));
   }
 
+  if (device.type() == torch::kCUDA) {
+    codecContext->hw_frames_ctx = av_hwframe_ctx_alloc(codecContext->hw_device_ctx);
+    AVHWFramesContext* hwframe_ctx = (AVHWFramesContext*)codecContext->hw_frames_ctx->data;
+    // in avcodec_open2, cuvid will not set sw_format until sequence callback in avcodec_send_packet()
+    hwframe_ctx->sw_format = codecContext->sw_pix_fmt;
+  }
   codecContext->time_base = streamInfo.stream->time_base;
   containerMetadata_.allStreamMetadata[activeStreamIndex_].codecName =
       std::string(avcodec_get_name(codecContext->codec_id));

diff --git a/src/audio/audio_reader.cc b/src/audio/audio_reader.cc
index be706f1..78c62fb 100644
--- a/src/audio/audio_reader.cc
+++ b/src/audio/audio_reader.cc
@@ -128,7 +128,11 @@ namespace decord {
                 pCodecParameters = tempCodecParameters;
                 originalSampleRate = tempCodecParameters->sample_rate;
                 if (targetSampleRate == -1) targetSampleRate = originalSampleRate;
+#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(57, 28, 100)
+                numChannels = tempCodecParameters->ch_layout.nb_channels;
+#else
                 numChannels = tempCodecParameters->channels;
+#endif
                 break;
             }
         }
@@ -229,7 +233,11 @@ namespace decord {
         // allocate resample buffer
         float** outBuffer;
         int outLinesize = 0;
+#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(57, 28, 100)
+        int outNumChannels = mono ? AV_CH_LAYOUT_MONO : pFrame->ch_layout.nb_channels;
+#else
         int outNumChannels = av_get_channel_layout_nb_channels(mono ? AV_CH_LAYOUT_MONO : pFrame->channel_layout);
+#endif
         numChannels = outNumChannels;
         int outNumSamples = av_rescale_rnd(pFrame->nb_samples,
                                            this->targetSampleRate, pFrame->sample_rate, AV_ROUND_UP);
@@ -281,11 +289,17 @@ namespace decord {
         if (!this->swr) {
             LOG(FATAL) << "ERROR Failed to allocate resample context";
         }
+#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(57, 28, 100)
+        av_channel_layout_default(&pCodecContext->ch_layout, pCodecContext->ch_layout.nb_channels);
+        av_opt_set_chlayout(this->swr, "in_channel_layout",  &pCodecContext->ch_layout, 0);
+        av_opt_set_chlayout(this->swr, "out_channel_layout", &pCodecContext->ch_layout, 0);
+#else
         if (pCodecContext->channel_layout == 0) {
             pCodecContext->channel_layout = av_get_default_channel_layout( pCodecContext->channels );
         }
         av_opt_set_channel_layout(this->swr, "in_channel_layout",  pCodecContext->channel_layout, 0);
         av_opt_set_channel_layout(this->swr, "out_channel_layout", mono ? AV_CH_LAYOUT_MONO : pCodecContext->channel_layout,  0);
+#endif
         av_opt_set_int(this->swr, "in_sample_rate",     pCodecContext->sample_rate,                0);
         av_opt_set_int(this->swr, "out_sample_rate",    this->targetSampleRate,                0);
         av_opt_set_sample_fmt(this->swr, "in_sample_fmt",  pCodecContext->sample_fmt, 0);
diff --git a/src/video/ffmpeg/ffmpeg_common.h b/src/video/ffmpeg/ffmpeg_common.h
index b0b973f..f0f7316 100644
--- a/src/video/ffmpeg/ffmpeg_common.h
+++ b/src/video/ffmpeg/ffmpeg_common.h
@@ -21,6 +21,7 @@
 extern "C" {
 #endif
 #include <libavcodec/avcodec.h>
+#include <libavcodec/bsf.h>
 #include <libavformat/avformat.h>
 #include <libavformat/avio.h>
 #include <libavfilter/avfilter.h>
diff --git a/src/video/nvcodec/cuda_threaded_decoder.cc b/src/video/nvcodec/cuda_threaded_decoder.cc
index 62bc7ee..957a90d 100644
--- a/src/video/nvcodec/cuda_threaded_decoder.cc
+++ b/src/video/nvcodec/cuda_threaded_decoder.cc
@@ -17,7 +17,7 @@ namespace decord {
 namespace cuda {
 using namespace runtime;
 
-CUThreadedDecoder::CUThreadedDecoder(int device_id, AVCodecParameters *codecpar, AVInputFormat *iformat)
+CUThreadedDecoder::CUThreadedDecoder(int device_id, AVCodecParameters *codecpar, const AVInputFormat *iformat)
     : device_id_(device_id), stream_({device_id, false}), device_{}, ctx_{}, parser_{}, decoder_{},
     pkt_queue_{}, frame_queue_{},
     run_(false), frame_count_(0), draining_(false),
@@ -70,7 +70,7 @@ CUThreadedDecoder::CUThreadedDecoder(int device_id, AVCodecParameters *codecpar,
     }
 }
 
-void CUThreadedDecoder::InitBitStreamFilter(AVCodecParameters *codecpar, AVInputFormat *iformat) {
+void CUThreadedDecoder::InitBitStreamFilter(AVCodecParameters *codecpar, const AVInputFormat *iformat) {
     const char* bsf_name = nullptr;
     if (AV_CODEC_ID_H264 == codecpar->codec_id) {
         // H.264
diff --git a/src/video/nvcodec/cuda_threaded_decoder.h b/src/video/nvcodec/cuda_threaded_decoder.h
index d7e6fcd..61958a1 100644
--- a/src/video/nvcodec/cuda_threaded_decoder.h
+++ b/src/video/nvcodec/cuda_threaded_decoder.h
@@ -46,7 +46,7 @@ class CUThreadedDecoder final : public ThreadedDecoderInterface {
     using FrameOrderQueuePtr = std::unique_ptr<FrameOrderQueue>;
 
     public:
-        CUThreadedDecoder(int device_id, AVCodecParameters *codecpar, AVInputFormat *iformat);
+        CUThreadedDecoder(int device_id, AVCodecParameters *codecpar, const AVInputFormat *iformat);
         void SetCodecContext(AVCodecContext *dec_ctx, int width = -1, int height = -1, int rotation = 0);
         bool Initialized() const;
         void Start();
@@ -70,7 +70,7 @@ class CUThreadedDecoder final : public ThreadedDecoderInterface {
         void LaunchThreadImpl();
         void RecordInternalError(std::string message);
         void CheckErrorStatus();
-        void InitBitStreamFilter(AVCodecParameters *codecpar, AVInputFormat *iformat);
+        void InitBitStreamFilter(AVCodecParameters *codecpar, const AVInputFormat *iformat);

         int device_id_;
         CUStream stream_;
diff --git a/src/video/video_reader.cc b/src/video/video_reader.cc
index af4858d..99c9635 100644
--- a/src/video/video_reader.cc
+++ b/src/video/video_reader.cc
@@ -145,7 +145,7 @@ VideoReader::~VideoReader(){

 void VideoReader::SetVideoStream(int stream_nb) {
     if (!fmt_ctx_) return;
-    AVCodec *dec;
+    const AVCodec *dec;
     int st_nb = av_find_best_stream(fmt_ctx_.get(), AVMEDIA_TYPE_VIDEO, stream_nb, -1, &dec, 0);
     // LOG(INFO) << "find best stream: " << st_nb;
     CHECK_GE(st_nb, 0) << "ERROR cannot find video stream with wanted index: " << stream_nb;