Skip to content

Commit 401901e

Browse files
authored
Initial C++ implementation of transforms (#902)
1 parent 0511d10 commit 401901e

22 files changed

+825
-551
lines changed

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -156,13 +156,20 @@ BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
156156
}
157157
}
158158

159-
void BetaCudaDeviceInterface::initializeInterface(AVStream* avStream) {
159+
void BetaCudaDeviceInterface::initialize(const AVStream* avStream) {
160160
torch::Tensor dummyTensorForCudaInitialization = torch::empty(
161161
{1}, torch::TensorOptions().dtype(torch::kUInt8).device(device_));
162162

163163
TORCH_CHECK(avStream != nullptr, "AVStream cannot be null");
164164
timeBase_ = avStream->time_base;
165165

166+
auto cudaDevice = torch::Device(torch::kCUDA);
167+
defaultCudaInterface_ =
168+
std::unique_ptr<DeviceInterface>(createDeviceInterface(cudaDevice));
169+
AVCodecContext dummyCodecContext = {};
170+
defaultCudaInterface_->initialize(avStream);
171+
defaultCudaInterface_->registerHardwareDeviceWithCodec(&dummyCodecContext);
172+
166173
const AVCodecParameters* codecpar = avStream->codecpar;
167174
TORCH_CHECK(codecpar != nullptr, "CodecParameters cannot be null");
168175

@@ -523,8 +530,6 @@ void BetaCudaDeviceInterface::flush() {
523530
}
524531

525532
void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
526-
const VideoStreamOptions& videoStreamOptions,
527-
const AVRational& timeBase,
528533
UniqueAVFrame& avFrame,
529534
FrameOutput& frameOutput,
530535
std::optional<torch::Tensor> preAllocatedOutputTensor) {
@@ -535,20 +540,8 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
535540
// TODONVDEC P1: we use the 'default' cuda device interface for color
536541
// conversion. That's a temporary hack to make things work. we should abstract
537542
// the color conversion stuff separately.
538-
if (!defaultCudaInterface_) {
539-
auto cudaDevice = torch::Device(torch::kCUDA);
540-
defaultCudaInterface_ =
541-
std::unique_ptr<DeviceInterface>(createDeviceInterface(cudaDevice));
542-
AVCodecContext dummyCodecContext = {};
543-
defaultCudaInterface_->initializeContext(&dummyCodecContext);
544-
}
545-
546543
defaultCudaInterface_->convertAVFrameToFrameOutput(
547-
videoStreamOptions,
548-
timeBase,
549-
avFrame,
550-
frameOutput,
551-
preAllocatedOutputTensor);
544+
avFrame, frameOutput, preAllocatedOutputTensor);
552545
}
553546

554547
BetaCudaDeviceInterface::FrameBuffer::Slot*

src/torchcodec/_core/BetaCudaDeviceInterface.h

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,9 @@ class BetaCudaDeviceInterface : public DeviceInterface {
3737
explicit BetaCudaDeviceInterface(const torch::Device& device);
3838
virtual ~BetaCudaDeviceInterface();
3939

40-
void initializeInterface(AVStream* stream) override;
40+
void initialize(const AVStream* avStream) override;
4141

4242
void convertAVFrameToFrameOutput(
43-
const VideoStreamOptions& videoStreamOptions,
44-
const AVRational& timeBase,
4543
UniqueAVFrame& avFrame,
4644
FrameOutput& frameOutput,
4745
std::optional<torch::Tensor> preAllocatedOutputTensor =

src/torchcodec/_core/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ function(make_torchcodec_libraries
9595
SingleStreamDecoder.cpp
9696
Encoder.cpp
9797
ValidationUtils.cpp
98+
Transform.cpp
9899
)
99100

100101
if(ENABLE_CUDA)

0 commit comments

Comments
 (0)