Add SimpleEncode::EncodeFrameWithQuantizeIndex()

Change-Id: I4442de01dfdbf13b0b9f7830f0fb393d3b935522

diff --git a/test/simple_encode_test.cc b/test/simple_encode_test.cc
index 3184d00588d1f8b9d151a7cea08d65e7b38291fa..15d13b18bb9b5d5be894288ee4f60d5bcb33e3ad 100644 (file)
--- a/test/simple_encode_test.cc
+++ b/test/simple_encode_test.cc
@@ -93,4 +93,28 @@ TEST(SimpleEncode, EncodeFrame) {
   simple_encode.EndEncode();
 }
 
+TEST(SimpleEncode, EncodeFrameWithQuantizeIndex) {
+  int w = 352;
+  int h = 288;
+  int frame_rate_num = 30;
+  int frame_rate_den = 1;
+  int target_bitrate = 1000;
+  int num_frames = 17;
+  // TODO(angiebird): Figure out how to upload test video to our codebase
+  FILE *file = fopen("bus_352x288_420_f20_b8.yuv", "r");
+  SimpleEncode simple_encode(w, h, frame_rate_num, frame_rate_den,
+                             target_bitrate, num_frames, file);
+  simple_encode.ComputeFirstPassStats();
+  int num_coding_frames = simple_encode.GetCodingFrameNum();
+  simple_encode.StartEncode();
+  for (int i = 0; i < num_coding_frames; ++i) {
+    int assigned_quantize_index = 100 + i;
+    EncodeFrameResult encode_frame_result;
+    simple_encode.EncodeFrameWithQuantizeIndex(&encode_frame_result,
+                                               assigned_quantize_index);
+    EXPECT_EQ(encode_frame_result.quantize_index, assigned_quantize_index);
+  }
+  simple_encode.EndEncode();
+}
+
 }  // namespace

diff --git a/vp9/encoder/vp9_encoder.c b/vp9/encoder/vp9_encoder.c
index 3104e71d1b2f0b62632b905b93c60c6308327ee1..ff47b21b5c551a7deb10d0403dfb8efba9140c96 100644 (file)
--- a/vp9/encoder/vp9_encoder.c
+++ b/vp9/encoder/vp9_encoder.c
@@ -2647,6 +2647,10 @@ VP9_COMP *vp9_create_compressor(const VP9EncoderConfig *oxcf,
 
   cm->error.setjmp = 0;
 
+#if CONFIG_RATE_CTRL
+  encode_command_init(&cpi->encode_command);
+#endif
+
   return cpi;
 }
 
@@ -4282,6 +4286,14 @@ static void encode_with_recode_loop(VP9_COMP *cpi, size_t *size,
       vp9_scale_references(cpi);
     }
 
+#if CONFIG_RATE_CTRL
+    // TODO(angiebird): This is a hack for making sure the encoder use the
+    // external_quantize_index exactly. Avoid this kind of hack later.
+    if (cpi->encode_command.use_external_quantize_index) {
+      q = cpi->encode_command.external_quantize_index;
+    }
+#endif
+
     vp9_set_quantizer(cm, q);
 
     if (loop_count == 0) setup_frame(cpi);
@@ -4307,6 +4319,13 @@ static void encode_with_recode_loop(VP9_COMP *cpi, size_t *size,
     // update_base_skip_probs(cpi);
 
     vpx_clear_system_state();
+#if CONFIG_RATE_CTRL
+    // TODO(angiebird): This is a hack for making sure the encoder use the
+    // external_quantize_index exactly. Avoid this kind of hack later.
+    if (cpi->encode_command.use_external_quantize_index) {
+      break;
+    }
+#endif
 
     // Dummy pack of the bitstream using up to date stats to get an
     // accurate estimate of output frame size to determine if we need

diff --git a/vp9/encoder/vp9_encoder.h b/vp9/encoder/vp9_encoder.h
index e12945f41838ac220fc58244c2885eb3658c037b..7c03ddeb8911f9249b25872d040fd7a7f5a60c62 100644 (file)
--- a/vp9/encoder/vp9_encoder.h
+++ b/vp9/encoder/vp9_encoder.h
@@ -516,6 +516,31 @@ typedef struct KMEANS_DATA {
   int group_idx;
 } KMEANS_DATA;
 
+#if CONFIG_RATE_CTRL
+typedef struct ENCODE_COMMAND {
+  int use_external_quantize_index;
+  int external_quantize_index;
+} ENCODE_COMMAND;
+
+static INLINE void encode_command_init(ENCODE_COMMAND *encode_command) {
+  vp9_zero(*encode_command);
+  encode_command->use_external_quantize_index = 0;
+  encode_command->external_quantize_index = -1;
+}
+
+static INLINE void encode_command_set_external_quantize_index(
+    ENCODE_COMMAND *encode_command, int quantize_index) {
+  encode_command->use_external_quantize_index = 1;
+  encode_command->external_quantize_index = quantize_index;
+}
+
+static INLINE void encode_command_reset_external_quantize_index(
+    ENCODE_COMMAND *encode_command) {
+  encode_command->use_external_quantize_index = 0;
+  encode_command->external_quantize_index = -1;
+}
+#endif
+
 typedef struct VP9_COMP {
   FRAME_INFO frame_info;
   QUANTS quants;
@@ -820,6 +845,9 @@ typedef struct VP9_COMP {
 
   int multi_layer_arf;
   vpx_roi_map_t roi;
+#if CONFIG_RATE_CTRL
+  ENCODE_COMMAND encode_command;
+#endif
 } VP9_COMP;
 
 typedef struct ENCODE_FRAME_RESULT {

diff --git a/vp9/simple_encode.cc b/vp9/simple_encode.cc
index 018d50a1eee90a302595779204d6f33ac1c0bcdc..38f87fc3a8403aa2b09717368a59d55686628db7 100644 (file)
--- a/vp9/simple_encode.cc
+++ b/vp9/simple_encode.cc
@@ -254,6 +254,14 @@ void SimpleEncode::EncodeFrame(EncodeFrameResult *encode_frame_result) {
   update_encode_frame_result(encode_frame_result, &encode_frame_info);
 }
 
+void SimpleEncode::EncodeFrameWithQuantizeIndex(
+    EncodeFrameResult *encode_frame_result, int quantize_index) {
+  encode_command_set_external_quantize_index(&pimpl->cpi->encode_command,
+                                             quantize_index);
+  EncodeFrame(encode_frame_result);
+  encode_command_reset_external_quantize_index(&pimpl->cpi->encode_command);
+}
+
 int SimpleEncode::GetCodingFrameNum() {
   assert(pimpl->first_pass_stats.size() - 1 > 0);
   // These are the default settings for now.

diff --git a/vp9/simple_encode.h b/vp9/simple_encode.h
index 48a86d0562d9b5b95d8f6879d7e7e0592794c240..e359f0553555028b043883c6acede63581f39771 100644 (file)
--- a/vp9/simple_encode.h
+++ b/vp9/simple_encode.h
@@ -47,6 +47,11 @@ class SimpleEncode {
   // This funtion should be called after StartEncode() before EndEncode()
   void EncodeFrame(EncodeFrameResult *encode_frame_result);
 
+  // Encode a frame with a specific quantize index
+  // This funtion should be called after StartEncode() before EndEncode()
+  void EncodeFrameWithQuantizeIndex(EncodeFrameResult *encode_frame_result,
+                                    int quantize_index);
+
   // Get the number of coding frames for the video. The coding frames include
   // show frame and no show frame.
   // This funtion should be called after ComputeFirstPassStats()