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

webrtc/modules/video_coding/utility/ivf_file_writer.cc

+/*
+ *  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/logging.h"
+#include "webrtc/modules/rtp_rtcp/source/byte_io.h"
+
+namespace webrtc {
+
+IvfFileWriter::IvfFileWriter(const char* file_name, FILE* file)
+    : num_frames_(0), file_name_(file_name), file_(file) {
+  RTC_CHECK(file != nullptr);
+}
+
+IvfFileWriter::~IvfFileWriter() {
+  Close();
+}
+
+const int kResolutionOffset = 12;
+const int kFrameCountOffset = 24;
+const size_t kIvfHeaderSize = 32;
+
+bool IvfFileWriter::SeekTo(size_t offset) {
+  int err = fseek(file_, offset, SEEK_SET);
+  RTC_DCHECK_EQ(0, err);
+  return err == 0;
+}
+
+std::unique_ptr<IvfFileWriter> IvfFileWriter::Open(const char* file_name,
+                                                   VideoCodecType codec_type) {
+  std::unique_ptr<IvfFileWriter> file_writer;
+
+  FILE* file = fopen(file_name, "wb");
+  if (file == nullptr) {
+    LOG(LS_ERROR) << "Unable to open IVF file " << file_name
+                  << " for binary output.";
+    return file_writer;
+  }
+
+  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 kVideoCodecVP8:
+      ivf_header[8] = 'V';
+      ivf_header[9] = 'P';
+      ivf_header[10] = '8';
+      ivf_header[11] = '0';
+      break;
+    case kVideoCodecVP9:
+      ivf_header[8] = 'V';
+      ivf_header[9] = 'P';
+      ivf_header[10] = '9';
+      ivf_header[11] = '0';
+      break;
+    case kVideoCodecH264:
+      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 file_writer;
+  }
+
+  // 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 (fwrite(ivf_header, sizeof(uint8_t), kIvfHeaderSize, file) !=
+      kIvfHeaderSize) {
+    LOG(LS_ERROR) << "Unable to write IVF header to file " << file_name;
+    fclose(file);
+    return file_writer;
+  }
+
+  file_writer.reset(new IvfFileWriter(file_name, file));
+  return file_writer;
+}
+
+bool IvfFileWriter::WriteFrame(const EncodedImage& encoded_image) {
+  RTC_CHECK(file_ != nullptr);
+
+  if (num_frames_ == 0) {
+    if (!SeekTo(kResolutionOffset))
+      return false;
+    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);
+    size_t bytes_written =
+        fwrite(buffer, sizeof(uint8_t), kResolutionDataSize, file_);
+    if (bytes_written != kResolutionDataSize) {
+      LOG(LS_ERROR) << "Unable to write resolution to IVF file " << file_name_;
+      return false;
+    }
+    if (!SeekTo(kIvfHeaderSize))
+      return false;
+  }
+
+  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.ntp_time_ms_);
+  size_t bytes_written =
+      fwrite(frame_header, sizeof(uint8_t), kFrameHeaderSize, file_);
+  bytes_written += fwrite(encoded_image._buffer, sizeof(uint8_t),
+                          encoded_image._length, file_);
+  if (bytes_written != kFrameHeaderSize + encoded_image._length) {
+    LOG(LS_ERROR) << "Unable to write frame to file " << file_name_;
+    return false;
+  }
+
+  ++num_frames_;
+  return true;
+}
+
+bool IvfFileWriter::Close() {
+  if (file_ == nullptr)
+    return false;
+
+  if (num_frames_ == 0) {
+    // No frame written to file, remove it entirely.
+    fclose(file_);
+    remove(file_name_.c_str());
+    return true;
+  }
+
+  if (!SeekTo(kFrameCountOffset))
+    return false;
+  const size_t kFrameCountLength = 4;
+  uint8_t size_data[kFrameCountLength];
+  ByteWriter<uint32_t>::WriteLittleEndian(size_data, num_frames_);
+  size_t bytes_written =
+      fwrite(size_data, sizeof(uint8_t), kFrameCountLength, file_);
+  RTC_DCHECK_EQ(kFrameCountLength, bytes_written);
+
+  fflush(file_);
+  fclose(file_);
+  file_ = nullptr;
+
+  return true;
+}
+
+}  // namespace webrtc