Let SimpleEncode be able to output bitstream

Add outfile_path to SimpleEncode() with default value NULL.
The encoder will only output bitstream when outfile_path is set.

Change-Id: Ic68e5358ea454358c510bb0ae214f4201cb3db39

4 files changed, 72 insertions, 23 deletions

diff --git a/ivfenc.c b/ivfenc.c
index a50d31839..2e8e04283 100644
--- a/ivfenc.c
+++ b/ivfenc.c
@@ -13,27 +13,35 @@
 #include "vpx/vpx_encoder.h"
 #include "vpx_ports/mem_ops.h"
 
-void ivf_write_file_header(FILE *outfile, const struct vpx_codec_enc_cfg *cfg,
-                           unsigned int fourcc, int frame_cnt) {
+void ivf_write_file_header_with_video_info(FILE *outfile, unsigned int fourcc,
+                                           int frame_cnt, int frame_width,
+                                           int frame_height,
+                                           vpx_rational_t timebase) {
   char header[32];
 
   header[0] = 'D';
   header[1] = 'K';
   header[2] = 'I';
   header[3] = 'F';
-  mem_put_le16(header + 4, 0);                     // version
-  mem_put_le16(header + 6, 32);                    // header size
-  mem_put_le32(header + 8, fourcc);                // fourcc
-  mem_put_le16(header + 12, cfg->g_w);             // width
-  mem_put_le16(header + 14, cfg->g_h);             // height
-  mem_put_le32(header + 16, cfg->g_timebase.den);  // rate
-  mem_put_le32(header + 20, cfg->g_timebase.num);  // scale
-  mem_put_le32(header + 24, frame_cnt);            // length
-  mem_put_le32(header + 28, 0);                    // unused
+  mem_put_le16(header + 4, 0);              // version
+  mem_put_le16(header + 6, 32);             // header size
+  mem_put_le32(header + 8, fourcc);         // fourcc
+  mem_put_le16(header + 12, frame_width);   // width
+  mem_put_le16(header + 14, frame_height);  // height
+  mem_put_le32(header + 16, timebase.den);  // rate
+  mem_put_le32(header + 20, timebase.num);  // scale
+  mem_put_le32(header + 24, frame_cnt);     // length
+  mem_put_le32(header + 28, 0);             // unused
 
   fwrite(header, 1, 32, outfile);
 }
 
+void ivf_write_file_header(FILE *outfile, const struct vpx_codec_enc_cfg *cfg,
+                           unsigned int fourcc, int frame_cnt) {
+  ivf_write_file_header_with_video_info(outfile, fourcc, frame_cnt, cfg->g_w,
+                                        cfg->g_h, cfg->g_timebase);
+}
+
 void ivf_write_frame_header(FILE *outfile, int64_t pts, size_t frame_size) {
   char header[12];
 
diff --git a/ivfenc.h b/ivfenc.h
index 483f2d2c5..27b691080 100644
--- a/ivfenc.h
+++ b/ivfenc.h
@@ -12,6 +12,8 @@
 
 #include "./tools_common.h"
 
+#include "vpx/vpx_encoder.h"
+
 struct vpx_codec_enc_cfg;
 struct vpx_codec_cx_pkt;
 
@@ -19,6 +21,11 @@ struct vpx_codec_cx_pkt;
 extern "C" {
 #endif
 
+void ivf_write_file_header_with_video_info(FILE *outfile, unsigned int fourcc,
+                                           int frame_cnt, int frame_width,
+                                           int frame_height,
+                                           vpx_rational_t timebase);
+
 void ivf_write_file_header(FILE *outfile, const struct vpx_codec_enc_cfg *cfg,
                            uint32_t fourcc, int frame_cnt);
 
diff --git a/vp9/simple_encode.cc b/vp9/simple_encode.cc
index e9c985df5..5229c3d8c 100644
--- a/vp9/simple_encode.cc
+++ b/vp9/simple_encode.cc
@@ -10,6 +10,7 @@
 
 #include <memory>
 #include <vector>
+#include "./ivfenc.h"
 #include "vp9/common/vp9_entropymode.h"
 #include "vp9/common/vp9_enums.h"
 #include "vp9/common/vp9_onyxc_int.h"
@@ -122,6 +123,13 @@ static INLINE vpx_rational_t make_vpx_rational(int num, int den) {
   return v;
 }
 
+static INLINE vpx_rational_t inverse_vpx_rational(vpx_rational_t v) {
+  vpx_rational_t inverse_v;
+  inverse_v.num = v.den;
+  inverse_v.den = v.num;
+  return inverse_v;
+}
+
 static INLINE FrameType
 get_frame_type_from_update_type(FRAME_UPDATE_TYPE update_type) {
   // TODO(angiebird): Figure out if we need frame type other than key frame,
@@ -517,7 +525,7 @@ static void UpdateGroupOfPicture(const VP9_COMP *cpi,
 SimpleEncode::SimpleEncode(int frame_width, int frame_height,
                            int frame_rate_num, int frame_rate_den,
                            int target_bitrate, int num_frames,
-                           const char *infile_path) {
+                           const char *infile_path, const char *outfile_path) {
   impl_ptr_ = std::unique_ptr<EncodeImpl>(new EncodeImpl());
   frame_width_ = frame_width;
   frame_height_ = frame_height;
@@ -526,7 +534,12 @@ SimpleEncode::SimpleEncode(int frame_width, int frame_height,
   target_bitrate_ = target_bitrate;
   num_frames_ = num_frames;
   // TODO(angirbid): Should we keep a file pointer here or keep the file_path?
-  file_ = fopen(infile_path, "r");
+  in_file_ = fopen(infile_path, "r");
+  if (outfile_path != NULL) {
+    out_file_ = fopen(outfile_path, "w");
+  } else {
+    out_file_ = NULL;
+  }
   impl_ptr_->cpi = NULL;
   impl_ptr_->img_fmt = VPX_IMG_FMT_I420;
 }
@@ -546,11 +559,11 @@ void SimpleEncode::ComputeFirstPassStats() {
 #endif
   vpx_image_t img;
   vpx_img_alloc(&img, impl_ptr_->img_fmt, frame_width_, frame_height_, 1);
-  rewind(file_);
+  rewind(in_file_);
   impl_ptr_->first_pass_stats.clear();
   for (i = 0; i < num_frames_; ++i) {
     assert(!vp9_lookahead_full(lookahead));
-    if (img_read(&img, file_)) {
+    if (img_read(&img, in_file_)) {
       int next_show_idx = vp9_lookahead_next_show_idx(lookahead);
       int64_t ts_start =
           timebase_units_to_ticks(&oxcf.g_timebase_in_ts, next_show_idx);
@@ -581,7 +594,7 @@ void SimpleEncode::ComputeFirstPassStats() {
   // TODO(angiebird): Store the total_stats apart form first_pass_stats
   impl_ptr_->first_pass_stats.push_back(vp9_get_total_stats(&cpi->twopass));
   free_encoder(cpi);
-  rewind(file_);
+  rewind(in_file_);
   vpx_img_free(&img);
 }
 
@@ -625,14 +638,22 @@ void SimpleEncode::StartEncode() {
   vpx_img_alloc(&impl_ptr_->tmp_img, impl_ptr_->img_fmt, frame_width_,
                 frame_height_, 1);
   UpdateGroupOfPicture(impl_ptr_->cpi, &group_of_picture_);
-  rewind(file_);
+  rewind(in_file_);
+
+  if (out_file_ != NULL) {
+    const char *fourcc = "VP90";
+    vpx_rational_t time_base = inverse_vpx_rational(frame_rate);
+    ivf_write_file_header_with_video_info(out_file_, *(const uint32_t *)fourcc,
+                                          num_frames_, frame_width_,
+                                          frame_height_, time_base);
+  }
 }
 
 void SimpleEncode::EndEncode() {
   free_encoder(impl_ptr_->cpi);
   impl_ptr_->cpi = nullptr;
   vpx_img_free(&impl_ptr_->tmp_img);
-  rewind(file_);
+  rewind(in_file_);
 }
 
 int SimpleEncode::GetKeyFrameGroupSize(int key_frame_index) const {
@@ -666,7 +687,7 @@ void SimpleEncode::EncodeFrame(EncodeFrameResult *encode_frame_result) {
   while (!vp9_lookahead_full(lookahead)) {
     // TODO(angiebird): Check whether we can move this file read logics to
     // lookahead
-    if (img_read(&impl_ptr_->tmp_img, file_)) {
+    if (img_read(&impl_ptr_->tmp_img, in_file_)) {
       int next_show_idx = vp9_lookahead_next_show_idx(lookahead);
       int64_t ts_start =
           timebase_units_to_ticks(&cpi->oxcf.g_timebase_in_ts, next_show_idx);
@@ -694,6 +715,13 @@ void SimpleEncode::EncodeFrame(EncodeFrameResult *encode_frame_result) {
                             &encode_frame_result->coding_data_byte_size,
                             encode_frame_result->coding_data.get(), &time_stamp,
                             &time_end, flush, &encode_frame_info);
+    if (out_file_ != NULL) {
+      ivf_write_frame_header(out_file_, time_stamp,
+                             encode_frame_result->coding_data_byte_size);
+      fwrite(encode_frame_result->coding_data.get(), 1,
+             encode_frame_result->coding_data_byte_size, out_file_);
+    }
+
     // vp9_get_compressed_data is expected to encode a frame every time, so the
     // data size should be greater than zero.
     if (encode_frame_result->coding_data_byte_size <= 0) {
@@ -763,8 +791,11 @@ uint64_t SimpleEncode::GetFramePixelCount() const {
 }
 
 SimpleEncode::~SimpleEncode() {
-  if (this->file_ != NULL) {
-    fclose(this->file_);
+  if (in_file_ != NULL) {
+    fclose(in_file_);
+  }
+  if (out_file_ != NULL) {
+    fclose(out_file_);
   }
 }
 
diff --git a/vp9/simple_encode.h b/vp9/simple_encode.h
index 741c49c7e..5d8f87676 100644
--- a/vp9/simple_encode.h
+++ b/vp9/simple_encode.h
@@ -202,9 +202,11 @@ struct GroupOfPicture {
 
 class SimpleEncode {
  public:
+  // When outfile_path is set, the encoder will output the bitstream in ivf
+  // format.
   SimpleEncode(int frame_width, int frame_height, int frame_rate_num,
                int frame_rate_den, int target_bitrate, int num_frames,
-               const char *infile_path);
+               const char *infile_path, const char *outfile_path = NULL);
   ~SimpleEncode();
   SimpleEncode(SimpleEncode &) = delete;
   SimpleEncode &operator=(const SimpleEncode &) = delete;
@@ -268,7 +270,8 @@ class SimpleEncode {
   int frame_rate_den_;
   int target_bitrate_;
   int num_frames_;
-  std::FILE *file_;
+  std::FILE *in_file_;
+  std::FILE *out_file_;
   std::unique_ptr<EncodeImpl> impl_ptr_;
 
   GroupOfPicture group_of_picture_;