Pass partition info to encode frame result

Init the memory for partition information in "EncodeFrameResult".
And pass the partition information of vp9 encoder to it through
the interface: "update_encode_frame_result()".

Change-Id: Iea049e661da79f54d41da7924b9ef28ff7cfbfa3

vp9/encoder/vp9_encoder.c

@@ -7229,6 +7229,9 @@ static void update_encode_frame_result(
     const YV12_BUFFER_CONFIG *source_frame,
     const YV12_BUFFER_CONFIG *coded_frame, int quantize_index,
     uint32_t bit_depth, uint32_t input_bit_depth, const FRAME_COUNTS *counts,
+#if CONFIG_RATE_CTRL
+    const PARTITION_INFO *partition_info,
+#endif  // CONFIG_RATE_CTRL
     ENCODE_FRAME_RESULT *encode_frame_result) {
 #if CONFIG_RATE_CTRL
   PSNR_STATS psnr;
@@ -7243,6 +7246,7 @@ static void update_encode_frame_result(
   encode_frame_result->psnr = psnr.psnr[0];
   encode_frame_result->sse = psnr.sse[0];
   copy_frame_counts(counts, &encode_frame_result->frame_counts);
+  encode_frame_result->partition_info = partition_info;
 #else   // CONFIG_RATE_CTRL
   (void)bit_depth;
   (void)input_bit_depth;
@@ -7551,6 +7555,9 @@ int vp9_get_compressed_data(VP9_COMP *cpi, unsigned int *frame_flags,
         cpi->twopass.gf_group.update_type[cpi->twopass.gf_group.index],
         cpi->Source, get_frame_new_buffer(cm), vp9_get_quantizer(cpi),
         cpi->oxcf.input_bit_depth, cm->bit_depth, cpi->td.counts,
+#if CONFIG_RATE_CTRL
+        cpi->partition_info,
+#endif  // CONFIG_RATE_CTRL
         encode_frame_result);
     vp9_twopass_postencode_update(cpi);
   } else if (cpi->use_svc) {

vp9/encoder/vp9_encoder.h

@@ -895,6 +895,7 @@ typedef struct ENCODE_FRAME_RESULT {
   double psnr;
   uint64_t sse;
   FRAME_COUNTS frame_counts;
+  const PARTITION_INFO *partition_info;
 #endif  // CONFIG_RATE_CTRL
   int quantize_index;
 } ENCODE_FRAME_RESULT;

vp9/simple_encode.cc

@@ -8,6 +8,7 @@
  *  be found in the AUTHORS file in the root of the source tree.
  */
 
+#include <memory>
 #include <vector>
 #include "vp9/common/vp9_entropymode.h"
 #include "vp9/common/vp9_enums.h"
@@ -101,6 +102,22 @@ get_frame_type_from_update_type(FRAME_UPDATE_TYPE update_type) {
   }
 }
 
+static void update_partition_info(const PARTITION_INFO *input_partition_info,
+                                  const int num_rows_4x4,
+                                  const int num_cols_4x4,
+                                  PartitionInfo *output_partition_info) {
+  const int num_units_4x4 = num_rows_4x4 * num_cols_4x4;
+  for (int i = 0; i < num_units_4x4; ++i) {
+    output_partition_info[i].row = input_partition_info[i].row;
+    output_partition_info[i].column = input_partition_info[i].column;
+    output_partition_info[i].row_start = input_partition_info[i].row_start;
+    output_partition_info[i].column_start =
+        input_partition_info[i].column_start;
+    output_partition_info[i].width = input_partition_info[i].width;
+    output_partition_info[i].height = input_partition_info[i].height;
+  }
+}
+
 static void update_frame_counts(const FRAME_COUNTS *input_counts,
                                 FrameCounts *output_counts) {
   // Init array sizes.
@@ -333,6 +350,10 @@ static void update_encode_frame_result(
   encode_frame_result->psnr = encode_frame_info->psnr;
   encode_frame_result->sse = encode_frame_info->sse;
   encode_frame_result->quantize_index = encode_frame_info->quantize_index;
+  update_partition_info(encode_frame_info->partition_info,
+                        encode_frame_result->num_rows_4x4,
+                        encode_frame_result->num_cols_4x4,
+                        encode_frame_result->partition_info.get());
   update_frame_counts(&encode_frame_info->frame_counts,
                       &encode_frame_result->frame_counts);
 }
@@ -405,6 +426,8 @@ SimpleEncode::SimpleEncode(int frame_width, int frame_height,
   impl_ptr_ = std::unique_ptr<EncodeImpl>(new EncodeImpl());
   frame_width_ = frame_width;
   frame_height_ = frame_height;
+  num_rows_4x4_ = get_num_unit_4x4(frame_width);
+  num_cols_4x4_ = get_num_unit_4x4(frame_height);
   frame_rate_num_ = frame_rate_num;
   frame_rate_den_ = frame_rate_den;
   target_bitrate_ = target_bitrate;
@@ -566,6 +589,11 @@ void SimpleEncode::EncodeFrame(EncodeFrameResult *encode_frame_result) {
   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]));
+  encode_frame_result->num_rows_4x4 = num_rows_4x4_;
+  encode_frame_result->num_cols_4x4 = num_cols_4x4_;
+  encode_frame_result->partition_info =
+      std::move(std::unique_ptr<PartitionInfo[]>(
+          new PartitionInfo[num_rows_4x4_ * num_cols_4x4_]));
   int64_t time_stamp;
   int64_t time_end;
   int flush = 1;  // Make vp9_get_compressed_data encode a frame

vp9/simple_encode.h

@@ -25,6 +25,17 @@ enum FrameType {
   kAlternateReference,
 };
 
+// The frame is split to 4x4 blocks.
+// This structure contains the information of each 4x4 block.
+struct PartitionInfo {
+  int row;           // row pixel offset of current 4x4 block
+  int column;        // column pixel offset of current 4x4 block
+  int row_start;     // row pixel offset of the start of the prediction block
+  int column_start;  // column pixel offset of the start of the prediction block
+  int width;         // prediction block width
+  int height;        // prediction block height
+};
+
 struct EncodeFrameInfo {
   int show_idx;
   FrameType frame_type;
@@ -126,6 +137,10 @@ struct EncodeFrameResult {
   uint64_t sse;
   int quantize_index;
   FrameCounts frame_counts;
+  int num_rows_4x4;  // number of row units, in size of 4.
+  int num_cols_4x4;  // number of column units, in size of 4.
+  // The pointer to the partition information of the frame.
+  std::unique_ptr<PartitionInfo[]> partition_info;
 };
 
 struct GroupOfPicture {
@@ -212,8 +227,10 @@ class SimpleEncode {
  private:
   class EncodeImpl;
 
-  int frame_width_;
+  int frame_width_;   // frame width in pixels.
-  int frame_height_;
+  int frame_height_;  // frame height in pixels.
+  int num_rows_4x4_;  // number of row units, in size of 4.
+  int num_cols_4x4_;  // number of column units, in size of 4.
   int frame_rate_num_;
   int frame_rate_den_;
   int target_bitrate_;