Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support specifying decoder and its options (#2327)
Browse files Browse the repository at this point in the history
Summary:
This commit adds support to specify decoder to Streamer's add stream method.
This is roughly equivalent to `ffmpeg`'s `-c:v foo` and `-c:a foo` options.

This allows to override the decoder codec and/or specify the option of
the decoder.

This change allows to specify Nvidia NVDEC codec for supported formats,
which uses dedicated hardware for decoding the video.

 ---

Note: The CL might look overwhelming, but it's essentially, add new parameters in Python, and pass them down all the way to  `AVCodecContextPtr`, which initializes the actual decoder implementation (`AVCodecContext`.)

Pull Request resolved: #2327

Differential Revision: D35626904

Pulled By: mthrok

fbshipit-source-id: 94ddbf5f1c3beb01194675d2792900cf873922f5
mthrok authored and facebook-github-bot committed Apr 14, 2022
1 parent c262758 commit 39c9213
Showing 10 changed files with 140 additions and 36 deletions.
6 changes: 5 additions & 1 deletion torchaudio/csrc/ffmpeg/decoder.cpp
Original file line number Diff line number Diff line change
@@ -6,7 +6,11 @@ namespace ffmpeg {
////////////////////////////////////////////////////////////////////////////////
// Decoder
////////////////////////////////////////////////////////////////////////////////
Decoder::Decoder(AVCodecParameters* pParam) : pCodecContext(pParam) {}
Decoder::Decoder(
AVCodecParameters* pParam,
const std::string& decoder_name,
const std::map<std::string, std::string>& decoder_option)
: pCodecContext(pParam, decoder_name, decoder_option) {}

int Decoder::process_packet(AVPacket* pPacket) {
return avcodec_send_packet(pCodecContext, pPacket);
5 changes: 4 additions & 1 deletion torchaudio/csrc/ffmpeg/decoder.h
Original file line number Diff line number Diff line change
@@ -10,7 +10,10 @@ class Decoder {

public:
// Default constructable
Decoder(AVCodecParameters* pParam);
Decoder(
AVCodecParameters* pParam,
const std::string& decoder_name,
const std::map<std::string, std::string>& decoder_option);
// Custom destructor to clean up the resources
~Decoder() = default;
// Non-copyable
45 changes: 36 additions & 9 deletions torchaudio/csrc/ffmpeg/ffmpeg.cpp
Original file line number Diff line number Diff line change
@@ -151,11 +151,22 @@ void AVCodecContextDeleter::operator()(AVCodecContext* p) {
};

namespace {
AVCodecContext* get_codec_context(AVCodecParameters* pParams) {
const AVCodec* pCodec = avcodec_find_decoder(pParams->codec_id);
AVCodecContext* get_codec_context(
enum AVCodecID codec_id,
const std::string& decoder_name) {
const AVCodec* pCodec = decoder_name.empty()
? avcodec_find_decoder(codec_id)
: avcodec_find_decoder_by_name(decoder_name.c_str());

if (!pCodec) {
throw std::runtime_error("Unknown codec.");
std::stringstream ss;
if (decoder_name.empty()) {
ss << "Unsupported codec: \"" << avcodec_get_name(codec_id) << "\", ("
<< codec_id << ").";
} else {
ss << "Unsupported codec: \"" << decoder_name << "\".";
}
throw std::runtime_error(ss.str());
}

AVCodecContext* pCodecContext = avcodec_alloc_context3(pCodec);
@@ -167,27 +178,43 @@ AVCodecContext* get_codec_context(AVCodecParameters* pParams) {

void init_codec_context(
AVCodecContext* pCodecContext,
AVCodecParameters* pParams) {
const AVCodec* pCodec = avcodec_find_decoder(pParams->codec_id);
AVCodecParameters* pParams,
const std::string& decoder_name,
std::map<std::string, std::string> decoder_option) {
const AVCodec* pCodec = decoder_name.empty()
? avcodec_find_decoder(pParams->codec_id)
: avcodec_find_decoder_by_name(decoder_name.c_str());

// No need to check if pCodec is null as it's been already checked in
// get_codec_context

if (avcodec_parameters_to_context(pCodecContext, pParams) < 0) {
throw std::runtime_error("Failed to set CodecContext parameter.");
}

if (avcodec_open2(pCodecContext, pCodec, NULL) < 0) {
AVDictionary* opts = get_option_dict(decoder_option);
if (avcodec_open2(pCodecContext, pCodec, &opts) < 0) {
throw std::runtime_error("Failed to initialize CodecContext.");
}
auto unused_keys = clean_up_dict(opts);
if (unused_keys.size()) {
throw std::runtime_error(
"Unexpected decoder options: " + join(unused_keys));
}

if (pParams->codec_type == AVMEDIA_TYPE_AUDIO && !pParams->channel_layout)
pParams->channel_layout =
av_get_default_channel_layout(pCodecContext->channels);
}
} // namespace

AVCodecContextPtr::AVCodecContextPtr(AVCodecParameters* pParam)
AVCodecContextPtr::AVCodecContextPtr(
AVCodecParameters* pParam,
const std::string& decoder_name,
const std::map<std::string, std::string>& decoder_option)
: Wrapper<AVCodecContext, AVCodecContextDeleter>(
get_codec_context(pParam)) {
init_codec_context(ptr.get(), pParam);
get_codec_context(pParam->codec_id, decoder_name)) {
init_codec_context(ptr.get(), pParam, decoder_name, decoder_option);
}
////////////////////////////////////////////////////////////////////////////////
// AVFilterGraph
5 changes: 4 additions & 1 deletion torchaudio/csrc/ffmpeg/ffmpeg.h
Original file line number Diff line number Diff line change
@@ -118,7 +118,10 @@ struct AVCodecContextDeleter {
};
struct AVCodecContextPtr
: public Wrapper<AVCodecContext, AVCodecContextDeleter> {
AVCodecContextPtr(AVCodecParameters* pParam);
AVCodecContextPtr(
AVCodecParameters* pParam,
const std::string& decoder,
const std::map<std::string, std::string>& decoder_option);
};

////////////////////////////////////////////////////////////////////////////////
36 changes: 26 additions & 10 deletions torchaudio/csrc/ffmpeg/prototype.cpp
Original file line number Diff line number Diff line change
@@ -7,8 +7,10 @@ namespace ffmpeg {

namespace {

using OptionDict = c10::Dict<std::string, std::string>;

std::map<std::string, std::string> convert_dict(
const c10::optional<c10::Dict<std::string, std::string>>& option) {
const c10::optional<OptionDict>& option) {
std::map<std::string, std::string> opts;
if (option) {
for (auto& it : option.value()) {
@@ -23,7 +25,7 @@ struct StreamerHolder : torch::CustomClassHolder {
StreamerHolder(
const std::string& src,
c10::optional<std::string> device,
c10::optional<c10::Dict<std::string, std::string>> option)
c10::optional<OptionDict> option)
: s(src, device.value_or(""), convert_dict(option)) {}
};

@@ -32,7 +34,7 @@ using S = c10::intrusive_ptr<StreamerHolder>;
S init(
const std::string& src,
c10::optional<std::string> device,
c10::optional<c10::Dict<std::string, std::string>> option) {
c10::optional<OptionDict> option) {
return c10::make_intrusive<StreamerHolder>(src, device, option);
}

@@ -216,7 +218,7 @@ void add_basic_audio_stream(
const c10::optional<int64_t>& sample_rate,
const c10::optional<c10::ScalarType>& dtype) {
std::string filter_desc = get_afilter_desc(sample_rate, dtype);
s->s.add_audio_stream(i, frames_per_chunk, num_chunks, filter_desc);
s->s.add_audio_stream(i, frames_per_chunk, num_chunks, filter_desc, "", {});
}

void add_basic_video_stream(
@@ -229,27 +231,41 @@ void add_basic_video_stream(
const c10::optional<int64_t>& height,
const c10::optional<std::string>& format) {
std::string filter_desc = get_vfilter_desc(frame_rate, width, height, format);
s->s.add_video_stream(i, frames_per_chunk, num_chunks, filter_desc);
s->s.add_video_stream(i, frames_per_chunk, num_chunks, filter_desc, "", {});
}

void add_audio_stream(
S s,
int64_t i,
int64_t frames_per_chunk,
int64_t num_chunks,
const c10::optional<std::string>& filter_desc) {
const c10::optional<std::string>& filter_desc,
const c10::optional<std::string>& decoder,
const c10::optional<OptionDict>& decoder_options) {
s->s.add_audio_stream(
i, frames_per_chunk, num_chunks, filter_desc.value_or(""));
i,
frames_per_chunk,
num_chunks,
filter_desc.value_or(""),
decoder.value_or(""),
convert_dict(decoder_options));
}

void add_video_stream(
S s,
int64_t i,
int64_t frames_per_chunk,
int64_t num_chunks,
const c10::optional<std::string>& filter_desc) {
const c10::optional<std::string>& filter_desc,
const c10::optional<std::string>& decoder,
const c10::optional<OptionDict>& decoder_options) {
s->s.add_video_stream(
i, frames_per_chunk, num_chunks, filter_desc.value_or(""));
i,
frames_per_chunk,
num_chunks,
filter_desc.value_or(""),
decoder.value_or(""),
convert_dict(decoder_options));
}

void remove_stream(S s, int64_t i) {
@@ -293,7 +309,7 @@ std::tuple<c10::optional<torch::Tensor>, int64_t> load(const std::string& src) {
int i = s.find_best_audio_stream();
auto sinfo = s.get_src_stream_info(i);
int64_t sample_rate = static_cast<int64_t>(sinfo.sample_rate);
s.add_audio_stream(i, -1, -1, "");
s.add_audio_stream(i, -1, -1, "", "", {});
process_all_packets(s);
auto tensors = s.pop_chunks();
return std::make_tuple<>(tensors[0], sample_rate);
7 changes: 5 additions & 2 deletions torchaudio/csrc/ffmpeg/stream_processor.cpp
Original file line number Diff line number Diff line change
@@ -6,8 +6,11 @@ namespace ffmpeg {

using KeyType = StreamProcessor::KeyType;

StreamProcessor::StreamProcessor(AVCodecParameters* codecpar)
: decoder(codecpar) {}
StreamProcessor::StreamProcessor(
AVCodecParameters* codecpar,
const std::string& decoder_name,
const std::map<std::string, std::string>& decoder_option)
: decoder(codecpar, decoder_name, decoder_option) {}

////////////////////////////////////////////////////////////////////////////////
// Configurations
5 changes: 4 additions & 1 deletion torchaudio/csrc/ffmpeg/stream_processor.h
Original file line number Diff line number Diff line change
@@ -25,7 +25,10 @@ class StreamProcessor {
std::map<KeyType, Sink> sinks;

public:
StreamProcessor(AVCodecParameters* codecpar);
StreamProcessor(
AVCodecParameters* codecpar,
const std::string& decoder_name,
const std::map<std::string, std::string>& decoder_option);
~StreamProcessor() = default;
// Non-copyable
StreamProcessor(const StreamProcessor&) = delete;
23 changes: 17 additions & 6 deletions torchaudio/csrc/ffmpeg/streamer.cpp
Original file line number Diff line number Diff line change
@@ -156,39 +156,50 @@ void Streamer::add_audio_stream(
int i,
int frames_per_chunk,
int num_chunks,
std::string filter_desc) {
std::string filter_desc,
const std::string& decoder,
const std::map<std::string, std::string>& decoder_option) {
add_stream(
i,
AVMEDIA_TYPE_AUDIO,
frames_per_chunk,
num_chunks,
std::move(filter_desc));
std::move(filter_desc),
decoder,
decoder_option);
}

void Streamer::add_video_stream(
int i,
int frames_per_chunk,
int num_chunks,
std::string filter_desc) {
std::string filter_desc,
const std::string& decoder,
const std::map<std::string, std::string>& decoder_option) {
add_stream(
i,
AVMEDIA_TYPE_VIDEO,
frames_per_chunk,
num_chunks,
std::move(filter_desc));
std::move(filter_desc),
decoder,
decoder_option);
}

void Streamer::add_stream(
int i,
AVMediaType media_type,
int frames_per_chunk,
int num_chunks,
std::string filter_desc) {
std::string filter_desc,
const std::string& decoder,
const std::map<std::string, std::string>& decoder_option) {
validate_src_stream_type(i, media_type);
AVStream* stream = pFormatContext->streams[i];
stream->discard = AVDISCARD_DEFAULT;
if (!processors[i])
processors[i] = std::make_unique<StreamProcessor>(stream->codecpar);
processors[i] = std::make_unique<StreamProcessor>(
stream->codecpar, decoder, decoder_option);
int key = processors[i]->add_stream(
stream->time_base,
stream->codecpar,
12 changes: 9 additions & 3 deletions torchaudio/csrc/ffmpeg/streamer.h
Original file line number Diff line number Diff line change
@@ -66,12 +66,16 @@ class Streamer {
int i,
int frames_per_chunk,
int num_chunks,
std::string filter_desc);
std::string filter_desc,
const std::string& decoder,
const std::map<std::string, std::string>& decoder_option);
void add_video_stream(
int i,
int frames_per_chunk,
int num_chunks,
std::string filter_desc);
std::string filter_desc,
const std::string& decoder,
const std::map<std::string, std::string>& decoder_option);
void remove_stream(int i);

private:
@@ -80,7 +84,9 @@ class Streamer {
AVMediaType media_type,
int frames_per_chunk,
int num_chunks,
std::string filter_desc);
std::string filter_desc,
const std::string& decoder,
const std::map<std::string, std::string>& decoder_option);

public:
//////////////////////////////////////////////////////////////////////////////
32 changes: 30 additions & 2 deletions torchaudio/prototype/io/streamer.py
Original file line number Diff line number Diff line change
@@ -354,6 +354,8 @@ def add_audio_stream(
buffer_chunk_size: int = 3,
stream_index: Optional[int] = None,
filter_desc: Optional[str] = None,
decoder: Optional[str] = None,
decoder_options: Optional[Dict[str, str]] = None,
):
"""Add output audio stream
@@ -374,10 +376,22 @@ def add_audio_stream(
The list of available filters can be found at
https://ffmpeg.org/ffmpeg-filters.html
Note that complex filters are not supported.
decoder (str or None, optional): The name of the decoder to be used.
When provided, use the specified decoder instead of the default one.
decoder_options (dict or None, optional): Options passed to decoder.
Mapping from str to str.
"""
i = self.default_audio_stream if stream_index is None else stream_index
torch.ops.torchaudio.ffmpeg_streamer_add_audio_stream(
self._s, i, frames_per_chunk, buffer_chunk_size, filter_desc
self._s,
i,
frames_per_chunk,
buffer_chunk_size,
filter_desc,
decoder,
decoder_options,
)

def add_video_stream(
@@ -386,6 +400,8 @@ def add_video_stream(
buffer_chunk_size: int = 3,
stream_index: Optional[int] = None,
filter_desc: Optional[str] = None,
decoder: Optional[str] = None,
decoder_options: Optional[Dict[str, str]] = None,
):
"""Add output video stream
@@ -406,10 +422,22 @@ def add_video_stream(
The list of available filters can be found at
https://ffmpeg.org/ffmpeg-filters.html
Note that complex filters are not supported.
decoder (str or None, optional): The name of the decoder to be used.
When provided, use the specified decoder instead of the default one.
decoder_options (dict or None, optional): Options passed to decoder.
Mapping from str to str.
"""
i = self.default_video_stream if stream_index is None else stream_index
torch.ops.torchaudio.ffmpeg_streamer_add_video_stream(
self._s, i, frames_per_chunk, buffer_chunk_size, filter_desc
self._s,
i,
frames_per_chunk,
buffer_chunk_size,
filter_desc,
decoder,
decoder_options,
)

def remove_stream(self, i: int):

0 comments on commit 39c9213

Please sign in to comment.