Skip to content

GitLab

Commit 014d7140 authored by moto's avatar moto Committed by Facebook GitHub Bot
Browse files

Refactor Tensor conversion in StreamReader (#3170) 

Summary:
Currently, when the Buffer converts AVFrame* to torch::Tensor,
it checks the format at each time a frame is passed, and
perform the conversion.

This commit changes it so that the conversion operation is
pre-instantiated at the time outside stream is configured.

It introduces Converter implementations for various formats,
and use template to embed them in Buffer class.
This way, branching like if/switch are eliminated from
decoding path.

Pull Request resolved: https://github.com/pytorch/audio/pull/3170

Reviewed By: xiaohui-zhang

Differential Revision: D44048293

Pulled By: mthrok

fbshipit-source-id: 30d8b240a5695d7513f499ce17853f2f0ffcab9f
parent 92f2ea89
Show whitespace changes
Inline Side-by-side
Showing with 1092 additions and 642 deletions
test/torchaudio_unittest/io/stream_reader_test.py
View file @ 014d7140
...@@ -832,11 +832,49 @@ class StreamReaderAudioTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase) ...@@ -832,11 +832,49 @@ class StreamReaderAudioTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase)
if self.test_type == "fileobj": if self.test_type == "fileobj":
src.seek(0) src.seek(0)
self._test_wav(src, original, fmt=None) self._test_wav(src, original, fmt=None)
# convert to float32
expected = _to_fltp(original) def test_audio_stream_format(self):
if self.test_type == "fileobj": num_channels = 2
src.seek(0) src, s32 = self.get_src(8000, dtype="int32", num_channels=num_channels)
self._test_wav(src, expected, fmt="fltp") args = {
"num_channels": num_channels,
"normalize": False,
"channels_first": False,
"num_frames": 1 << 16,
}
u8 = get_wav_data("uint8", **args)
s16 = get_wav_data("int16", **args)
s64 = s32.to(torch.int64) * (1 << 32)
f32 = get_wav_data("float32", **args)
f64 = get_wav_data("float64", **args)
s = StreamReader(src)
s.add_basic_audio_stream(frames_per_chunk=-1, format="u8")
s.add_basic_audio_stream(frames_per_chunk=-1, format="u8p")
s.add_basic_audio_stream(frames_per_chunk=-1, format="s16")
s.add_basic_audio_stream(frames_per_chunk=-1, format="s16p")
s.add_basic_audio_stream(frames_per_chunk=-1, format="s32")
s.add_basic_audio_stream(frames_per_chunk=-1, format="s32p")
s.add_basic_audio_stream(frames_per_chunk=-1, format="s64")
s.add_basic_audio_stream(frames_per_chunk=-1, format="s64p")
s.add_basic_audio_stream(frames_per_chunk=-1, format="flt")
s.add_basic_audio_stream(frames_per_chunk=-1, format="fltp")
s.add_basic_audio_stream(frames_per_chunk=-1, format="dbl")
s.add_basic_audio_stream(frames_per_chunk=-1, format="dblp")
s.process_all_packets()
chunks = s.pop_chunks()
self.assertEqual(chunks[0], u8, atol=1, rtol=0)
self.assertEqual(chunks[1], u8, atol=1, rtol=0)
self.assertEqual(chunks[2], s16)
self.assertEqual(chunks[3], s16)
self.assertEqual(chunks[4], s32)
self.assertEqual(chunks[5], s32)
self.assertEqual(chunks[6], s64)
self.assertEqual(chunks[7], s64)
self.assertEqual(chunks[8], f32)
self.assertEqual(chunks[9], f32)
self.assertEqual(chunks[10], f64)
self.assertEqual(chunks[11], f64)
@nested_params( @nested_params(
["int16", "uint8", "int32"], # "float", "double", "int64"] ["int16", "uint8", "int32"], # "float", "double", "int64"]
...@@ -969,6 +1007,7 @@ class StreamReaderImageTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase) ...@@ -969,6 +1007,7 @@ class StreamReaderImageTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase)
rgb = torch.empty(1, 3, 256, 256, dtype=torch.uint8) rgb = torch.empty(1, 3, 256, 256, dtype=torch.uint8)
rgb[0, 0] = torch.arange(256, dtype=torch.uint8).reshape([1, -1]) rgb[0, 0] = torch.arange(256, dtype=torch.uint8).reshape([1, -1])
rgb[0, 1] = torch.arange(256, dtype=torch.uint8).reshape([-1, 1]) rgb[0, 1] = torch.arange(256, dtype=torch.uint8).reshape([-1, 1])
alpha = torch.full((1, 1, 256, 256), 255, dtype=torch.uint8)
for i in range(256): for i in range(256):
rgb[0, 2] = i rgb[0, 2] = i
path = self.get_temp_path(f"ref_{i}.png") path = self.get_temp_path(f"ref_{i}.png")
...@@ -979,6 +1018,10 @@ class StreamReaderImageTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase) ...@@ -979,6 +1018,10 @@ class StreamReaderImageTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase)
yuv = rgb_to_yuv_ccir(rgb) yuv = rgb_to_yuv_ccir(rgb)
bgr = rgb[:, [2, 1, 0], :, :] bgr = rgb[:, [2, 1, 0], :, :]
gray = rgb_to_gray(rgb) gray = rgb_to_gray(rgb)
argb = torch.cat([alpha, rgb], dim=1)
rgba = torch.cat([rgb, alpha], dim=1)
abgr = torch.cat([alpha, bgr], dim=1)
bgra = torch.cat([bgr, alpha], dim=1)
s = StreamReader(path) s = StreamReader(path)
s.add_basic_video_stream(frames_per_chunk=-1, format="yuv444p") s.add_basic_video_stream(frames_per_chunk=-1, format="yuv444p")
...@@ -988,12 +1031,20 @@ class StreamReaderImageTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase) ...@@ -988,12 +1031,20 @@ class StreamReaderImageTest(_MediaSourceMixin, TempDirMixin, TorchaudioTestCase)
s.add_basic_video_stream(frames_per_chunk=-1, format="bgr24") s.add_basic_video_stream(frames_per_chunk=-1, format="bgr24")
s.add_basic_video_stream(frames_per_chunk=-1, format="gray8") s.add_basic_video_stream(frames_per_chunk=-1, format="gray8")
s.add_basic_video_stream(frames_per_chunk=-1, format="rgb48le") s.add_basic_video_stream(frames_per_chunk=-1, format="rgb48le")
s.add_basic_video_stream(frames_per_chunk=-1, format="argb")
s.add_basic_video_stream(frames_per_chunk=-1, format="rgba")
s.add_basic_video_stream(frames_per_chunk=-1, format="abgr")
s.add_basic_video_stream(frames_per_chunk=-1, format="bgra")
s.process_all_packets() s.process_all_packets()
yuv444, yuv420, nv12, rgb24, bgr24, gray8, rgb48le = s.pop_chunks() chunks = s.pop_chunks()
self.assertEqual(yuv, yuv444, atol=1, rtol=0) self.assertEqual(chunks[0], yuv, atol=1, rtol=0)
self.assertEqual(yuv, yuv420, atol=1, rtol=0) self.assertEqual(chunks[1], yuv, atol=1, rtol=0)
self.assertEqual(yuv, nv12, atol=1, rtol=0) self.assertEqual(chunks[2], yuv, atol=1, rtol=0)
self.assertEqual(rgb, rgb24, atol=0, rtol=0) self.assertEqual(chunks[3], rgb, atol=0, rtol=0)
self.assertEqual(bgr, bgr24, atol=0, rtol=0) self.assertEqual(chunks[4], bgr, atol=0, rtol=0)
self.assertEqual(gray, gray8, atol=1, rtol=0) self.assertEqual(chunks[5], gray, atol=1, rtol=0)
self.assertEqual(rgb16, rgb48le, atol=256, rtol=0) self.assertEqual(chunks[6], rgb16, atol=256, rtol=0)
self.assertEqual(chunks[7], argb, atol=0, rtol=0)
self.assertEqual(chunks[8], rgba, atol=0, rtol=0)
self.assertEqual(chunks[9], abgr, atol=0, rtol=0)
self.assertEqual(chunks[10], bgra, atol=0, rtol=0)
torchaudio/csrc/ffmpeg/CMakeLists.txt
View file @ 014d7140
...@@ -9,9 +9,9 @@ set( ...@@ -9,9 +9,9 @@ set(
sources sources
ffmpeg.cpp ffmpeg.cpp
filter_graph.cpp filter_graph.cpp
stream_reader/buffer/common.cpp
stream_reader/buffer/chunked_buffer.cpp stream_reader/buffer/chunked_buffer.cpp
stream_reader/buffer/unchunked_buffer.cpp stream_reader/buffer/unchunked_buffer.cpp
stream_reader/conversion.cpp
stream_reader/sink.cpp stream_reader/sink.cpp
stream_reader/stream_processor.cpp stream_reader/stream_processor.cpp
stream_reader/stream_reader.cpp stream_reader/stream_reader.cpp
......
torchaudio/csrc/ffmpeg/stream_reader/buffer/chunked_buffer.cpp
View file @ 014d7140
#include <torchaudio/csrc/ffmpeg/stream_reader/buffer/chunked_buffer.h> #include <torchaudio/csrc/ffmpeg/stream_reader/buffer/chunked_buffer.h>
#include <torchaudio/csrc/ffmpeg/stream_reader/buffer/common.h> #include <torchaudio/csrc/ffmpeg/stream_reader/conversion.h>
namespace torchaudio { namespace torchaudio::io::detail {
namespace io {
namespace detail {
ChunkedBuffer::ChunkedBuffer( template <typename Converter>
int frames_per_chunk, ChunkedBuffer<Converter>::ChunkedBuffer(
int num_chunks, int frames_per_chunk_,
double frame_duration) int num_chunks_,
: frame_duration(frame_duration), double frame_duration_,
frames_per_chunk(frames_per_chunk), Converter&& converter_)
num_chunks(num_chunks) {} : frame_duration(frame_duration_),
frames_per_chunk(frames_per_chunk_),
num_chunks(num_chunks_),
converter(std::move(converter_)){};
ChunkedAudioBuffer::ChunkedAudioBuffer( template <typename Converter>
int frames_per_chunk, bool ChunkedBuffer<Converter>::is_ready() const {
int num_chunks,
double frame_duration)
: ChunkedBuffer(frames_per_chunk, num_chunks, frame_duration) {}
ChunkedVideoBuffer::ChunkedVideoBuffer(
int frames_per_chunk,
int num_chunks,
double frame_duration,
const torch::Device& device_)
: ChunkedBuffer(frames_per_chunk, num_chunks, frame_duration),
device(device_) {}
bool ChunkedBuffer::is_ready() const {
return num_buffered_frames >= frames_per_chunk; return num_buffered_frames >= frames_per_chunk;
} }
void ChunkedBuffer::push_tensor(torch::Tensor frame, double pts_) { template <typename Converter>
void ChunkedBuffer<Converter>::push_frame(AVFrame* frame_, double pts_) {
torch::Tensor frame = converter.convert(frame_);
using namespace torch::indexing; using namespace torch::indexing;
// Note: // Note:
// Audio tensors contain multiple frames while video tensors contain only // Audio tensors contain multiple frames while video tensors contain only
...@@ -122,7 +113,8 @@ void ChunkedBuffer::push_tensor(torch::Tensor frame, double pts_) { ...@@ -122,7 +113,8 @@ void ChunkedBuffer::push_tensor(torch::Tensor frame, double pts_) {
} }
} }
c10::optional<Chunk> ChunkedBuffer::pop_chunk() { template <typename Converter>
c10::optional<Chunk> ChunkedBuffer<Converter>::pop_chunk() {
using namespace torch::indexing; using namespace torch::indexing;
if (!num_buffered_frames) { if (!num_buffered_frames) {
return {}; return {};
...@@ -138,19 +130,171 @@ c10::optional<Chunk> ChunkedBuffer::pop_chunk() { ...@@ -138,19 +130,171 @@ c10::optional<Chunk> ChunkedBuffer::pop_chunk() {
return {Chunk{chunk, pts_val}}; return {Chunk{chunk, pts_val}};
} }
void ChunkedAudioBuffer::push_frame(AVFrame* frame, double pts_) { template <typename Converter>
push_tensor(convert_audio(frame), pts_); void ChunkedBuffer<Converter>::flush() {
num_buffered_frames = 0;
chunks.clear();
} }
void ChunkedVideoBuffer::push_frame(AVFrame* frame, double pts_) { std::unique_ptr<Buffer> get_chunked_buffer(
push_tensor(convert_image(frame, device), pts_); int frames_per_chunk,
int num_chunks,
double frame_duration,
AVSampleFormat fmt,
int channels) {
switch (fmt) {
case AV_SAMPLE_FMT_U8: {
using Converter = AudioConverter<torch::kUInt8, false>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_S16: {
using Converter = AudioConverter<torch::kInt16, false>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_S32: {
using Converter = AudioConverter<torch::kInt32, false>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_S64: {
using Converter = AudioConverter<torch::kInt64, false>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_FLT: {
using Converter = AudioConverter<torch::kFloat32, false>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_DBL: {
using Converter = AudioConverter<torch::kFloat64, false>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_U8P: {
using Converter = AudioConverter<torch::kUInt8, true>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_S16P: {
using Converter = AudioConverter<torch::kInt16, true>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_S32P: {
using Converter = AudioConverter<torch::kInt32, true>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_S64P: {
using Converter = AudioConverter<torch::kInt64, true>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_FLTP: {
using Converter = AudioConverter<torch::kFloat32, true>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
case AV_SAMPLE_FMT_DBLP: {
using Converter = AudioConverter<torch::kFloat64, true>;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{channels});
}
default:
TORCH_INTERNAL_ASSERT(
false, "Unexpected audio type:", av_get_sample_fmt_name(fmt));
}
} }
void ChunkedBuffer::flush() { std::unique_ptr<Buffer> get_chunked_buffer(
num_buffered_frames = 0; int frames_per_chunk,
chunks.clear(); int num_chunks,
double frame_duration,
AVPixelFormat fmt,
int h,
int w,
const torch::Device& device) {
if (device.type() == at::DeviceType::CUDA) {
#ifndef USE_CUDA
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
false,
"USE_CUDA is not defined, and it should be guarded before here.");
#else
switch (fmt) {
case AV_PIX_FMT_NV12: {
using Conv = NV12CudaConverter;
return std::make_unique<ChunkedBuffer<Conv>>(
frames_per_chunk, num_chunks, frame_duration, Conv{h, w, device});
}
case AV_PIX_FMT_P010: {
using Conv = P010CudaConverter;
return std::make_unique<ChunkedBuffer<Conv>>(
frames_per_chunk, num_chunks, frame_duration, Conv{h, w, device});
}
case AV_PIX_FMT_P016: {
TORCH_CHECK(
false,
"Unsupported video format found in CUDA HW: ",
av_get_pix_fmt_name(fmt));
}
default: {
TORCH_CHECK(
false,
"Unexpected video format found in CUDA HW: ",
av_get_pix_fmt_name(fmt));
}
}
#endif
}
switch (fmt) {
case AV_PIX_FMT_RGB24:
case AV_PIX_FMT_BGR24: {
using Converter = InterlacedImageConverter;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{h, w, 3});
}
case AV_PIX_FMT_ARGB:
case AV_PIX_FMT_RGBA:
case AV_PIX_FMT_ABGR:
case AV_PIX_FMT_BGRA: {
using Converter = InterlacedImageConverter;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{h, w, 4});
}
case AV_PIX_FMT_GRAY8: {
using Converter = InterlacedImageConverter;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{h, w, 1});
}
case AV_PIX_FMT_RGB48LE: {
using Converter = Interlaced16BitImageConverter;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{h, w, 3});
}
case AV_PIX_FMT_YUV444P: {
using Converter = PlanarImageConverter;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{h, w, 3});
}
case AV_PIX_FMT_YUV420P: {
using Converter = YUV420PConverter;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{h, w});
}
case AV_PIX_FMT_NV12: {
using Converter = NV12Converter;
return std::make_unique<ChunkedBuffer<Converter>>(
frames_per_chunk, num_chunks, frame_duration, Converter{h, w});
}
default: {
TORCH_INTERNAL_ASSERT(
false, "Unexpected video format found: ", av_get_pix_fmt_name(fmt));
}
}
} }
} // namespace detail } // namespace torchaudio::io::detail
} // namespace io
} // namespace torchaudio
torchaudio/csrc/ffmpeg/stream_reader/buffer/chunked_buffer.h
View file @ 014d7140
...@@ -2,14 +2,13 @@ ...@@ -2,14 +2,13 @@
#include <torchaudio/csrc/ffmpeg/ffmpeg.h> #include <torchaudio/csrc/ffmpeg/ffmpeg.h>
#include <torchaudio/csrc/ffmpeg/stream_reader/buffer.h> #include <torchaudio/csrc/ffmpeg/stream_reader/buffer.h>
namespace torchaudio { namespace torchaudio::io::detail {
namespace io {
namespace detail {
////////////////////////////////////////////////////////////////////////////// //////////////////////////////////////////////////////////////////////////////
// Chunked Buffer Implementation // Chunked Buffer Implementation
////////////////////////////////////////////////////////////////////////////// //////////////////////////////////////////////////////////////////////////////
// Common to both audio and video // Common to both audio and video
template <typename Converter>
class ChunkedBuffer : public Buffer { class ChunkedBuffer : public Buffer {
// Each AVFrame is converted to a Tensor and stored here. // Each AVFrame is converted to a Tensor and stored here.
std::deque<torch::Tensor> chunks; std::deque<torch::Tensor> chunks;
...@@ -28,40 +27,35 @@ class ChunkedBuffer : public Buffer { ...@@ -28,40 +27,35 @@ class ChunkedBuffer : public Buffer {
// one Tensor contains multiple samples, so we track here. // one Tensor contains multiple samples, so we track here.
int64_t num_buffered_frames = 0; int64_t num_buffered_frames = 0;
protected: Converter converter;
ChunkedBuffer(int frames_per_chunk, int num_chunks, double frame_duration);
void push_tensor(torch::Tensor frame, double pts);
public: public:
ChunkedBuffer(
int frames_per_chunk,
int num_chunks,
double frame_duration,
Converter&& converter);
bool is_ready() const override; bool is_ready() const override;
void flush() override; void flush() override;
c10::optional<Chunk> pop_chunk() override; c10::optional<Chunk> pop_chunk() override;
void push_frame(AVFrame* frame_, double pts_) override;
}; };
class ChunkedAudioBuffer : public ChunkedBuffer { std::unique_ptr<Buffer> get_chunked_buffer(
public:
ChunkedAudioBuffer(
int frames_per_chunk, int frames_per_chunk,
int num_chunks, int num_chunks,
double frame_duration); double frame_duration,
AVSampleFormat fmt,
void push_frame(AVFrame* frame, double pts) override; int num_channels);
};
class ChunkedVideoBuffer : public ChunkedBuffer {
const torch::Device device;
public: std::unique_ptr<Buffer> get_chunked_buffer(
ChunkedVideoBuffer(
int frames_per_chunk, int frames_per_chunk,
int num_chunks, int num_chunks,
double frame_duration, double frame_duration,
AVPixelFormat fmt,
int height,
int width,
const torch::Device& device); const torch::Device& device);
void push_frame(AVFrame* frame, double pts) override; } // namespace torchaudio::io::detail
};
} // namespace detail
} // namespace io
} // namespace torchaudio
torchaudio/csrc/ffmpeg/stream_reader/buffer/common.cpp deleted 100644 → 0
View file @ 92f2ea89
#include <torchaudio/csrc/ffmpeg/stream_reader/buffer/common.h>
#include <stdexcept>
#include <vector>
#ifdef USE_CUDA
#include <c10/cuda/CUDAStream.h>
#endif
namespace torchaudio {
namespace io {
namespace detail {
torch::Tensor convert_audio(AVFrame* pFrame) {
// ref: https://ffmpeg.org/doxygen/4.1/filter__audio_8c_source.html#l00215
AVSampleFormat format = static_cast<AVSampleFormat>(pFrame->format);
int num_channels = pFrame->channels;
int bps = av_get_bytes_per_sample(format);
// Note
// FFMpeg's `nb_samples` represnts the number of samples par channel.
// This corresponds to `num_frames` in torchaudio's notation.
// Also torchaudio uses `num_samples` as the number of samples
// across channels.
int num_frames = pFrame->nb_samples;
int is_planar = av_sample_fmt_is_planar(format);
int num_planes = is_planar ? num_channels : 1;
int plane_size = bps * num_frames * (is_planar ? 1 : num_channels);
std::vector<int64_t> shape = is_planar
? std::vector<int64_t>{num_channels, num_frames}
: std::vector<int64_t>{num_frames, num_channels};
torch::Tensor t;
uint8_t* ptr = nullptr;
switch (format) {
case AV_SAMPLE_FMT_U8:
case AV_SAMPLE_FMT_U8P: {
t = torch::empty(shape, torch::kUInt8);
ptr = t.data_ptr<uint8_t>();
break;
}
case AV_SAMPLE_FMT_S16:
case AV_SAMPLE_FMT_S16P: {
t = torch::empty(shape, torch::kInt16);
ptr = reinterpret_cast<uint8_t*>(t.data_ptr<int16_t>());
break;
}
case AV_SAMPLE_FMT_S32:
case AV_SAMPLE_FMT_S32P: {
t = torch::empty(shape, torch::kInt32);
ptr = reinterpret_cast<uint8_t*>(t.data_ptr<int32_t>());
break;
}
case AV_SAMPLE_FMT_S64:
case AV_SAMPLE_FMT_S64P: {
t = torch::empty(shape, torch::kInt64);
ptr = reinterpret_cast<uint8_t*>(t.data_ptr<int64_t>());
break;
}
case AV_SAMPLE_FMT_FLT:
case AV_SAMPLE_FMT_FLTP: {
t = torch::empty(shape, torch::kFloat32);
ptr = reinterpret_cast<uint8_t*>(t.data_ptr<float>());
break;
}
case AV_SAMPLE_FMT_DBL:
case AV_SAMPLE_FMT_DBLP: {
t = torch::empty(shape, torch::kFloat64);
ptr = reinterpret_cast<uint8_t*>(t.data_ptr<double>());
break;
}
default:
TORCH_CHECK(
false,
"Unsupported audio format: " +
std::string(av_get_sample_fmt_name(format)));
}
for (int i = 0; i < num_planes; ++i) {
memcpy(ptr, pFrame->extended_data[i], plane_size);
ptr += plane_size;
}
if (is_planar) {
t = t.t();
}
return t;
}
namespace {
torch::Tensor get_buffer(
at::IntArrayRef shape,
const torch::Device& device = torch::Device(torch::kCPU),
const torch::Dtype dtype = torch::kUInt8) {
auto options = torch::TensorOptions()
.dtype(dtype)
.layout(torch::kStrided)
.device(device.type(), device.index());
return torch::empty(shape, options);
}
std::tuple<torch::Tensor, bool> get_image_buffer(
AVFrame* frame,
int num_frames,
const torch::Device& device) {
auto fmt = static_cast<AVPixelFormat>(frame->format);
const AVPixFmtDescriptor* desc = [&]() {
if (fmt == AV_PIX_FMT_CUDA) {
AVHWFramesContext* hwctx = (AVHWFramesContext*)frame->hw_frames_ctx->data;
return av_pix_fmt_desc_get(hwctx->sw_format);
}
return av_pix_fmt_desc_get(fmt);
}();
int channels = desc->nb_components;
// Note
// AVPixFmtDescriptor::nb_components represents the number of
// color components. This is different from the number of planes.
//
// For example, YUV420P has three color components Y, U and V, but
// U and V are squashed into the same plane, so there are only
// two planes.
//
// In our application, we cannot express the bare YUV420P as a
// single tensor, so we convert it to 3 channel tensor.
// For this reason, we use nb_components for the number of channels,
// instead of the number of planes.
//
// The actual number of planes can be retrieved with
// av_pix_fmt_count_planes.
int height = frame->height;
int width = frame->width;
int depth = desc->comp[0].depth;
auto dtype = (depth > 8) ? torch::kInt16 : torch::kUInt8;
if (desc->flags & AV_PIX_FMT_FLAG_PLANAR) {
auto buffer =
get_buffer({num_frames, channels, height, width}, device, dtype);
return std::make_tuple(buffer, true);
}
auto buffer =
get_buffer({num_frames, height, width, channels}, device, dtype);
return std::make_tuple(buffer, false);
}
void write_interlaced_image(AVFrame* pFrame, torch::Tensor& frame) {
auto ptr = frame.data_ptr<uint8_t>();
uint8_t* buf = pFrame->data[0];
size_t height = frame.size(1);
size_t stride = frame.size(2) * frame.size(3);
for (int i = 0; i < height; ++i) {
memcpy(ptr, buf, stride);
buf += pFrame->linesize[0];
ptr += stride;
}
}
void write_interlaced_image16(AVFrame* pFrame, torch::Tensor& frame) {
auto ptr = frame.data_ptr<int16_t>();
uint8_t* buf = pFrame->data[0];
size_t height = frame.size(1);
size_t stride = frame.size(2) * frame.size(3);
for (int i = 0; i < height; ++i) {
memcpy(ptr, buf, stride * 2);
buf += pFrame->linesize[0];
ptr += stride;
}
// correct for int16
frame += 32768;
}
void write_planar_image(AVFrame* pFrame, torch::Tensor& frame) {
int num_planes = static_cast<int>(frame.size(1));
int height = static_cast<int>(frame.size(2));
int width = static_cast<int>(frame.size(3));
for (int i = 0; i < num_planes; ++i) {
torch::Tensor plane = frame.index({0, i});
uint8_t* tgt = plane.data_ptr<uint8_t>();
uint8_t* src = pFrame->data[i];
int linesize = pFrame->linesize[i];
for (int h = 0; h < height; ++h) {
memcpy(tgt, src, width);
tgt += width;
src += linesize;
}
}
}
void write_yuv420p(AVFrame* pFrame, torch::Tensor& yuv) {
int height = static_cast<int>(yuv.size(2));
int width = static_cast<int>(yuv.size(3));
// Write Y plane directly
{
uint8_t* tgt = yuv.data_ptr<uint8_t>();
uint8_t* src = pFrame->data[0];
int linesize = pFrame->linesize[0];
for (int h = 0; h < height; ++h) {
memcpy(tgt, src, width);
tgt += width;
src += linesize;
}
}
// Prepare intermediate UV plane
torch::Tensor uv = get_buffer({1, 2, height / 2, width / 2});
{
uint8_t* tgt = uv.data_ptr<uint8_t>();
uint8_t* src = pFrame->data[1];
int linesize = pFrame->linesize[1];
for (int h = 0; h < height / 2; ++h) {
memcpy(tgt, src, width / 2);
tgt += width / 2;
src += linesize;
}
src = pFrame->data[2];
linesize = pFrame->linesize[2];
for (int h = 0; h < height / 2; ++h) {
memcpy(tgt, src, width / 2);
tgt += width / 2;
src += linesize;
}
}
// Upsample width and height
namespace F = torch::nn::functional;
using namespace torch::indexing;
uv = F::interpolate(
uv,
F::InterpolateFuncOptions()
.mode(torch::kNearest)
.size(std::vector<int64_t>({height, width})));
// Write to the UV plane
// yuv[:, 1:] = uv
yuv.index_put_({Slice(), Slice(1)}, uv);
}
void write_nv12_cpu(AVFrame* pFrame, torch::Tensor& yuv) {
int height = static_cast<int>(yuv.size(2));
int width = static_cast<int>(yuv.size(3));
// Write Y plane directly
{
uint8_t* tgt = yuv.data_ptr<uint8_t>();
uint8_t* src = pFrame->data[0];
int linesize = pFrame->linesize[0];
for (int h = 0; h < height; ++h) {
memcpy(tgt, src, width);
tgt += width;
src += linesize;
}
}
// Prepare intermediate UV plane
torch::Tensor uv = get_buffer({1, height / 2, width / 2, 2});
{
uint8_t* tgt = uv.data_ptr<uint8_t>();
uint8_t* src = pFrame->data[1];
int linesize = pFrame->linesize[1];
for (int h = 0; h < height / 2; ++h) {
memcpy(tgt, src, width);
tgt += width;
src += linesize;
}
}
// Upsample width and height
namespace F = torch::nn::functional;
using namespace torch::indexing;
uv = F::interpolate(
uv.permute({0, 3, 1, 2}),
F::InterpolateFuncOptions()
.mode(torch::kNearest)
.size(std::vector<int64_t>({height, width})));
// Write to the UV plane
// yuv[:, 1:] = uv
yuv.index_put_({Slice(), Slice(1)}, uv);
}
#ifdef USE_CUDA
void write_nv12_cuda(AVFrame* pFrame, torch::Tensor& yuv) {
int height = static_cast<int>(yuv.size(2));
int width = static_cast<int>(yuv.size(3));
// Write Y plane directly
{
uint8_t* tgt = yuv.data_ptr<uint8_t>();
CUdeviceptr src = (CUdeviceptr)pFrame->data[0];
int linesize = pFrame->linesize[0];
TORCH_CHECK(
cudaSuccess ==
cudaMemcpy2D(
(void*)tgt,
width,
(const void*)src,
linesize,
width,
height,
cudaMemcpyDeviceToDevice),
"Failed to copy Y plane to Cuda tensor.");
}
// Preapare intermediate UV planes
torch::Tensor uv = get_buffer({1, height / 2, width / 2, 2}, yuv.device());
{
uint8_t* tgt = uv.data_ptr<uint8_t>();
CUdeviceptr src = (CUdeviceptr)pFrame->data[1];
int linesize = pFrame->linesize[1];
TORCH_CHECK(
cudaSuccess ==
cudaMemcpy2D(
(void*)tgt,
width,
(const void*)src,
linesize,
width,
height / 2,
cudaMemcpyDeviceToDevice),
"Failed to copy UV plane to Cuda tensor.");
}
// Upsample width and height
namespace F = torch::nn::functional;
using namespace torch::indexing;
uv = F::interpolate(
uv.permute({0, 3, 1, 2}),
F::InterpolateFuncOptions()
.mode(torch::kNearest)
.size(std::vector<int64_t>({height, width})));
// Write to the UV plane
// yuv[:, 1:] = uv
yuv.index_put_({Slice(), Slice(1)}, uv);
}
void write_p010_cuda(AVFrame* pFrame, torch::Tensor& yuv) {
int height = static_cast<int>(yuv.size(2));
int width = static_cast<int>(yuv.size(3));
// Write Y plane directly
{
int16_t* tgt = yuv.data_ptr<int16_t>();
CUdeviceptr src = (CUdeviceptr)pFrame->data[0];
int linesize = pFrame->linesize[0];
TORCH_CHECK(
cudaSuccess ==
cudaMemcpy2D(
(void*)tgt,
width * 2,
(const void*)src,
linesize,
width * 2,
height,
cudaMemcpyDeviceToDevice),
"Failed to copy Y plane to Cuda tensor.");
}
// Prepare intermediate UV planes
torch::Tensor uv =
get_buffer({1, height / 2, width / 2, 2}, yuv.device(), torch::kInt16);
{
int16_t* tgt = uv.data_ptr<int16_t>();
CUdeviceptr src = (CUdeviceptr)pFrame->data[1];
int linesize = pFrame->linesize[1];
TORCH_CHECK(
cudaSuccess ==
cudaMemcpy2D(
(void*)tgt,
width * 2,
(const void*)src,
linesize,
width * 2,
height / 2,
cudaMemcpyDeviceToDevice),
"Failed to copy UV plane to Cuda tensor.");
}
uv = uv.permute({0, 3, 1, 2});
using namespace torch::indexing;
// Write to the UV plane
// very simplistic upscale using indexing since interpolate doesn't support
// shorts
yuv.index_put_(
{Slice(), Slice(1, 3), Slice(None, None, 2), Slice(None, None, 2)}, uv);
yuv.index_put_(
{Slice(), Slice(1, 3), Slice(1, None, 2), Slice(None, None, 2)}, uv);
yuv.index_put_(
{Slice(), Slice(1, 3), Slice(None, None, 2), Slice(1, None, 2)}, uv);
yuv.index_put_(
{Slice(), Slice(1, 3), Slice(1, None, 2), Slice(1, None, 2)}, uv);
// correct for int16
yuv += 32768;
}
#endif
void write_image(AVFrame* frame, torch::Tensor& buf) {
// ref:
// https://ffmpeg.org/doxygen/4.1/filtering__video_8c_source.html#l00179
// https://ffmpeg.org/doxygen/4.1/decode__video_8c_source.html#l00038
AVPixelFormat format = static_cast<AVPixelFormat>(frame->format);
switch (format) {
case AV_PIX_FMT_RGB24:
case AV_PIX_FMT_BGR24:
case AV_PIX_FMT_ARGB:
case AV_PIX_FMT_RGBA:
case AV_PIX_FMT_ABGR:
case AV_PIX_FMT_BGRA:
case AV_PIX_FMT_GRAY8: {
write_interlaced_image(frame, buf);
return;
}
case AV_PIX_FMT_YUV444P: {
write_planar_image(frame, buf);
return;
}
case AV_PIX_FMT_YUV420P: {
write_yuv420p(frame, buf);
return;
}
case AV_PIX_FMT_NV12: {
write_nv12_cpu(frame, buf);
return;
}
case AV_PIX_FMT_RGB48LE: {
write_interlaced_image16(frame, buf);
return;
}
#ifdef USE_CUDA
case AV_PIX_FMT_CUDA: {
AVHWFramesContext* hwctx = (AVHWFramesContext*)frame->hw_frames_ctx->data;
AVPixelFormat sw_format = hwctx->sw_format;
// cuvid decoder (nvdec frontend of ffmpeg) only supports the following
// output formats
// https://github.com/FFmpeg/FFmpeg/blob/072101bd52f7f092ee976f4e6e41c19812ad32fd/libavcodec/cuviddec.c#L1121-L1124
switch (sw_format) {
case AV_PIX_FMT_NV12: {
write_nv12_cuda(frame, buf);
return;
}
case AV_PIX_FMT_P010: {
write_p010_cuda(frame, buf);
return;
}
case AV_PIX_FMT_P016:
TORCH_CHECK(
false,
"Unsupported video format found in CUDA HW: " +
std::string(av_get_pix_fmt_name(sw_format)));
default:
TORCH_CHECK(
false,
"Unexpected video format found in CUDA HW: " +
std::string(av_get_pix_fmt_name(sw_format)));
}
}
#endif
default:
TORCH_CHECK(
false,
"Unexpected video format: " +
std::string(av_get_pix_fmt_name(format)));
}
}
} // namespace
torch::Tensor convert_image(AVFrame* frame, const torch::Device& device) {
auto [buffer, is_planar] = get_image_buffer(frame, 1, device);
write_image(frame, buffer);
return is_planar ? buffer : buffer.permute({0, 3, 1, 2});
}
} // namespace detail
} // namespace io
} // namespace torchaudio
torchaudio/csrc/ffmpeg/stream_reader/buffer/common.h deleted 100644 → 0
View file @ 92f2ea89
#pragma once
#include <torch/torch.h>
#include <torchaudio/csrc/ffmpeg/ffmpeg.h>
namespace torchaudio {
namespace io {
namespace detail {
//////////////////////////////////////////////////////////////////////////////
// Helper functions
//////////////////////////////////////////////////////////////////////////////
torch::Tensor convert_audio(AVFrame* frame);
torch::Tensor convert_image(AVFrame* frame, const torch::Device& device);
} // namespace detail
} // namespace io
} // namespace torchaudio
torchaudio/csrc/ffmpeg/stream_reader/buffer/unchunked_buffer.cpp
View file @ 014d7140
#include <torchaudio/csrc/ffmpeg/stream_reader/buffer/common.h>
#include <torchaudio/csrc/ffmpeg/stream_reader/buffer/unchunked_buffer.h> #include <torchaudio/csrc/ffmpeg/stream_reader/buffer/unchunked_buffer.h>
#include <torchaudio/csrc/ffmpeg/stream_reader/conversion.h>
namespace torchaudio { namespace torchaudio {
namespace io { namespace io {
namespace detail { namespace detail {
UnchunkedVideoBuffer::UnchunkedVideoBuffer(const torch::Device& device) template <typename Converter>
: device(device) {} UnchunkedBuffer<Converter>::UnchunkedBuffer(Converter&& converter_)
: converter(std::move(converter_)) {}
bool UnchunkedBuffer::is_ready() const { template <typename Converter>
bool UnchunkedBuffer<Converter>::is_ready() const {
return chunks.size() > 0; return chunks.size() > 0;
} }
void UnchunkedBuffer::push_tensor(const torch::Tensor& t, double pts_) { template <typename Converter>
void UnchunkedBuffer<Converter>::push_frame(AVFrame* frame, double pts_) {
if (chunks.size() == 0) { if (chunks.size() == 0) {
pts = pts_; pts = pts_;
} }
chunks.push_back(t); chunks.push_back(converter.convert(frame));
} }
void UnchunkedAudioBuffer::push_frame(AVFrame* frame, double pts_) { template <typename Converter>
push_tensor(convert_audio(frame), pts_); c10::optional<Chunk> UnchunkedBuffer<Converter>::pop_chunk() {
}
void UnchunkedVideoBuffer::push_frame(AVFrame* frame, double pts_) {
push_tensor(convert_image(frame, device), pts_);
}
c10::optional<Chunk> UnchunkedBuffer::pop_chunk() {
if (chunks.size() == 0) { if (chunks.size() == 0) {
return {}; return {};
} }
...@@ -38,10 +34,143 @@ c10::optional<Chunk> UnchunkedBuffer::pop_chunk() { ...@@ -38,10 +34,143 @@ c10::optional<Chunk> UnchunkedBuffer::pop_chunk() {
return {Chunk{frames, pts}}; return {Chunk{frames, pts}};
} }
void UnchunkedBuffer::flush() { template <typename Converter>
void UnchunkedBuffer<Converter>::flush() {
chunks.clear(); chunks.clear();
} }
std::unique_ptr<Buffer> get_unchunked_buffer(AVSampleFormat fmt, int channels) {
switch (fmt) {
case AV_SAMPLE_FMT_U8: {
using Converter = AudioConverter<torch::kUInt8, false>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_S16: {
using Converter = AudioConverter<torch::kInt16, false>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_S32: {
using Converter = AudioConverter<torch::kInt32, false>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_S64: {
using Converter = AudioConverter<torch::kInt64, false>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_FLT: {
using Converter = AudioConverter<torch::kFloat32, false>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_DBL: {
using Converter = AudioConverter<torch::kFloat64, false>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_U8P: {
using Converter = AudioConverter<torch::kUInt8, true>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_S16P: {
using Converter = AudioConverter<torch::kInt16, true>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_S32P: {
using Converter = AudioConverter<torch::kInt32, true>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_S64P: {
using Converter = AudioConverter<torch::kInt64, true>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_FLTP: {
using Converter = AudioConverter<torch::kFloat32, true>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
case AV_SAMPLE_FMT_DBLP: {
using Converter = AudioConverter<torch::kFloat64, true>;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{channels});
}
default:
TORCH_INTERNAL_ASSERT(
false, "Unexpected audio type:", av_get_sample_fmt_name(fmt));
}
}
std::unique_ptr<Buffer> get_unchunked_buffer(
AVPixelFormat fmt,
int h,
int w,
const torch::Device& device) {
if (device.type() == at::DeviceType::CUDA) {
#ifndef USE_CUDA
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
false,
"USE_CUDA is not defined, and it should be guarded before here.");
#else
switch (fmt) {
case AV_PIX_FMT_NV12: {
using Conv = NV12CudaConverter;
return std::make_unique<UnchunkedBuffer<Conv>>(Conv{h, w, device});
}
case AV_PIX_FMT_P010: {
using Conv = P010CudaConverter;
return std::make_unique<UnchunkedBuffer<Conv>>(Conv{h, w, device});
}
case AV_PIX_FMT_P016: {
TORCH_CHECK(
false,
"Unsupported video format found in CUDA HW: ",
av_get_pix_fmt_name(fmt));
}
default: {
TORCH_CHECK(
false,
"Unexpected video format found in CUDA HW: ",
av_get_pix_fmt_name(fmt));
}
}
#endif
}
switch (fmt) {
case AV_PIX_FMT_RGB24:
case AV_PIX_FMT_BGR24: {
using Converter = InterlacedImageConverter;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{h, w, 3});
}
case AV_PIX_FMT_ARGB:
case AV_PIX_FMT_RGBA:
case AV_PIX_FMT_ABGR:
case AV_PIX_FMT_BGRA: {
using Converter = InterlacedImageConverter;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{h, w, 4});
}
case AV_PIX_FMT_GRAY8: {
using Converter = InterlacedImageConverter;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{h, w, 1});
}
case AV_PIX_FMT_RGB48LE: {
using Converter = Interlaced16BitImageConverter;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{h, w, 3});
}
case AV_PIX_FMT_YUV444P: {
using Converter = PlanarImageConverter;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{h, w, 3});
}
case AV_PIX_FMT_YUV420P: {
using Converter = YUV420PConverter;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{h, w});
}
case AV_PIX_FMT_NV12: {
using Converter = NV12Converter;
return std::make_unique<UnchunkedBuffer<Converter>>(Converter{h, w});
}
default: {
TORCH_INTERNAL_ASSERT(
false, "Unexpected video format found: ", av_get_pix_fmt_name(fmt));
}
}
}
} // namespace detail } // namespace detail
} // namespace io } // namespace io
} // namespace torchaudio } // namespace torchaudio
torchaudio/csrc/ffmpeg/stream_reader/buffer/unchunked_buffer.h
View file @ 014d7140
...@@ -4,43 +4,36 @@ ...@@ -4,43 +4,36 @@
#include <torchaudio/csrc/ffmpeg/stream_reader/buffer.h> #include <torchaudio/csrc/ffmpeg/stream_reader/buffer.h>
#include <deque> #include <deque>
namespace torchaudio { namespace torchaudio::io::detail {
namespace io {
namespace detail {
////////////////////////////////////////////////////////////////////////////// //////////////////////////////////////////////////////////////////////////////
// Unchunked Buffer Interface // Unchunked Buffer Interface
////////////////////////////////////////////////////////////////////////////// //////////////////////////////////////////////////////////////////////////////
// Partial implementation for unchunked buffer common to both audio and video // Partial implementation for unchunked buffer common to both audio and video
// Used for buffering audio/video streams without chunking // Used for buffering audio/video streams without chunking
template <typename Converter>
class UnchunkedBuffer : public Buffer { class UnchunkedBuffer : public Buffer {
// Each AVFrame is converted to a Tensor and stored here. // Each AVFrame is converted to a Tensor and stored here.
std::deque<torch::Tensor> chunks; std::deque<torch::Tensor> chunks;
double pts = -1.; double pts = -1.;
Converter converter;
protected:
void push_tensor(const torch::Tensor& t, double pts);
public: public:
UnchunkedBuffer(Converter&& converter);
bool is_ready() const override; bool is_ready() const override;
void push_frame(AVFrame* frame, double pts_) override;
c10::optional<Chunk> pop_chunk() override; c10::optional<Chunk> pop_chunk() override;
void flush() override; void flush() override;
}; };
class UnchunkedAudioBuffer : public UnchunkedBuffer { std::unique_ptr<Buffer> get_unchunked_buffer(
public: AVSampleFormat fmt,
void push_frame(AVFrame* frame, double pts) override; int num_channels);
};
class UnchunkedVideoBuffer : public UnchunkedBuffer {
const torch::Device device;
public: std::unique_ptr<Buffer> get_unchunked_buffer(
explicit UnchunkedVideoBuffer(const torch::Device& device); AVPixelFormat fmt,
int height,
void push_frame(AVFrame* frame, double pts) override; int width,
}; const torch::Device& device);
} // namespace detail } // namespace torchaudio::io::detail
} // namespace io
} // namespace torchaudio
torchaudio/csrc/ffmpeg/stream_reader/conversion.cpp 0 → 100644
View file @ 014d7140
#include <torch/torch.h>
#include <torchaudio/csrc/ffmpeg/stream_reader/conversion.h>
#ifdef USE_CUDA
#include <c10/cuda/CUDAStream.h>
#endif
namespace torchaudio::io {
////////////////////////////////////////////////////////////////////////////////
// Audio
////////////////////////////////////////////////////////////////////////////////
template <c10::ScalarType dtype, bool is_planar>
AudioConverter<dtype, is_planar>::AudioConverter(int c) : num_channels(c) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(num_channels > 0);
}
template <c10::ScalarType dtype, bool is_planar>
torch::Tensor AudioConverter<dtype, is_planar>::convert(const AVFrame* src) {
if constexpr (is_planar) {
torch::Tensor dst = torch::empty({num_channels, src->nb_samples}, dtype);
convert(src, dst);
return dst.permute({1, 0});
} else {
torch::Tensor dst = torch::empty({src->nb_samples, num_channels}, dtype);
convert(src, dst);
return dst;
}
}
// Converts AVFrame* into pre-allocated Tensor.
// The shape must be [C, T] if is_planar otherwise [T, C]
template <c10::ScalarType dtype, bool is_planar>
void AudioConverter<dtype, is_planar>::convert(
const AVFrame* src,
torch::Tensor& dst) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(num_channels == src->channels);
constexpr int bps = []() {
switch (dtype) {
case torch::kUInt8:
return 1;
case torch::kInt16:
return 2;
case torch::kInt32:
case torch::kFloat32:
return 4;
case torch::kInt64:
case torch::kFloat64:
return 8;
}
}();
// Note
// FFMpeg's `nb_samples` represnts the number of samples par channel.
// whereas, in torchaudio, `num_samples` is used to represent the number of
// samples across channels. torchaudio uses `num_frames` for per-channel
// samples.
if constexpr (is_planar) {
int plane_size = bps * src->nb_samples;
uint8_t* p_dst = static_cast<uint8_t*>(dst.data_ptr());
for (int i = 0; i < num_channels; ++i) {
memcpy(p_dst, src->extended_data[i], plane_size);
p_dst += plane_size;
}
} else {
int plane_size = bps * src->nb_samples * num_channels;
memcpy(dst.data_ptr(), src->extended_data[0], plane_size);
}
}
// Explicit instantiation
template class AudioConverter<torch::kUInt8, false>;
template class AudioConverter<torch::kUInt8, true>;
template class AudioConverter<torch::kInt16, false>;
template class AudioConverter<torch::kInt16, true>;
template class AudioConverter<torch::kInt32, false>;
template class AudioConverter<torch::kInt32, true>;
template class AudioConverter<torch::kInt64, false>;
template class AudioConverter<torch::kInt64, true>;
template class AudioConverter<torch::kFloat32, false>;
template class AudioConverter<torch::kFloat32, true>;
template class AudioConverter<torch::kFloat64, false>;
template class AudioConverter<torch::kFloat64, true>;
////////////////////////////////////////////////////////////////////////////////
// Image
////////////////////////////////////////////////////////////////////////////////
namespace {
torch::Tensor get_image_buffer(
at::IntArrayRef shape,
const torch::Dtype dtype = torch::kUInt8) {
return torch::empty(
shape, torch::TensorOptions().dtype(dtype).layout(torch::kStrided));
}
torch::Tensor get_image_buffer(
at::IntArrayRef shape,
torch::Device device,
const torch::Dtype dtype = torch::kUInt8) {
return torch::empty(
shape,
torch::TensorOptions()
.dtype(dtype)
.layout(torch::kStrided)
.device(device));
}
} // namespace
ImageConverterBase::ImageConverterBase(int h, int w, int c)
: height(h), width(w), num_channels(c) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(height > 0);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(width > 0);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(num_channels > 0);
}
////////////////////////////////////////////////////////////////////////////////
// Interlaced Image
////////////////////////////////////////////////////////////////////////////////
void InterlacedImageConverter::convert(const AVFrame* src, torch::Tensor& dst) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->height == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(1) == height);
int stride = width * num_channels;
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(2) * dst.size(3) == stride);
auto p_dst = dst.data_ptr<uint8_t>();
uint8_t* p_src = src->data[0];
for (int i = 0; i < height; ++i) {
memcpy(p_dst, p_src, stride);
p_src += src->linesize[0];
p_dst += stride;
}
}
torch::Tensor InterlacedImageConverter::convert(const AVFrame* src) {
torch::Tensor buffer = get_image_buffer({1, height, width, num_channels});
convert(src, buffer);
return buffer.permute({0, 3, 1, 2});
}
////////////////////////////////////////////////////////////////////////////////
// Interlaced 16 Bit Image
////////////////////////////////////////////////////////////////////////////////
void Interlaced16BitImageConverter::convert(
const AVFrame* src,
torch::Tensor& dst) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->height == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(1) == height);
int stride = width * num_channels;
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(2) * dst.size(3) == stride);
auto p_dst = dst.data_ptr<int16_t>();
uint8_t* p_src = src->data[0];
for (int i = 0; i < height; ++i) {
memcpy(p_dst, p_src, stride * 2);
p_src += src->linesize[0];
p_dst += stride;
}
// correct for int16
dst += 32768;
}
torch::Tensor Interlaced16BitImageConverter::convert(const AVFrame* src) {
torch::Tensor buffer =
get_image_buffer({1, height, width, num_channels}, torch::kInt16);
convert(src, buffer);
return buffer.permute({0, 3, 1, 2});
}
////////////////////////////////////////////////////////////////////////////////
// Planar Image
////////////////////////////////////////////////////////////////////////////////
void PlanarImageConverter::convert(const AVFrame* src, torch::Tensor& dst) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->height == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->width == width);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(1) == num_channels);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(2) == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(3) == width);
for (int i = 0; i < num_channels; ++i) {
torch::Tensor plane = dst.index({0, i});
uint8_t* p_dst = plane.data_ptr<uint8_t>();
uint8_t* p_src = src->data[i];
int linesize = src->linesize[i];
for (int h = 0; h < height; ++h) {
memcpy(p_dst, p_src, width);
p_src += linesize;
p_dst += width;
}
}
}
torch::Tensor PlanarImageConverter::convert(const AVFrame* src) {
torch::Tensor buffer = get_image_buffer({1, num_channels, height, width});
convert(src, buffer);
return buffer;
}
////////////////////////////////////////////////////////////////////////////////
// YUV420P
////////////////////////////////////////////////////////////////////////////////
YUV420PConverter::YUV420PConverter(int h, int w)
: ImageConverterBase(h, w, 3),
tmp_uv(get_image_buffer({1, 2, height / 2, width / 2})) {}
void YUV420PConverter::convert(const AVFrame* src, torch::Tensor& dst) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
(AVPixelFormat)(src->format) == AV_PIX_FMT_YUV420P);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->height == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->width == width);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(1) == 3);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(2) == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(3) == width);
// Write Y plane directly
{
uint8_t* p_dst = dst.data_ptr<uint8_t>();
uint8_t* p_src = src->data[0];
for (int h = 0; h < height; ++h) {
memcpy(p_dst, p_src, width);
p_dst += width;
p_src += src->linesize[0];
}
}
// Write intermediate UV plane
{
uint8_t* p_dst = tmp_uv.data_ptr<uint8_t>();
uint8_t* p_src = src->data[1];
for (int h = 0; h < height / 2; ++h) {
memcpy(p_dst, p_src, width / 2);
p_dst += width / 2;
p_src += src->linesize[1];
}
p_src = src->data[2];
for (int h = 0; h < height / 2; ++h) {
memcpy(p_dst, p_src, width / 2);
p_dst += width / 2;
p_src += src->linesize[2];
}
}
// Upsample width and height
namespace F = torch::nn::functional;
torch::Tensor uv = F::interpolate(
tmp_uv,
F::InterpolateFuncOptions()
.mode(torch::kNearest)
.size(std::vector<int64_t>({height, width})));
// Write to the UV plane
// dst[:, 1:] = uv
using namespace torch::indexing;
dst.index_put_({Slice(), Slice(1)}, uv);
}
torch::Tensor YUV420PConverter::convert(const AVFrame* src) {
torch::Tensor buffer = get_image_buffer({1, num_channels, height, width});
convert(src, buffer);
return buffer;
}
////////////////////////////////////////////////////////////////////////////////
// NV12
////////////////////////////////////////////////////////////////////////////////
NV12Converter::NV12Converter(int h, int w)
: ImageConverterBase(h, w, 3),
tmp_uv(get_image_buffer({1, height / 2, width / 2, 2})) {}
void NV12Converter::convert(const AVFrame* src, torch::Tensor& dst) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
(AVPixelFormat)(src->format) == AV_PIX_FMT_NV12);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->height == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->width == width);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(1) == 3);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(2) == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(3) == width);
// Write Y plane directly
{
uint8_t* p_dst = dst.data_ptr<uint8_t>();
uint8_t* p_src = src->data[0];
for (int h = 0; h < height; ++h) {
memcpy(p_dst, p_src, width);
p_dst += width;
p_src += src->linesize[0];
}
}
// Write intermediate UV plane
{
uint8_t* p_dst = tmp_uv.data_ptr<uint8_t>();
uint8_t* p_src = src->data[1];
for (int h = 0; h < height / 2; ++h) {
memcpy(p_dst, p_src, width);
p_dst += width;
p_src += src->linesize[1];
}
}
// Upsample width and height
namespace F = torch::nn::functional;
torch::Tensor uv = F::interpolate(
tmp_uv.permute({0, 3, 1, 2}),
F::InterpolateFuncOptions()
.mode(torch::kNearest)
.size(std::vector<int64_t>({height, width})));
// Write to the UV plane
// dst[:, 1:] = uv
using namespace torch::indexing;
dst.index_put_({Slice(), Slice(1)}, uv);
}
torch::Tensor NV12Converter::convert(const AVFrame* src) {
torch::Tensor buffer = get_image_buffer({1, num_channels, height, width});
convert(src, buffer);
return buffer;
}
#ifdef USE_CUDA
////////////////////////////////////////////////////////////////////////////////
// NV12 CUDA
////////////////////////////////////////////////////////////////////////////////
NV12CudaConverter::NV12CudaConverter(int h, int w, const torch::Device& device)
: ImageConverterBase(h, w, 3),
tmp_uv(get_image_buffer(
{1, height / 2, width / 2, 2},
device,
torch::kUInt8)) {}
void NV12CudaConverter::convert(const AVFrame* src, torch::Tensor& dst) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->height == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->width == width);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(1) == 3);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(2) == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(3) == width);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.dtype() == torch::kUInt8);
auto fmt = (AVPixelFormat)(src->format);
AVHWFramesContext* hwctx = (AVHWFramesContext*)src->hw_frames_ctx->data;
AVPixelFormat sw_fmt = hwctx->sw_format;
TORCH_INTERNAL_ASSERT(
AV_PIX_FMT_CUDA == fmt,
"Expected CUDA frame. Found: ",
av_get_pix_fmt_name(fmt));
TORCH_INTERNAL_ASSERT(
AV_PIX_FMT_NV12 == sw_fmt,
"Expected NV12 format. Found: ",
av_get_pix_fmt_name(sw_fmt));
// Write Y plane directly
auto status = cudaMemcpy2D(
dst.data_ptr(),
width,
src->data[0],
src->linesize[0],
width,
height,
cudaMemcpyDeviceToDevice);
TORCH_CHECK(cudaSuccess == status, "Failed to copy Y plane to Cuda tensor.");
// Preapare intermediate UV planes
status = cudaMemcpy2D(
tmp_uv.data_ptr(),
width,
src->data[1],
src->linesize[1],
width,
height / 2,
cudaMemcpyDeviceToDevice);
TORCH_CHECK(cudaSuccess == status, "Failed to copy UV plane to Cuda tensor.");
// Upsample width and height
namespace F = torch::nn::functional;
torch::Tensor uv = F::interpolate(
tmp_uv.permute({0, 3, 1, 2}),
F::InterpolateFuncOptions()
.mode(torch::kNearest)
.size(std::vector<int64_t>({height, width})));
// Write to the UV plane
// dst[:, 1:] = uv
using namespace torch::indexing;
dst.index_put_({Slice(), Slice(1)}, uv);
}
torch::Tensor NV12CudaConverter::convert(const AVFrame* src) {
torch::Tensor buffer =
get_image_buffer({1, num_channels, height, width}, tmp_uv.device());
convert(src, buffer);
return buffer;
}
////////////////////////////////////////////////////////////////////////////////
// P010 CUDA
////////////////////////////////////////////////////////////////////////////////
P010CudaConverter::P010CudaConverter(int h, int w, const torch::Device& device)
: ImageConverterBase(h, w, 3),
tmp_uv(get_image_buffer(
{1, height / 2, width / 2, 2},
device,
torch::kInt16)) {}
void P010CudaConverter::convert(const AVFrame* src, torch::Tensor& dst) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->height == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(src->width == width);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(1) == 3);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(2) == height);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.size(3) == width);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dst.dtype() == torch::kInt16);
auto fmt = (AVPixelFormat)(src->format);
AVHWFramesContext* hwctx = (AVHWFramesContext*)src->hw_frames_ctx->data;
AVPixelFormat sw_fmt = hwctx->sw_format;
TORCH_INTERNAL_ASSERT(
AV_PIX_FMT_CUDA == fmt,
"Expected CUDA frame. Found: ",
av_get_pix_fmt_name(fmt));
TORCH_INTERNAL_ASSERT(
AV_PIX_FMT_P010 == sw_fmt,
"Expected P010 format. Found: ",
av_get_pix_fmt_name(sw_fmt));
// Write Y plane directly
auto status = cudaMemcpy2D(
dst.data_ptr(),
width * 2,
src->data[0],
src->linesize[0],
width * 2,
height,
cudaMemcpyDeviceToDevice);
TORCH_CHECK(cudaSuccess == status, "Failed to copy Y plane to CUDA tensor.");
// Prepare intermediate UV planes
status = cudaMemcpy2D(
tmp_uv.data_ptr(),
width * 2,
src->data[1],
src->linesize[1],
width * 2,
height / 2,
cudaMemcpyDeviceToDevice);
TORCH_CHECK(cudaSuccess == status, "Failed to copy UV plane to CUDA tensor.");
// Write to the UV plane
torch::Tensor uv = tmp_uv.permute({0, 3, 1, 2});
using namespace torch::indexing;
// very simplistic upscale using indexing since interpolate doesn't support
// shorts
dst.index_put_(
{Slice(), Slice(1, 3), Slice(None, None, 2), Slice(None, None, 2)}, uv);
dst.index_put_(
{Slice(), Slice(1, 3), Slice(1, None, 2), Slice(None, None, 2)}, uv);
dst.index_put_(
{Slice(), Slice(1, 3), Slice(None, None, 2), Slice(1, None, 2)}, uv);
dst.index_put_(
{Slice(), Slice(1, 3), Slice(1, None, 2), Slice(1, None, 2)}, uv);
// correct for int16
dst += 32768;
}
torch::Tensor P010CudaConverter::convert(const AVFrame* src) {
torch::Tensor buffer = get_image_buffer(
{1, num_channels, height, width}, tmp_uv.device(), torch::kInt16);
convert(src, buffer);
return buffer;
}
#endif
} // namespace torchaudio::io
torchaudio/csrc/ffmpeg/stream_reader/conversion.h 0 → 100644
View file @ 014d7140
#pragma once
#include <torch/types.h>
#include <torchaudio/csrc/ffmpeg/ffmpeg.h>
namespace torchaudio::io {
////////////////////////////////////////////////////////////////////////////////
// Audio
////////////////////////////////////////////////////////////////////////////////
template <c10::ScalarType dtype, bool is_planar>
class AudioConverter {
const int num_channels;
public:
AudioConverter(int num_channels);
// Converts AVFrame* into Tensor of [T, C]
torch::Tensor convert(const AVFrame* src);
// Converts AVFrame* into pre-allocated Tensor.
// The shape must be [C, T] if is_planar otherwise [T, C]
void convert(const AVFrame* src, torch::Tensor& dst);
};
////////////////////////////////////////////////////////////////////////////////
// Image
////////////////////////////////////////////////////////////////////////////////
struct ImageConverterBase {
const int height;
const int width;
const int num_channels;
ImageConverterBase(int h, int w, int c);
};
////////////////////////////////////////////////////////////////////////////////
// Interlaced Images - NHWC
////////////////////////////////////////////////////////////////////////////////
struct InterlacedImageConverter : public ImageConverterBase {
using ImageConverterBase::ImageConverterBase;
// convert AVFrame* into Tensor of NCHW format
torch::Tensor convert(const AVFrame* src);
// convert AVFrame* into pre-allocated Tensor of NHWC format
void convert(const AVFrame* src, torch::Tensor& dst);
};
struct Interlaced16BitImageConverter : public ImageConverterBase {
using ImageConverterBase::ImageConverterBase;
// convert AVFrame* into Tensor of NCHW format
torch::Tensor convert(const AVFrame* src);
// convert AVFrame* into pre-allocated Tensor of NHWC format
void convert(const AVFrame* src, torch::Tensor& dst);
};
////////////////////////////////////////////////////////////////////////////////
// Planar Images - NCHW
////////////////////////////////////////////////////////////////////////////////
struct PlanarImageConverter : public ImageConverterBase {
using ImageConverterBase::ImageConverterBase;
void convert(const AVFrame* src, torch::Tensor& dst);
torch::Tensor convert(const AVFrame* src);
};
////////////////////////////////////////////////////////////////////////////////
// Family of YUVs - NCHW
////////////////////////////////////////////////////////////////////////////////
class YUV420PConverter : public ImageConverterBase {
torch::Tensor tmp_uv;
public:
YUV420PConverter(int height, int width);
void convert(const AVFrame* src, torch::Tensor& dst);
torch::Tensor convert(const AVFrame* src);
};
class NV12Converter : public ImageConverterBase {
torch::Tensor tmp_uv;
public:
NV12Converter(int height, int width);
void convert(const AVFrame* src, torch::Tensor& dst);
torch::Tensor convert(const AVFrame* src);
};
#ifdef USE_CUDA
class NV12CudaConverter : ImageConverterBase {
torch::Tensor tmp_uv;
public:
NV12CudaConverter(int height, int width, const torch::Device& device);
void convert(const AVFrame* src, torch::Tensor& dst);
torch::Tensor convert(const AVFrame* src);
};
class P010CudaConverter : ImageConverterBase {
torch::Tensor tmp_uv;
public:
P010CudaConverter(int height, int width, const torch::Device& device);
void convert(const AVFrame* src, torch::Tensor& dst);
torch::Tensor convert(const AVFrame* src);
};
#endif
} // namespace torchaudio::io
torchaudio/csrc/ffmpeg/stream_reader/sink.cpp
View file @ 014d7140
...@@ -8,10 +8,10 @@ namespace io { ...@@ -8,10 +8,10 @@ namespace io {
namespace { namespace {
std::unique_ptr<Buffer> get_buffer( std::unique_ptr<Buffer> get_buffer(
AVMediaType type, AVCodecContext* codec_ctx,
FilterGraph& filter,
int frames_per_chunk, int frames_per_chunk,
int num_chunks, int num_chunks,
double frame_duration,
const torch::Device& device) { const torch::Device& device) {
TORCH_CHECK( TORCH_CHECK(
frames_per_chunk > 0 || frames_per_chunk == -1, frames_per_chunk > 0 || frames_per_chunk == -1,
...@@ -23,47 +23,71 @@ std::unique_ptr<Buffer> get_buffer( ...@@ -23,47 +23,71 @@ std::unique_ptr<Buffer> get_buffer(
"`num_chunks` must be positive or -1. Found: ", "`num_chunks` must be positive or -1. Found: ",
num_chunks); num_chunks);
TORCH_INTERNAL_ASSERT( auto info = filter.get_output_info();
type == AVMEDIA_TYPE_AUDIO || type == AVMEDIA_TYPE_VIDEO,
TORCH_CHECK(
info.type == AVMEDIA_TYPE_AUDIO || info.type == AVMEDIA_TYPE_VIDEO,
"Unsupported media type: ", "Unsupported media type: ",
av_get_media_type_string(type), av_get_media_type_string(info.type),
". Only video or audio is supported "); ". Only video or audio is supported ");
// Chunked Mode auto time_base = filter.get_output_timebase();
if (frames_per_chunk > 0) { double frame_duration = double(time_base.num) / time_base.den;
if (type == AVMEDIA_TYPE_AUDIO) {
return std::unique_ptr<Buffer>(new detail::ChunkedAudioBuffer( if (info.type == AVMEDIA_TYPE_AUDIO) {
frames_per_chunk, num_chunks, frame_duration)); AVSampleFormat fmt = (AVSampleFormat)(info.format);
if (frames_per_chunk == -1) {
return detail::get_unchunked_buffer(fmt, codec_ctx->channels);
} else {
return detail::get_chunked_buffer(
frames_per_chunk,
num_chunks,
frame_duration,
fmt,
codec_ctx->channels);
}
} else { } else {
return std::unique_ptr<Buffer>(new detail::ChunkedVideoBuffer( // Note
frames_per_chunk, num_chunks, frame_duration, device)); // When using HW decoder, the pixel format is CUDA, and FilterGraph does
// not yet support CUDA frames, nor propagating the software pixel format,
// so here, we refer to AVCodecContext* to look at the pixel format.
AVPixelFormat fmt = (AVPixelFormat)(info.format);
if (fmt == AV_PIX_FMT_CUDA) {
fmt = codec_ctx->sw_pix_fmt;
} }
} else { // unchunked mode
if (type == AVMEDIA_TYPE_AUDIO) { if (frames_per_chunk == -1) {
return std::unique_ptr<Buffer>(new detail::UnchunkedAudioBuffer()); return detail::get_unchunked_buffer(fmt, info.height, info.width, device);
} else { } else {
return std::unique_ptr<Buffer>(new detail::UnchunkedVideoBuffer(device)); return detail::get_chunked_buffer(
frames_per_chunk,
num_chunks,
frame_duration,
fmt,
info.height,
info.width,
device);
} }
} }
} }
std::unique_ptr<FilterGraph> get_filter_graph( FilterGraph get_filter_graph(
AVRational input_time_base, AVRational input_time_base,
AVCodecContext* codec_ctx, AVCodecContext* codec_ctx,
AVRational frame_rate, AVRational frame_rate,
const std::string& filter_description) { const std::string& filter_description) {
auto p = std::make_unique<FilterGraph>(codec_ctx->codec_type); auto p = FilterGraph{codec_ctx->codec_type};
switch (codec_ctx->codec_type) { switch (codec_ctx->codec_type) {
case AVMEDIA_TYPE_AUDIO: case AVMEDIA_TYPE_AUDIO:
p->add_audio_src( p.add_audio_src(
codec_ctx->sample_fmt, codec_ctx->sample_fmt,
input_time_base, input_time_base,
codec_ctx->sample_rate, codec_ctx->sample_rate,
codec_ctx->channel_layout); codec_ctx->channel_layout);
break; break;
case AVMEDIA_TYPE_VIDEO: case AVMEDIA_TYPE_VIDEO:
p->add_video_src( p.add_video_src(
codec_ctx->pix_fmt, codec_ctx->pix_fmt,
input_time_base, input_time_base,
frame_rate, frame_rate,
...@@ -74,9 +98,9 @@ std::unique_ptr<FilterGraph> get_filter_graph( ...@@ -74,9 +98,9 @@ std::unique_ptr<FilterGraph> get_filter_graph(
default: default:
TORCH_CHECK(false, "Only audio/video are supported."); TORCH_CHECK(false, "Only audio/video are supported.");
} }
p->add_sink(); p.add_sink();
p->add_process(filter_description); p.add_process(filter_description);
p->create_filter(); p.create_filter();
return p; return p;
} }
...@@ -100,20 +124,17 @@ Sink::Sink( ...@@ -100,20 +124,17 @@ Sink::Sink(
codec_ctx, codec_ctx,
frame_rate, frame_rate,
filter_description)), filter_description)),
output_time_base(filter->get_output_timebase()), output_time_base(filter.get_output_timebase()),
buffer(get_buffer( buffer(
codec_ctx->codec_type, get_buffer(codec_ctx, filter, frames_per_chunk, num_chunks, device)) {
frames_per_chunk, }
num_chunks,
double(output_time_base.num) / output_time_base.den,
device)) {}
// 0: some kind of success // 0: some kind of success
// <0: Some error happened // <0: Some error happened
int Sink::process_frame(AVFrame* pFrame) { int Sink::process_frame(AVFrame* pFrame) {
int ret = filter->add_frame(pFrame); int ret = filter.add_frame(pFrame);
while (ret >= 0) { while (ret >= 0) {
ret = filter->get_frame(frame); ret = filter.get_frame(frame);
// AVERROR(EAGAIN) means that new input data is required to return new // AVERROR(EAGAIN) means that new input data is required to return new
// output. // output.
if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) { if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
...@@ -134,7 +155,7 @@ std::string Sink::get_filter_description() const { ...@@ -134,7 +155,7 @@ std::string Sink::get_filter_description() const {
} }
FilterGraphOutputInfo Sink::get_filter_output_info() const { FilterGraphOutputInfo Sink::get_filter_output_info() const {
return filter->get_output_info(); return filter.get_output_info();
} }
void Sink::flush() { void Sink::flush() {
......
torchaudio/csrc/ffmpeg/stream_reader/sink.h
View file @ 014d7140
...@@ -15,7 +15,7 @@ class Sink { ...@@ -15,7 +15,7 @@ class Sink {
AVCodecContext* codec_ctx; AVCodecContext* codec_ctx;
AVRational frame_rate; AVRational frame_rate;
std::string filter_description; std::string filter_description;
std::unique_ptr<FilterGraph> filter; FilterGraph filter;
// time_base of filter graph output, used for PTS calc // time_base of filter graph output, used for PTS calc
AVRational output_time_base; AVRational output_time_base;
......
torchaudio/csrc/ffmpeg/stream_reader/stream_processor.cpp
View file @ 014d7140
...@@ -50,6 +50,29 @@ enum AVPixelFormat get_hw_format( ...@@ -50,6 +50,29 @@ enum AVPixelFormat get_hw_format(
const AVCodecHWConfig* cfg = static_cast<AVCodecHWConfig*>(codec_ctx->opaque); const AVCodecHWConfig* cfg = static_cast<AVCodecHWConfig*>(codec_ctx->opaque);
for (const enum AVPixelFormat* p = pix_fmts; *p != -1; p++) { for (const enum AVPixelFormat* p = pix_fmts; *p != -1; p++) {
if (*p == cfg->pix_fmt) { if (*p == cfg->pix_fmt) {
// Note
// The HW decode example uses generic approach
// https://ffmpeg.org/doxygen/4.1/hw__decode_8c_source.html#l00063
// But this approach finalizes the codec configuration when the first
// frame comes in.
// We need to inspect the codec configuration right after the codec is
// opened.
// So we add short cut for known patterns.
// yuv420p (h264) -> nv12
// yuv420p10le (hevc/h265) -> p010le
switch (codec_ctx->pix_fmt) {
case AV_PIX_FMT_YUV420P: {
codec_ctx->pix_fmt = AV_PIX_FMT_CUDA;
codec_ctx->sw_pix_fmt = AV_PIX_FMT_NV12;
break;
}
case AV_PIX_FMT_YUV420P10LE: {
codec_ctx->pix_fmt = AV_PIX_FMT_CUDA;
codec_ctx->sw_pix_fmt = AV_PIX_FMT_P010LE;
break;
}
default:;
}
return *p; return *p;
} }
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment