Dump codec output to ivf/y4m

Bug: b/261160916, webrtc:14852
Change-Id: I19de2210aa03b56752db5ce8b6fd94498123d6f5
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/296260
Reviewed-by: Rasmus Brandt <brandtr@webrtc.org>
Commit-Queue: Sergey Silkin <ssilkin@webrtc.org>
Reviewed-by: Mirko Bonadei <mbonadei@webrtc.org>
Cr-Commit-Position: refs/heads/main@{#39490}
diff --git a/api/test/video_codec_tester.h b/api/test/video_codec_tester.h
index 149f0de..ab5a915 100644
--- a/api/test/video_codec_tester.h
+++ b/api/test/video_codec_tester.h
@@ -12,6 +12,7 @@
 #define API_TEST_VIDEO_CODEC_TESTER_H_
 
 #include <memory>
+#include <string>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/types/optional.h"
@@ -46,10 +47,12 @@
 
   struct DecoderSettings {
     PacingSettings pacing;
+    absl::optional<std::string> decoded_y4m_base_path;
   };
 
   struct EncoderSettings {
     PacingSettings pacing;
+    absl::optional<std::string> encoded_ivf_base_path;
   };
 
   virtual ~VideoCodecTester() = default;
diff --git a/modules/video_coding/codecs/test/video_codec_test.cc b/modules/video_coding/codecs/test/video_codec_test.cc
index 27a8a2b..89665ae 100644
--- a/modules/video_coding/codecs/test/video_codec_test.cc
+++ b/modules/video_coding/codecs/test/video_codec_test.cc
@@ -98,7 +98,9 @@
         frame_settings_(frame_settings),
         num_frames_(num_frames),
         frame_num_(0),
-        timestamp_rtp_(0) {
+        // Start with non-zero timestamp to force using frame RTP timestamps in
+        // IvfFrameWriter.
+        timestamp_rtp_(90000) {
     // Ensure settings for the first frame are provided.
     RTC_CHECK_GT(frame_settings_.size(), 0u);
     RTC_CHECK_EQ(frame_settings_.begin()->first, 0);
@@ -439,6 +441,16 @@
     f.target_framerate = layer_settings.framerate;
   }
 }
+
+std::string TestOutputPath() {
+  std::string output_path =
+      OutputPath() +
+      ::testing::UnitTest::GetInstance()->current_test_info()->name();
+  std::string output_dir = DirName(output_path);
+  bool result = CreateDir(output_dir);
+  RTC_CHECK(result) << "Cannot create " << output_dir;
+  return output_path;
+}
 }  // namespace
 
 std::unique_ptr<VideoCodecStats> RunEncodeDecodeTest(
@@ -446,7 +458,8 @@
     std::string codec_impl,
     const VideoInfo& video_info,
     const std::map<int, EncodingSettings>& frame_settings,
-    int num_frames) {
+    int num_frames,
+    bool save_codec_output) {
   std::unique_ptr<TestRawVideoSource> video_source =
       CreateVideoSource(video_info, frame_settings, num_frames);
 
@@ -468,6 +481,12 @@
           ? PacingMode::kRealTime
           : PacingMode::kNoPacing;
 
+  if (save_codec_output) {
+    std::string output_path = TestOutputPath();
+    encoder_settings.encoded_ivf_base_path = output_path;
+    decoder_settings.decoded_y4m_base_path = output_path;
+  }
+
   std::unique_ptr<VideoCodecTester> tester = CreateVideoCodecTester();
   return tester->RunEncodeDecodeTest(video_source.get(), encoder.get(),
                                      decoder.get(), encoder_settings,
@@ -479,7 +498,8 @@
     std::string codec_impl,
     const VideoInfo& video_info,
     const std::map<int, EncodingSettings>& frame_settings,
-    int num_frames) {
+    int num_frames,
+    bool save_codec_output) {
   std::unique_ptr<TestRawVideoSource> video_source =
       CreateVideoSource(video_info, frame_settings, num_frames);
 
@@ -492,6 +512,10 @@
           ? PacingMode::kRealTime
           : PacingMode::kNoPacing;
 
+  if (save_codec_output) {
+    encoder_settings.encoded_ivf_base_path = TestOutputPath();
+  }
+
   std::unique_ptr<VideoCodecTester> tester = CreateVideoCodecTester();
   return tester->RunEncodeTest(video_source.get(), encoder.get(),
                                encoder_settings);
@@ -535,8 +559,9 @@
   int duration_s = 10;
   int num_frames = duration_s * framerate_fps;
 
-  std::unique_ptr<VideoCodecStats> stats = RunEncodeDecodeTest(
-      codec_type, codec_impl, video_info, frame_settings, num_frames);
+  std::unique_ptr<VideoCodecStats> stats =
+      RunEncodeDecodeTest(codec_type, codec_impl, video_info, frame_settings,
+                          num_frames, /*save_codec_output=*/false);
 
   std::vector<VideoCodecStats::Frame> frames = stats->Slice();
   SetTargetRates(frame_settings, frames);
@@ -613,8 +638,9 @@
               .framerate = video_info.framerate,
               .bitrate = DataRate::KilobitsPerSec(bitrate_kbps.second)}}}}}};
 
-  std::unique_ptr<VideoCodecStats> stats = RunEncodeTest(
-      codec_type, codec_impl, video_info, frame_settings, num_frames);
+  std::unique_ptr<VideoCodecStats> stats =
+      RunEncodeTest(codec_type, codec_impl, video_info, frame_settings,
+                    num_frames, /*save_codec_output=*/false);
 
   std::vector<VideoCodecStats::Frame> frames =
       stats->Slice(VideoCodecStats::Filter{.first_frame = first_frame});
@@ -686,8 +712,9 @@
               .framerate = Frequency::MilliHertz(1000 * framerate_fps.second),
               .bitrate = DataRate::KilobitsPerSec(512)}}}}}};
 
-  std::unique_ptr<VideoCodecStats> stats = RunEncodeTest(
-      codec_type, codec_impl, video_info, frame_settings, num_frames);
+  std::unique_ptr<VideoCodecStats> stats =
+      RunEncodeTest(codec_type, codec_impl, video_info, frame_settings,
+                    num_frames, /*save_codec_output=*/false);
 
   std::vector<VideoCodecStats::Frame> frames =
       stats->Slice(VideoCodecStats::Filter{.first_frame = first_frame});
diff --git a/modules/video_coding/codecs/test/video_codec_tester_impl.cc b/modules/video_coding/codecs/test/video_codec_tester_impl.cc
index 3583b59..1d79933 100644
--- a/modules/video_coding/codecs/test/video_codec_tester_impl.cc
+++ b/modules/video_coding/codecs/test/video_codec_tester_impl.cc
@@ -12,6 +12,7 @@
 
 #include <map>
 #include <memory>
+#include <string>
 #include <utility>
 
 #include "api/task_queue/default_task_queue_factory.h"
@@ -20,11 +21,14 @@
 #include "api/units/timestamp.h"
 #include "api/video/encoded_image.h"
 #include "api/video/i420_buffer.h"
+#include "api/video/video_codec_type.h"
 #include "api/video/video_frame.h"
 #include "modules/video_coding/codecs/test/video_codec_analyzer.h"
+#include "modules/video_coding/utility/ivf_file_writer.h"
 #include "rtc_base/event.h"
 #include "rtc_base/time_utils.h"
 #include "system_wrappers/include/sleep.h"
+#include "test/testsupport/video_frame_writer.h"
 
 namespace webrtc {
 namespace test {
@@ -150,6 +154,39 @@
   rtc::Event task_executed_;
 };
 
+class TesterY4mWriter {
+ public:
+  explicit TesterY4mWriter(absl::string_view base_path)
+      : base_path_(base_path) {}
+
+  ~TesterY4mWriter() {
+    task_queue_.SendTask([] {});
+  }
+
+  void Write(const VideoFrame& frame, int spatial_idx) {
+    task_queue_.PostTask([this, frame, spatial_idx] {
+      if (y4m_writers_.find(spatial_idx) == y4m_writers_.end()) {
+        std::string file_path =
+            base_path_ + "_s" + std::to_string(spatial_idx) + ".y4m";
+
+        Y4mVideoFrameWriterImpl* y4m_writer = new Y4mVideoFrameWriterImpl(
+            file_path, frame.width(), frame.height(), /*fps=*/30);
+        RTC_CHECK(y4m_writer);
+
+        y4m_writers_[spatial_idx] =
+            std::unique_ptr<VideoFrameWriter>(y4m_writer);
+      }
+
+      y4m_writers_.at(spatial_idx)->WriteFrame(frame);
+    });
+  }
+
+ protected:
+  std::string base_path_;
+  std::map<int, std::unique_ptr<VideoFrameWriter>> y4m_writers_;
+  TaskQueueForTest task_queue_;
+};
+
 class TesterDecoder {
  public:
   TesterDecoder(Decoder* decoder,
@@ -160,6 +197,11 @@
         settings_(settings),
         pacer_(settings.pacing) {
     RTC_CHECK(analyzer_) << "Analyzer must be provided";
+
+    if (settings.decoded_y4m_base_path) {
+      y4m_writer_ =
+          std::make_unique<TesterY4mWriter>(*settings.decoded_y4m_base_path);
+    }
   }
 
   void Decode(const EncodedImage& frame) {
@@ -168,10 +210,15 @@
     task_queue_.PostScheduledTask(
         [this, frame] {
           analyzer_->StartDecode(frame);
+
           decoder_->Decode(
               frame, [this, spatial_idx = frame.SpatialIndex().value_or(0)](
                          const VideoFrame& decoded_frame) {
                 analyzer_->FinishDecode(decoded_frame, spatial_idx);
+
+                if (y4m_writer_) {
+                  y4m_writer_->Write(decoded_frame, spatial_idx);
+                }
               });
         },
         pacer_.Schedule(timestamp));
@@ -189,6 +236,45 @@
   const DecoderSettings& settings_;
   Pacer pacer_;
   LimitedTaskQueue task_queue_;
+  std::unique_ptr<TesterY4mWriter> y4m_writer_;
+};
+
+class TesterIvfWriter {
+ public:
+  explicit TesterIvfWriter(absl::string_view base_path)
+      : base_path_(base_path) {}
+
+  ~TesterIvfWriter() {
+    task_queue_.SendTask([] {});
+  }
+
+  void Write(const EncodedImage& encoded_frame) {
+    task_queue_.PostTask([this, encoded_frame] {
+      int spatial_idx = encoded_frame.SpatialIndex().value_or(0);
+      if (ivf_file_writers_.find(spatial_idx) == ivf_file_writers_.end()) {
+        std::string ivf_path =
+            base_path_ + "_s" + std::to_string(spatial_idx) + ".ivf";
+
+        FileWrapper ivf_file = FileWrapper::OpenWriteOnly(ivf_path);
+        RTC_CHECK(ivf_file.is_open());
+
+        std::unique_ptr<IvfFileWriter> ivf_writer =
+            IvfFileWriter::Wrap(std::move(ivf_file), /*byte_limit=*/0);
+        RTC_CHECK(ivf_writer);
+
+        ivf_file_writers_[spatial_idx] = std::move(ivf_writer);
+      }
+
+      // To play: ffplay -vcodec vp8|vp9|av1|hevc|h264 filename
+      ivf_file_writers_.at(spatial_idx)
+          ->WriteFrame(encoded_frame, VideoCodecType::kVideoCodecGeneric);
+    });
+  }
+
+ protected:
+  std::string base_path_;
+  std::map<int, std::unique_ptr<IvfFileWriter>> ivf_file_writers_;
+  TaskQueueForTest task_queue_;
 };
 
 class TesterEncoder {
@@ -203,6 +289,10 @@
         settings_(settings),
         pacer_(settings.pacing) {
     RTC_CHECK(analyzer_) << "Analyzer must be provided";
+    if (settings.encoded_ivf_base_path) {
+      ivf_writer_ =
+          std::make_unique<TesterIvfWriter>(*settings.encoded_ivf_base_path);
+    }
   }
 
   void Encode(const VideoFrame& frame) {
@@ -213,9 +303,14 @@
           analyzer_->StartEncode(frame);
           encoder_->Encode(frame, [this](const EncodedImage& encoded_frame) {
             analyzer_->FinishEncode(encoded_frame);
+
             if (decoder_ != nullptr) {
               decoder_->Decode(encoded_frame);
             }
+
+            if (ivf_writer_ != nullptr) {
+              ivf_writer_->Write(encoded_frame);
+            }
           });
         },
         pacer_.Schedule(timestamp));
@@ -232,6 +327,7 @@
   TesterDecoder* const decoder_;
   VideoCodecAnalyzer* const analyzer_;
   const EncoderSettings& settings_;
+  std::unique_ptr<TesterIvfWriter> ivf_writer_;
   Pacer pacer_;
   LimitedTaskQueue task_queue_;
 };