Add support for writing raw encoder output to .ivf files.

webrtc/modules/video_coding/utility/ivf_file_writer.cc

Index: webrtc/modules/video_coding/utility/ivf_file_writer.cc
diff --git a/webrtc/modules/video_coding/utility/ivf_file_writer.cc b/webrtc/modules/video_coding/utility/ivf_file_writer.cc
new file mode 100644
index 0000000000000000000000000000000000000000..c977491c81e8397ebeebe71fe3c78d7e36f87c1e
--- /dev/null
+++ b/webrtc/modules/video_coding/utility/ivf_file_writer.cc
@@ -0,0 +1,195 @@
+/*
+ *  Copyright (c) 2016 The WebRTC project authors. All Rights Reserved.
+ *
+ *  Use of this source code is governed by a BSD-style license
+ *  that can be found in the LICENSE file in the root of the source
+ *  tree. An additional intellectual property rights grant can be found
+ *  in the file PATENTS.  All contributing project authors may
+ *  be found in the AUTHORS file in the root of the source tree.
+ */
+
+#include "webrtc/modules/video_coding/utility/ivf_file_writer.h"
+
+#include "webrtc/base/checks.h"
+#include "webrtc/base/fileutils.h"
+#include "webrtc/base/logging.h"
+#include "webrtc/modules/rtp_rtcp/source/byte_io.h"
+
+namespace webrtc {
+
+IvfFileWriter::IvfFileWriter(rtc::Pathname file_name,
+                             rtc::StreamInterface* out_stream)

[pbos-webrtc, 2016/04/07 16:16:04]

Should this take an unique_ptr as part of the interface?

[sprang_webrtc, 2016/04/11 15:57:28]

Done.

+    : num_frames_(0),
+      width_(0),
+      height_(0),
+      last_timestamp_(-1),
+      file_name_(file_name),
+      out_stream_(out_stream) {
+  RTC_CHECK(out_stream != nullptr);

[pbos-webrtc, 2016/04/07 16:16:04]

drop != nullptr

[sprang_webrtc, 2016/04/11 15:57:28]

Done.

+}
+
+IvfFileWriter::~IvfFileWriter() {
+  Close();
+}
+
+const int kResolutionOffset = 12;
+const int kFrameCountOffset = 24;
+const size_t kIvfHeaderSize = 32;
+
+std::unique_ptr<IvfFileWriter> IvfFileWriter::Open(
+    rtc::Pathname file_name,
+    RtpVideoCodecTypes codec_type) {
+  std::unique_ptr<IvfFileWriter> file_writer;
+  std::unique_ptr<rtc::FileStream> out_stream(new rtc::FileStream());
+  if (!out_stream->Open(file_name.pathname(), "wb", nullptr))
+    return file_writer;
+
+  file_writer.reset(new IvfFileWriter(file_name, out_stream.release()));
+  if (!file_writer->WriteHeader(codec_type))
+    file_writer.reset();
+
+  return file_writer;
+}
+
+std::unique_ptr<IvfFileWriter> IvfFileWriter::Open(
+    rtc::StreamInterface* out_stream,
+    RtpVideoCodecTypes codec_type) {
+  std::unique_ptr<IvfFileWriter> file_writer(
+      new IvfFileWriter(rtc::Pathname(), out_stream));
+  if (!file_writer->WriteHeader(codec_type))
+    file_writer.reset();
+
+  return file_writer;
+}
+
+bool IvfFileWriter::WriteHeader(RtpVideoCodecTypes codec_type) {
+  uint8_t ivf_header[kIvfHeaderSize] = {0};
+  ivf_header[0] = 'D';
+  ivf_header[1] = 'K';
+  ivf_header[2] = 'I';
+  ivf_header[3] = 'F';
+  ByteWriter<uint16_t>::WriteLittleEndian(&ivf_header[4], 0);   // Verison.
+  ByteWriter<uint16_t>::WriteLittleEndian(&ivf_header[6], 32);  // Header size.
+
+  switch (codec_type) {
+    case kRtpVideoVp8:
+      ivf_header[8] = 'V';
+      ivf_header[9] = 'P';
+      ivf_header[10] = '8';
+      ivf_header[11] = '0';
+      break;
+    case kRtpVideoVp9:
+      ivf_header[8] = 'V';
+      ivf_header[9] = 'P';
+      ivf_header[10] = '9';
+      ivf_header[11] = '0';
+      break;
+    case kRtpVideoH264:
+      ivf_header[8] = 'H';
+      ivf_header[9] = '2';
+      ivf_header[10] = '6';
+      ivf_header[11] = '4';
+      break;
+    default:
+      LOG(LS_ERROR) << "Unknown CODEC type: " << codec_type;
+      return false;
+  }
+
+  // Width and Height (populated on first frame).
+  ByteWriter<uint16_t>::WriteLittleEndian(&ivf_header[12], 0);
+  ByteWriter<uint16_t>::WriteLittleEndian(&ivf_header[14], 0);
+  // Timestamps in milliseconds => time scale of 1/1000.
+  ByteWriter<uint32_t>::WriteLittleEndian(&ivf_header[16], 1000);
+  ByteWriter<uint32_t>::WriteLittleEndian(&ivf_header[20], 1);
+  // Number of frames (populated on close).
+  ByteWriter<uint32_t>::WriteLittleEndian(&ivf_header[24], 0);
+  ByteWriter<uint32_t>::WriteLittleEndian(&ivf_header[28], 0);  // Reserved.
+
+  if (out_stream_->WriteAll(ivf_header, kIvfHeaderSize, nullptr, nullptr) !=
+      rtc::SR_SUCCESS) {
+    LOG(LS_ERROR) << "Unable to write IVF header.";
+    return false;
+  }
+
+  return true;
+}
+
+bool IvfFileWriter::WriteFrame(const EncodedImage& encoded_image) {
+  RTC_DCHECK(out_stream_->GetState() == rtc::SS_OPEN);
+
+  if (num_frames_ == 0) {
+    if (!out_stream_->SetPosition(kResolutionOffset))
+      return false;
+
+    width_ = encoded_image._encodedWidth;
+    height_ = encoded_image._encodedHeight;
+
+    const size_t kResolutionDataSize = 4;
+    uint8_t buffer[kResolutionDataSize];
+    ByteWriter<uint16_t>::WriteLittleEndian(&buffer[0],
+                                            encoded_image._encodedWidth);
+    ByteWriter<uint16_t>::WriteLittleEndian(&buffer[2],
+                                            encoded_image._encodedHeight);
+
+    if (out_stream_->WriteAll(buffer, kResolutionDataSize, nullptr, nullptr) !=
+        rtc::SR_SUCCESS) {
+      LOG(LS_ERROR) << "Unable to write resolution to IVF file "
+                    << file_name_.filename();
+      return false;
+    }
+    if (!out_stream_->SetPosition(kIvfHeaderSize))
+      return false;
+  }
+
+  RTC_CHECK_EQ(height_, encoded_image._encodedHeight);
+  RTC_CHECK_EQ(width_, encoded_image._encodedWidth);
+  if (last_timestamp_ != -1)
+    RTC_CHECK_GT(encoded_image.capture_time_ms_, last_timestamp_);
+  last_timestamp_ = encoded_image.capture_time_ms_;
+
+  const size_t kFrameHeaderSize = 12;
+  uint8_t frame_header[kFrameHeaderSize] = {};
+  ByteWriter<uint32_t>::WriteLittleEndian(&frame_header[0],
+                                          encoded_image._length);
+  ByteWriter<uint64_t>::WriteLittleEndian(&frame_header[4],
+                                          encoded_image.capture_time_ms_);
+  if (out_stream_->WriteAll(frame_header, kFrameHeaderSize, nullptr, nullptr) !=
+          rtc::SR_SUCCESS ||
+      out_stream_->WriteAll(encoded_image._buffer, encoded_image._length,
+                            nullptr, nullptr) != rtc::SR_SUCCESS) {
+    LOG(LS_ERROR) << "Unable to write frame to file " << file_name_.filename();
+    return false;
+  }
+
+  ++num_frames_;
+  return true;
+}
+
+bool IvfFileWriter::Close() {
+  if (out_stream_->GetState() != rtc::SS_OPEN)
+    return false;
+
+  if (num_frames_ == 0) {
+    // No frame written to file, close and remove it entirely if possible.
+    out_stream_->Close();
+    if (!file_name_.empty())
+      rtc::Filesystem::default_filesystem()->DeleteFile(file_name_);
+    return true;
+  }
+
+  if (!out_stream_->SetPosition(kFrameCountOffset))
+    return false;
+  const size_t kFrameCountLength = 4;
+  uint8_t size_data[kFrameCountLength];
+  ByteWriter<uint32_t>::WriteLittleEndian(size_data, num_frames_);
+  if (out_stream_->WriteAll(size_data, kFrameCountLength, nullptr, nullptr) !=
+      rtc::SR_SUCCESS) {
+    LOG(LS_ERROR) << "Unable to update header of IVF file when closing "
+                  << file_name_.filename();
+  }
+
+  out_stream_->Close();
+  return true;
+}
+
+}  // namespace webrtc