Pass motion vector info to encode frame result

Pass the motion vector info stored to the encode frame result
through the interface "update_encode_frame_result()".

Change-Id: I589affa0c4c4d0fd4d639edff9068e44a715beff

3 files changed, 76 insertions, 1 deletions

diff --git a/vp9/encoder/vp9_encoder.c b/vp9/encoder/vp9_encoder.c
index 49514f8e9..bcad5ad31 100644
--- a/vp9/encoder/vp9_encoder.c
+++ b/vp9/encoder/vp9_encoder.c
@@ -7257,6 +7257,7 @@ static void update_encode_frame_result(
     uint32_t bit_depth, uint32_t input_bit_depth, const FRAME_COUNTS *counts,
 #if CONFIG_RATE_CTRL
     const PARTITION_INFO *partition_info,
+    const MOTION_VECTOR_INFO *motion_vector_info,
 #endif  // CONFIG_RATE_CTRL
     ENCODE_FRAME_RESULT *encode_frame_result) {
 #if CONFIG_RATE_CTRL
@@ -7273,6 +7274,7 @@ static void update_encode_frame_result(
   encode_frame_result->sse = psnr.sse[0];
   copy_frame_counts(counts, &encode_frame_result->frame_counts);
   encode_frame_result->partition_info = partition_info;
+  encode_frame_result->motion_vector_info = motion_vector_info;
   if (encode_frame_result->coded_frame.allocated) {
     yv12_buffer_to_image_buffer(coded_frame, &encode_frame_result->coded_frame);
   }
@@ -7593,7 +7595,7 @@ int vp9_get_compressed_data(VP9_COMP *cpi, unsigned int *frame_flags,
         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,
+        cpi->partition_info, cpi->motion_vector_info,
 #endif  // CONFIG_RATE_CTRL
         encode_frame_result);
     vp9_twopass_postencode_update(cpi);
diff --git a/vp9/simple_encode.cc b/vp9/simple_encode.cc
index 5229c3d8c..fded3befc 100644
--- a/vp9/simple_encode.cc
+++ b/vp9/simple_encode.cc
@@ -157,6 +157,35 @@ static void update_partition_info(const PARTITION_INFO *input_partition_info,
   }
 }
 
+static void update_motion_vector_info(
+    const MOTION_VECTOR_INFO *input_motion_vector_info,
+    const FrameType frame_type, const int num_rows_4x4, const int num_cols_4x4,
+    MotionVectorInfo *output_motion_vector_info) {
+  const int num_units_4x4 = num_rows_4x4 * num_cols_4x4;
+  for (int i = 0; i < num_units_4x4; ++i) {
+    output_motion_vector_info[i].mv_count =
+        (frame_type == kKeyFrame)
+            ? 0
+            : ((input_motion_vector_info[i].ref_frame[1] == -1) ? 1 : 2);
+    output_motion_vector_info[i].ref_frame[0] =
+        static_cast<RefFrameType>(input_motion_vector_info[i].ref_frame[0]);
+    output_motion_vector_info[i].ref_frame[1] =
+        static_cast<RefFrameType>(input_motion_vector_info[i].ref_frame[1]);
+    output_motion_vector_info[i].mv_row[0] =
+        (double)input_motion_vector_info[i].mv[0].as_mv.row /
+        kMotionVectorPrecision;
+    output_motion_vector_info[i].mv_column[0] =
+        (double)input_motion_vector_info[i].mv[0].as_mv.col /
+        kMotionVectorPrecision;
+    output_motion_vector_info[i].mv_row[1] =
+        (double)input_motion_vector_info[i].mv[1].as_mv.row /
+        kMotionVectorPrecision;
+    output_motion_vector_info[i].mv_column[1] =
+        (double)input_motion_vector_info[i].mv[1].as_mv.col /
+        kMotionVectorPrecision;
+  }
+}
+
 static void update_frame_counts(const FRAME_COUNTS *input_counts,
                                 FrameCounts *output_counts) {
   // Init array sizes.
@@ -434,6 +463,8 @@ static bool init_encode_frame_result(EncodeFrameResult *encode_frame_result,
   encode_frame_result->num_cols_4x4 = get_num_unit_4x4(frame_height);
   encode_frame_result->partition_info.resize(encode_frame_result->num_rows_4x4 *
                                              encode_frame_result->num_cols_4x4);
+  encode_frame_result->motion_vector_info.resize(
+      encode_frame_result->num_rows_4x4 * encode_frame_result->num_cols_4x4);
 
   if (encode_frame_result->coding_data.get() == nullptr) {
     return false;
@@ -457,6 +488,10 @@ static void update_encode_frame_result(
                         encode_frame_result->num_rows_4x4,
                         encode_frame_result->num_cols_4x4,
                         &encode_frame_result->partition_info[0]);
+  update_motion_vector_info(
+      encode_frame_info->motion_vector_info, encode_frame_result->frame_type,
+      encode_frame_result->num_rows_4x4, encode_frame_result->num_cols_4x4,
+      &encode_frame_result->motion_vector_info[0]);
   update_frame_counts(&encode_frame_info->frame_counts,
                       &encode_frame_result->frame_counts);
 }
diff --git a/vp9/simple_encode.h b/vp9/simple_encode.h
index 5d8f87676..03c361deb 100644
--- a/vp9/simple_encode.h
+++ b/vp9/simple_encode.h
@@ -25,6 +25,15 @@ enum FrameType {
   kAlternateReference,
 };
 
+// The enum type is similar to vp9: |MV_REFERENCE_FRAME|.
+enum RefFrameType {
+  kIntraFrame = 0,
+  kLastFrame = 1,
+  kGoldenFrame = 2,
+  kAltRefFrame = 3,
+  kNoneRefFrame = -1,
+};
+
 // The frame is split to 4x4 blocks.
 // This structure contains the information of each 4x4 block.
 struct PartitionInfo {
@@ -36,6 +45,25 @@ struct PartitionInfo {
   int height;        // prediction block height
 };
 
+constexpr int kMotionVectorPrecision = 8;
+
+// The frame is split to 4x4 blocks.
+// This structure contains the information of each 4x4 block.
+struct MotionVectorInfo {
+  // Number of valid motion vectors, always 0 if this block is in the key frame.
+  // For inter frames, it could be 1 or 2.
+  int mv_count;
+  // The reference frame for motion vectors. If the second motion vector does
+  // not exist (mv_count = 1), the reference frame is kNoneRefFrame.
+  // Otherwise, the reference frame is either kLastFrame, or kGoldenFrame,
+  // or kAltRefFrame.
+  RefFrameType ref_frame[2];
+  // The row offset of motion vectors in the unit of pixel.
+  double mv_row[2];
+  // The column offset of motion vectors in the unit of pixel.
+  double mv_column[2];
+};
+
 struct EncodeFrameInfo {
   int show_idx;
   FrameType frame_type;
@@ -175,6 +203,16 @@ struct EncodeFrameResult {
   // Horizontal next: |column_start| + |width|,
   // Vertical next: |row_start| + |height|.
   std::vector<PartitionInfo> partition_info;
+  // A vector of the motion vector information of the frame.
+  // The number of elements is |num_rows_4x4| * |num_cols_4x4|.
+  // The frame is divided 4x4 blocks of |num_rows_4x4| rows and
+  // |num_cols_4x4| columns.
+  // Each 4x4 block contains 0 motion vector if this is an intra predicted
+  // frame (for example, the key frame). If the frame is inter predicted,
+  // each 4x4 block contains either 1 or 2 motion vectors.
+  // Similar to partition info, all 4x4 blocks inside the same partition block
+  // share the same motion vector information.
+  std::vector<MotionVectorInfo> motion_vector_info;
   ImageBuffer coded_frame;
 };