Move timeBase to convertAVFrameToFrameOutput

Signed-off-by: Dmitry Rogozhkin <[email protected]>

Author: dvrogozh

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -16,9 +16,7 @@ namespace {
 
 bool g_cpu = registerDeviceInterface(
     torch::kCPU,
-    [](const torch::Device& device, const AVRational& timeBase) {
-      return new CpuDeviceInterface(device, timeBase);
-    });
+    [](const torch::Device& device) { return new CpuDeviceInterface(device); });
 
 } // namespace
 
@@ -36,10 +34,8 @@ bool CpuDeviceInterface::DecodedFrameContext::operator!=(
   return !(*this == other);
 }
 
-CpuDeviceInterface::CpuDeviceInterface(
-    const torch::Device& device,
-    const AVRational& timeBase)
-    : DeviceInterface(device, timeBase) {
+CpuDeviceInterface::CpuDeviceInterface(const torch::Device& device)
+    : DeviceInterface(device) {
   if (device_.type() != torch::kCPU) {
     throw std::runtime_error("Unsupported device: " + device_.str());
   }
@@ -56,6 +52,7 @@ CpuDeviceInterface::CpuDeviceInterface(
 // `dimension_order` parameter. It's up to callers to re-shape it if needed.
 void CpuDeviceInterface::convertAVFrameToFrameOutput(
     const VideoStreamOptions& videoStreamOptions,
+    const AVRational& timeBase,
     UniqueAVFrame& avFrame,
     FrameOutput& frameOutput,
     std::optional<torch::Tensor> preAllocatedOutputTensor) {
@@ -136,7 +133,7 @@ void CpuDeviceInterface::convertAVFrameToFrameOutput(
     frameOutput.data = outputTensor;
   } else if (colorConversionLibrary == ColorConversionLibrary::FILTERGRAPH) {
     if (!filterGraphContext_.filterGraph || prevFrameContext_ != frameContext) {
-      createFilterGraph(frameContext, videoStreamOptions);
+      createFilterGraph(frameContext, videoStreamOptions, timeBase);
       prevFrameContext_ = frameContext;
     }
     outputTensor = convertAVFrameToTensorUsingFilterGraph(avFrame);
@@ -215,7 +212,8 @@ torch::Tensor CpuDeviceInterface::convertAVFrameToTensorUsingFilterGraph(
 
 void CpuDeviceInterface::createFilterGraph(
     const DecodedFrameContext& frameContext,
-    const VideoStreamOptions& videoStreamOptions) {
+    const VideoStreamOptions& videoStreamOptions,
+    const AVRational& timeBase) {
   filterGraphContext_.filterGraph.reset(avfilter_graph_alloc());
   TORCH_CHECK(filterGraphContext_.filterGraph.get() != nullptr);
 
@@ -231,7 +229,7 @@ void CpuDeviceInterface::createFilterGraph(
   filterArgs << "video_size=" << frameContext.decodedWidth << "x"
              << frameContext.decodedHeight;
   filterArgs << ":pix_fmt=" << frameContext.decodedFormat;
-  filterArgs << ":time_base=" << timeBase_.num << "/" << timeBase_.den;
+  filterArgs << ":time_base=" << timeBase.num << "/" << timeBase.den;
   filterArgs << ":pixel_aspect=" << frameContext.decodedAspectRatio.num << "/"
              << frameContext.decodedAspectRatio.den;
 

src/torchcodec/_core/CpuDeviceInterface.h

@@ -13,7 +13,7 @@ namespace facebook::torchcodec {
 
 class CpuDeviceInterface : public DeviceInterface {
  public:
-  CpuDeviceInterface(const torch::Device& device, const AVRational& timeBase);
+  CpuDeviceInterface(const torch::Device& device);
 
   virtual ~CpuDeviceInterface() {}
 
@@ -27,6 +27,7 @@ class CpuDeviceInterface : public DeviceInterface {
 
   void convertAVFrameToFrameOutput(
       const VideoStreamOptions& videoStreamOptions,
+      const AVRational& timeBase,
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
       std::optional<torch::Tensor> preAllocatedOutputTensor =
@@ -63,7 +64,8 @@ class CpuDeviceInterface : public DeviceInterface {
 
   void createFilterGraph(
       const DecodedFrameContext& frameContext,
-      const VideoStreamOptions& videoStreamOptions);
+      const VideoStreamOptions& videoStreamOptions,
+      const AVRational& timeBase);
 
   // color-conversion fields. Only one of FilterGraphContext and
   // UniqueSwsContext should be non-null.

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -15,10 +15,9 @@ extern "C" {
 namespace facebook::torchcodec {
 namespace {
 
-bool g_cuda = registerDeviceInterface(
-    torch::kCUDA,
-    [](const torch::Device& device, const AVRational& timeBase) {
-      return new CudaDeviceInterface(device, timeBase);
+bool g_cuda =
+    registerDeviceInterface(torch::kCUDA, [](const torch::Device& device) {
+      return new CudaDeviceInterface(device);
     });
 
 // We reuse cuda contexts across VideoDeoder instances. This is because
@@ -164,10 +163,8 @@ AVBufferRef* getCudaContext(const torch::Device& device) {
 }
 } // namespace
 
-CudaDeviceInterface::CudaDeviceInterface(
-    const torch::Device& device,
-    const AVRational& timeBase)
-    : DeviceInterface(device, timeBase) {
+CudaDeviceInterface::CudaDeviceInterface(const torch::Device& device)
+    : DeviceInterface(device) {
   if (device_.type() != torch::kCUDA) {
     throw std::runtime_error("Unsupported device: " + device_.str());
   }
@@ -195,6 +192,7 @@ void CudaDeviceInterface::initializeContext(AVCodecContext* codecContext) {
 
 void CudaDeviceInterface::convertAVFrameToFrameOutput(
     const VideoStreamOptions& videoStreamOptions,
+    [[maybe_unused]] const AVRational& timeBase,
     UniqueAVFrame& avFrame,
     FrameOutput& frameOutput,
     std::optional<torch::Tensor> preAllocatedOutputTensor) {

src/torchcodec/_core/CudaDeviceInterface.h

@@ -12,7 +12,7 @@ namespace facebook::torchcodec {
 
 class CudaDeviceInterface : public DeviceInterface {
  public:
-  CudaDeviceInterface(const torch::Device& device, const AVRational& timeBase);
+  CudaDeviceInterface(const torch::Device& device);
 
   virtual ~CudaDeviceInterface();
 
@@ -22,6 +22,7 @@ class CudaDeviceInterface : public DeviceInterface {
 
   void convertAVFrameToFrameOutput(
       const VideoStreamOptions& videoStreamOptions,
+      const AVRational& timeBase,
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
       std::optional<torch::Tensor> preAllocatedOutputTensor =

src/torchcodec/_core/DeviceInterface.cpp

@@ -63,8 +63,7 @@ torch::Device createTorchDevice(const std::string device) {
 }
 
 std::unique_ptr<DeviceInterface> createDeviceInterface(
-    const torch::Device& device,
-    const AVRational& timeBase) {
+    const torch::Device& device) {
   auto deviceType = device.type();
   std::scoped_lock lock(g_interface_mutex);
   TORCH_CHECK(
@@ -73,7 +72,7 @@ std::unique_ptr<DeviceInterface> createDeviceInterface(
       device);
 
   return std::unique_ptr<DeviceInterface>(
-      (*g_interface_map)[deviceType](device, timeBase));
+      (*g_interface_map)[deviceType](device));
 }
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/DeviceInterface.h

@@ -27,8 +27,7 @@ namespace facebook::torchcodec {
 
 class DeviceInterface {
  public:
-  DeviceInterface(const torch::Device& device, const AVRational& timeBase)
-      : device_(device), timeBase_(timeBase) {}
+  DeviceInterface(const torch::Device& device) : device_(device) {}
 
   virtual ~DeviceInterface(){};
 
@@ -44,17 +43,17 @@ class DeviceInterface {
 
   virtual void convertAVFrameToFrameOutput(
       const VideoStreamOptions& videoStreamOptions,
+      const AVRational& timeBase,
       UniqueAVFrame& avFrame,
       FrameOutput& frameOutput,
       std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt) = 0;
 
  protected:
   torch::Device device_;
-  AVRational timeBase_;
 };
 
-using CreateDeviceInterfaceFn = std::function<
-    DeviceInterface*(const torch::Device& device, const AVRational& timeBase)>;
+using CreateDeviceInterfaceFn =
+    std::function<DeviceInterface*(const torch::Device& device)>;
 
 bool registerDeviceInterface(
     torch::DeviceType deviceType,
@@ -63,7 +62,6 @@ bool registerDeviceInterface(
 torch::Device createTorchDevice(const std::string device);
 
 std::unique_ptr<DeviceInterface> createDeviceInterface(
-    const torch::Device& device,
-    const AVRational& timeBase);
+    const torch::Device& device);
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -363,7 +363,7 @@ void SingleStreamDecoder::addStream(
   streamInfo.stream = formatContext_->streams[activeStreamIndex_];
   streamInfo.avMediaType = mediaType;
 
-  deviceInterface_ = createDeviceInterface(device, streamInfo.timeBase);
+  deviceInterface_ = createDeviceInterface(device);
 
   // This should never happen, checking just to be safe.
   TORCH_CHECK(
@@ -1151,6 +1151,7 @@ FrameOutput SingleStreamDecoder::convertAVFrameToFrameOutput(
   } else if (deviceInterface_) {
     deviceInterface_->convertAVFrameToFrameOutput(
         streamInfo.videoStreamOptions,
+        streamInfo.timeBase,
         avFrame,
         frameOutput,
         preAllocatedOutputTensor);