Cleanup BSF logic

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -252,6 +252,8 @@ int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) {
   CUVIDSOURCEDATAPACKET cuvidPacket = {};
 
   if (packet.get() && packet->data && packet->size > 0) {
+    applyBSF(packet);
+
     // Regular packet with data
     cuvidPacket.payload = packet->data;
     cuvidPacket.payload_size = packet->size;
@@ -275,27 +277,36 @@ int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) {
   return AVSUCCESS;
 }
 
-// TODONVDEC P0: cleanup this raw pointer / reference monstruosity.
-ReferenceAVPacket* BetaCudaDeviceInterface::applyBSF(
-    ReferenceAVPacket& packet,
-    [[maybe_unused]] AutoAVPacket& filteredAutoPacket,
-    ReferenceAVPacket& filteredPacket) {
+void BetaCudaDeviceInterface::applyBSF(ReferenceAVPacket& packet) {
   if (!bitstreamFilter_) {
-    return &packet;
+    return;
   }
+
   int retVal = av_bsf_send_packet(bitstreamFilter_.get(), packet.get());
   TORCH_CHECK(
       retVal >= AVSUCCESS,
       "Failed to send packet to bitstream filter: ",
       getFFMPEGErrorStringFromErrorCode(retVal));
 
+  // Create a temporary packet to receive the filtered data
+  // TODO P1: the docs mention there can theoretically be multiple output
+  // packets for a single input, i.e. we may need to call av_bsf_receive_packet
+  // more than once. We should figure out whether that applies to the BSF we're
+  // using.
+  AutoAVPacket filteredAutoPacket;
+  ReferenceAVPacket filteredPacket(filteredAutoPacket);
   retVal = av_bsf_receive_packet(bitstreamFilter_.get(), filteredPacket.get());
   TORCH_CHECK(
       retVal >= AVSUCCESS,
       "Failed to receive packet from bitstream filter: ",
       getFFMPEGErrorStringFromErrorCode(retVal));
 
-  return &filteredPacket;
+  // Free the original packet's data which isn't needed anymore, and move the
+  // fields of the filtered packet into the original packet. The filtered packet
+  // fields are re-set by av_packet_move_ref, so when it goes out of scope and
+  // gets destructed, it's not going to affect the original packet.
+  av_packet_unref(packet.get());
+  av_packet_move_ref(packet.get(), filteredPacket.get());
 }
 
 // Parser triggers this callback within cuvidParseVideoData when a frame is

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -54,16 +54,15 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   int sendPacket(ReferenceAVPacket& packet) override;
   int receiveFrame(UniqueAVFrame& avFrame, int64_t desiredPts) override;
   void flush() override;
-  ReferenceAVPacket* applyBSF(
-      ReferenceAVPacket& packet,
-      AutoAVPacket& filteredAutoPacket,
-      ReferenceAVPacket& filteredPacket) override;
 
   // NVDEC callback functions (must be public for C callbacks)
   unsigned char streamPropertyChange(CUVIDEOFORMAT* videoFormat);
   int frameReadyForDecoding(CUVIDPICPARAMS* pPicParams);
 
  private:
+  // Apply bitstream filter, modifies packet in-place
+  void applyBSF(ReferenceAVPacket& packet);
+
   class FrameBuffer {
    public:
     struct Slot {

src/torchcodec/_core/DeviceInterface.h

@@ -102,15 +102,6 @@ class DeviceInterface {
     // Custom decoders can override this method
   }
 
-  // Apply bitstream filter if needed, returns pointer to packet to use
-  // Default implementation returns the original packet (no filtering)
-  virtual ReferenceAVPacket* applyBSF(
-      ReferenceAVPacket& packet,
-      [[maybe_unused]] AutoAVPacket& filteredAutoPacket,
-      [[maybe_unused]] ReferenceAVPacket& filteredPacket) {
-    return &packet; // No filtering by default
-  }
-
  protected:
   torch::Device device_;
 };

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -1222,15 +1222,7 @@ UniqueAVFrame SingleStreamDecoder::decodeAVFrame(
     // We got a valid packet. Send it to the decoder, and we'll receive it in
     // the next iteration.
     if (useCustomInterface) {
-      // TODONVDEC P0:
-      // - cleanup this raw pointer / reference monstruosity.
-      // - don't even expose applyBSF in the interface. This should just be part
-      //   of sendPacket().
-      AutoAVPacket filteredAutoPacket;
-      ReferenceAVPacket filteredPacket(filteredAutoPacket);
-      ReferenceAVPacket* packetToSend = deviceInterface_->applyBSF(
-          packet, filteredAutoPacket, filteredPacket);
-      status = deviceInterface_->sendPacket(*packetToSend);
+      status = deviceInterface_->sendPacket(packet);
     } else {
       status = avcodec_send_packet(streamInfo.codecContext.get(), packet.get());
     }