shithub: libvpx

--- a/test/simple_encode_test.cc

+++ b/test/simple_encode_test.cc

@@ -1,3 +1,4 @@

+#include <math.h>

 #include <memory>

 #include <vector>

 #include "third_party/googletest/src/include/gtest/gtest.h"

@@ -4,18 +5,24 @@

 #include "vp9/simple_encode.h"

 namespace {

+const int w = 352;

+const int h = 288;

+const int frame_rate_num = 30;

+const int frame_rate_den = 1;

+const int target_bitrate = 1000;

+const int num_frames = 17;

+// TODO(angiebird): Figure out how to upload test video to our codebase

+const char infile_path[] = "bus_352x288_420_f20_b8.yuv";

+static double get_bit_rate_in_kpbs(size_t bit_size, int num_frames,

+                                   int frame_rate_num, int frame_rate_den) {

+  return static_cast<double>(bit_size) / num_frames * frame_rate_num /

+         frame_rate_den / 1000.;

+}

 TEST(SimpleEncode, ComputeFirstPassStats) {

-  int w = 352;

-  int h = 288;

-  int frame_rate_num = 30;

-  int frame_rate_den = 1;

-  int target_bitrate = 200;

-  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);

+                             target_bitrate, num_frames, infile_path);

   simple_encode.ComputeFirstPassStats();

   std::vector<std::vector<double>> frame_stats =

       simple_encode.ObserveFirstPassStats();

@@ -34,16 +41,8 @@

 TEST(SimpleEncode, GetCodingFrameNum) {

-  int w = 352;

-  int h = 288;

-  int frame_rate_num = 30;

-  int frame_rate_den = 1;

-  int target_bitrate = 200;

-  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);

+                             target_bitrate, num_frames, infile_path);

   simple_encode.ComputeFirstPassStats();

   int num_coding_frames = simple_encode.GetCodingFrameNum();

   EXPECT_EQ(num_coding_frames, 19);

@@ -50,16 +49,8 @@

 TEST(SimpleEncode, EncodeFrame) {

-  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);

+                             target_bitrate, num_frames, infile_path);

   simple_encode.ComputeFirstPassStats();

   int num_coding_frames = simple_encode.GetCodingFrameNum();

   EXPECT_GE(num_coding_frames, num_frames);

@@ -68,6 +59,7 @@

   int ref_num_alternate_refereces = num_coding_frames - num_frames;

   int num_alternate_refereces = 0;

   simple_encode.StartEncode();

+  size_t total_data_bit_size = 0;

   for (int i = 0; i < num_coding_frames; ++i) {

     EncodeFrameResult encode_frame_result;

     if (i == 0) {

@@ -88,8 +80,28 @@

     EXPECT_GE(encode_frame_result.psnr, 34)

         << "The psnr is supposed to be greater than 34 given the "

            "target_bitrate 1000 kbps";

+    total_data_bit_size += encode_frame_result.coding_data_bit_size;

   EXPECT_EQ(num_alternate_refereces, ref_num_alternate_refereces);

+  double bitrate = get_bit_rate_in_kpbs(total_data_bit_size, num_frames,

+                                        frame_rate_num, frame_rate_den);

+  EXPECT_LE(fabs(target_bitrate - bitrate), 150);

+  simple_encode.EndEncode();

+}

+TEST(SimpleEncode, EncodeFrameWithQuantizeIndex) {

+  SimpleEncode simple_encode(w, h, frame_rate_num, frame_rate_den,

+                             target_bitrate, num_frames, infile_path);

+  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();

--- a/vp9/encoder/vp9_encoder.c

+++ b/vp9/encoder/vp9_encoder.c

@@ -2647,6 +2647,10 @@

   cm->error.setjmp = 0;

+#if CONFIG_RATE_CTRL

+  encode_command_init(&cpi->encode_command);

+#endif

   return cpi;

@@ -4282,6 +4286,14 @@

       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 @@

     // 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

--- a/vp9/encoder/vp9_encoder.h

+++ b/vp9/encoder/vp9_encoder.h

@@ -516,6 +516,31 @@

   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 @@

   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 {

--- a/vp9/simple_encode.cc

+++ b/vp9/simple_encode.cc

@@ -87,6 +87,8 @@

 static void update_encode_frame_result(

     EncodeFrameResult *encode_frame_result,

     const ENCODE_FRAME_RESULT *encode_frame_info) {

+  encode_frame_result->coding_data_bit_size =

+      encode_frame_result->coding_data_byte_size * 8;

   encode_frame_result->show_idx = encode_frame_info->show_idx;

   encode_frame_result->frame_type =

       get_frame_type_from_update_type(encode_frame_info->update_type);

@@ -97,7 +99,8 @@

 SimpleEncode::SimpleEncode(int frame_width, int frame_height,

                            int frame_rate_num, int frame_rate_den,

-                           int target_bitrate, int num_frames, FILE *file)

+                           int target_bitrate, int num_frames,

+                           const char *infile_path)

     : pimpl{ std::unique_ptr<impl>(new impl()) } {

   this->frame_width = frame_width;

   this->frame_height = frame_height;

@@ -105,7 +108,8 @@

   this->frame_rate_den = frame_rate_den;

   this->target_bitrate = target_bitrate;

   this->num_frames = num_frames;

-  this->file = file;

+  // TODO(angirbid): Should we keep a file pointer here or keep the file_path?

+  this->file = fopen(infile_path, "r");

   pimpl->cpi = NULL;

   pimpl->img_fmt = VPX_IMG_FMT_I420;

@@ -234,9 +238,9 @@

   assert(encode_frame_result->coding_data.get() == nullptr);

-  const size_t max_coding_data_size = frame_width * frame_height * 3;

-  encode_frame_result->coding_data =

-      std::move(std::unique_ptr<uint8_t[]>(new uint8_t[max_coding_data_size]));

+  const size_t max_coding_data_byte_size = frame_width * frame_height * 3;

+  encode_frame_result->coding_data = std::move(

+      std::unique_ptr<uint8_t[]>(new uint8_t[max_coding_data_byte_size]));

   int64_t time_stamp;

   int64_t time_end;

   int flush = 1;  // Make vp9_get_compressed_data encode a frame

@@ -243,15 +247,24 @@

   unsigned int frame_flags = 0;

   ENCODE_FRAME_RESULT encode_frame_info;

   vp9_get_compressed_data(cpi, &frame_flags,

-                          &encode_frame_result->coding_data_size,

+                          &encode_frame_result->coding_data_byte_size,

                           encode_frame_result->coding_data.get(), &time_stamp,

                           &time_end, flush, &encode_frame_info);

   // vp9_get_compressed_data is expected to encode a frame every time, so the

   // data size should be greater than zero.

-  assert(encode_frame_result->coding_data_size > 0);

-  assert(encode_frame_result->coding_data_size < max_coding_data_size);

+  assert(encode_frame_result->coding_data_byte_size > 0);

+  assert(encode_frame_result->coding_data_byte_size <

+         max_coding_data_byte_size);

   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() {

--- a/vp9/simple_encode.h

+++ b/vp9/simple_encode.h

@@ -10,7 +10,8 @@

 struct EncodeFrameResult {

   int show_idx;

   FrameType frame_type;

-  size_t coding_data_size;

+  size_t coding_data_bit_size;

+  size_t coding_data_byte_size;

   // The EncodeFrame will allocate a buffer, write the coding data into the

   // buffer and give the ownership of the buffer to coding_data

   std::unique_ptr<unsigned char[]> coding_data;

@@ -23,7 +24,7 @@

  public:

   SimpleEncode(int frame_width, int frame_height, int frame_rate_num,

                int frame_rate_den, int target_bitrate, int num_frames,

-               FILE *file);

+               const char *infile_path);

   ~SimpleEncode();

   SimpleEncode(SimpleEncode &&) = delete;

   SimpleEncode &operator=(SimpleEncode &&) = delete;

@@ -46,6 +47,11 @@

   // Encode a frame

   // 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.