Stuff

Author: NicolasHug

src/torchcodec/_core/Encoder.cpp

@@ -10,6 +10,8 @@ namespace facebook::torchcodec {
 
 Encoder::~Encoder() {}
 
+// TODO-ENCODING: disable ffmpeg logs by default
+
 Encoder::Encoder(int sampleRate, std::string_view fileName)
     : sampleRate_(sampleRate) {
   AVFormatContext* avFormatContext = nullptr;
@@ -40,8 +42,8 @@ Encoder::Encoder(int sampleRate, std::string_view fileName)
 
   // This will use the default bit rate
   // TODO-ENCODING Should let user choose for compressed formats like mp3.
-  //   avCodecContext_->bit_rate = 0;
-  avCodecContext_->bit_rate = 24000;
+    // avCodecContext_->bit_rate = 64000;
+  avCodecContext_->bit_rate = 0;
 
   // FFmpeg will raise a reasonably informative error if the desired sample rate
   // isn't supported by the encoder.
@@ -134,6 +136,7 @@ void Encoder::encode(const torch::Tensor& wf) {
     auto numSamplesToEncode =
         std::min(numSamplesPerFrame, numSamples - numEncodedSamples);
     auto numBytesToEncode = numSamplesToEncode * numBytesPerSample;
+    avFrame->nb_samples = std::min(static_cast<int64_t>(avCodecContext_->frame_size), numSamplesToEncode);
 
     for (int ch = 0; ch < numChannels; ch++) {
       memcpy(
@@ -160,6 +163,11 @@ void Encoder::encode_inner_loop(
     AutoAVPacket& autoAVPacket,
     const UniqueAVFrame& avFrame) {
   auto status = avcodec_send_frame(avCodecContext_.get(), avFrame.get());
+//   if (avFrame.get()) {
+//     printf("Sending frame with %d samples\n", avFrame->nb_samples);
+//   } else {
+//     printf("Flushing\n");
+//   }
   TORCH_CHECK(
       status == AVSUCCESS,
       "Error while sending frame: ",

test/decoders/test_ops.py

@@ -940,8 +940,22 @@ def decode(self, source) -> torch.Tensor:
         )
         return frames
 
-    def test_round_trip(self, tmp_path):
-        asset = SINE_MONO_S32
+    # def test_round_trip(self, tmp_path):
+    #     asset = NASA_AUDIO_MP3
+
+    #     encoded_path = tmp_path / "output.mp3"
+    #     encoder = create_encoder(
+    #         sample_rate=asset.sample_rate, filename=str(encoded_path)
+    #     )
+
+    #     source_samples = self.decode(asset)
+    #     encode(encoder, source_samples)
+
+    #     torch.testing.assert_close(self.decode(encoded_path), source_samples)
+
+    def test_against_cli(self, tmp_path):
+
+        asset = NASA_AUDIO_MP3
 
         encoded_by_ffmpeg = tmp_path / "ffmpeg_output.mp3"
         encoded_by_us = tmp_path / "our_output.mp3"
@@ -951,9 +965,10 @@ def test_round_trip(self, tmp_path):
             "-i",
             str(asset.path),
             # '-vn',
-            # '-ar', '44100',    # Set audio sampling rate
+            # '-ar', '16000',    # Set audio sampling rate
             # '-ac', '2',        # Set number of audio channels
             # '-b:a', '192k',    # Set audio bitrate
+            '-b:a', '0',    # Set audio bitrate
             str(encoded_by_ffmpeg),
         ]
         subprocess.run(command, check=True)
@@ -964,8 +979,6 @@ def test_round_trip(self, tmp_path):
 
         encode(encoder, self.decode(asset))
 
-        print(encoded_by_ffmpeg)
-        print(encoded_by_us)
         from_ffmpeg = self.decode(encoded_by_ffmpeg)
         from_us = self.decode(encoded_by_us)
         torch.testing.assert_close(from_us, from_ffmpeg)