Add codec control to export TPL stats

new codec control: VP9E_GET_TPL_STATS with unit test

Bug: b/273736974
Change-Id: I27343bd3f6dffafc86925234537bcdb557bc4079

diff --git a/test/encode_api_test.cc b/test/encode_api_test.cc
index ecdf928..eac9626 100644 (file)
--- a/test/encode_api_test.cc
+++ b/test/encode_api_test.cc
@@ -13,9 +13,12 @@
 #include <initializer_list>
 
 #include "third_party/googletest/src/include/gtest/gtest.h"
+#include "test/codec_factory.h"
+#include "test/encode_test_driver.h"
+#include "test/i420_video_source.h"
+#include "test/video_source.h"
 
 #include "./vpx_config.h"
-#include "test/video_source.h"
 #include "vpx/vp8cx.h"
 #include "vpx/vpx_encoder.h"
 
@@ -360,4 +363,79 @@ TEST(EncodeAPI, ConfigChangeThreadCount) {
   }
 }
 
+#if CONFIG_VP9_ENCODER
+class EncodeApiGetTplStatsTest
+    : public ::libvpx_test::EncoderTest,
+      public ::testing::TestWithParam<const libvpx_test::CodecFactory *> {
+ public:
+  EncodeApiGetTplStatsTest() : EncoderTest(GetParam()) {}
+  ~EncodeApiGetTplStatsTest() override {}
+
+ protected:
+  void SetUp() override {
+    InitializeConfig();
+    SetMode(::libvpx_test::kTwoPassGood);
+  }
+
+  void PreEncodeFrameHook(::libvpx_test::VideoSource *video,
+                          ::libvpx_test::Encoder *encoder) override {
+    if (video->frame() == 0) {
+      encoder->Control(VP9E_SET_TPL, 1);
+    }
+  }
+
+  vpx_codec_err_t AllocateTplList(TplFrameStats **data) {
+    // Allocate MAX_ARF_GOP_SIZE * sizeof(TplFrameStats) that will be filled
+    // by VP9E_GET_TPL_STATS
+    *data = static_cast<TplFrameStats *>(calloc(50, sizeof(TplFrameStats)));
+    if (*data == nullptr) return VPX_CODEC_MEM_ERROR;
+    return VPX_CODEC_OK;
+  }
+
+  void PostEncodeFrameHook(::libvpx_test::Encoder *encoder) override {
+    ::libvpx_test::CxDataIterator iter = encoder->GetCxData();
+    while (const vpx_codec_cx_pkt_t *pkt = iter.Next()) {
+      switch (pkt->kind) {
+        case VPX_CODEC_CX_FRAME_PKT: {
+          TplFrameStats *tpl_stats = NULL;
+          EXPECT_EQ(AllocateTplList(&tpl_stats), VPX_CODEC_OK);
+          encoder->Control(VP9E_GET_TPL_STATS, tpl_stats);
+          bool stats_not_all_zero = false;
+          for (unsigned int i = 0; i < cfg_.g_lag_in_frames; i++) {
+            if (tpl_stats[i].frame_width != 0) {
+              ASSERT_EQ(tpl_stats[i].frame_width, width_);
+              ASSERT_EQ(tpl_stats[i].frame_height, height_);
+              ASSERT_NE(tpl_stats[i].block_stats_list, nullptr);
+              stats_not_all_zero = true;
+            }
+          }
+          ASSERT_TRUE(stats_not_all_zero);
+          // Free the memory right away now as this is only a test.
+          free(tpl_stats);
+          break;
+        }
+        default: break;
+      }
+    }
+  }
+
+  int width_;
+  int height_;
+};
+
+TEST_P(EncodeApiGetTplStatsTest, GetTplStats) {
+  cfg_.g_lag_in_frames = 25;
+  width_ = 352;
+  height_ = 288;
+  ::libvpx_test::I420VideoSource video("hantro_collage_w352h288.yuv", width_,
+                                       height_, 30, 1, 0, 150);
+  ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    VP9, EncodeApiGetTplStatsTest,
+    ::testing::Values(
+        static_cast<const libvpx_test::CodecFactory *>(&libvpx_test::kVP9)));
+#endif  // CONFIG_VP9_ENCODER
+
 }  // namespace

diff --git a/test/encode_test_driver.h b/test/encode_test_driver.h
index b57df85..27a78e6 100644 (file)
--- a/test/encode_test_driver.h
+++ b/test/encode_test_driver.h
@@ -153,6 +153,11 @@ class Encoder {
     const vpx_codec_err_t res = vpx_codec_control_(&encoder_, ctrl_id, arg);
     ASSERT_EQ(VPX_CODEC_OK, res) << EncoderError();
   }
+
+  void Control(int ctrl_id, TplFrameStats *arg) {
+    const vpx_codec_err_t res = vpx_codec_control_(&encoder_, ctrl_id, arg);
+    ASSERT_EQ(VPX_CODEC_OK, res) << EncoderError();
+  }
 #endif  // CONFIG_VP9_ENCODER
 
 #if CONFIG_VP8_ENCODER || CONFIG_VP9_ENCODER

diff --git a/vp9/encoder/vp9_encoder.h b/vp9/encoder/vp9_encoder.h
index 7c22c80..742cf0b 100644 (file)
--- a/vp9/encoder/vp9_encoder.h
+++ b/vp9/encoder/vp9_encoder.h
@@ -505,6 +505,7 @@ typedef struct EncFrameBuf {
 } EncFrameBuf;
 
 // Maximum operating frame buffer size needed for a GOP using ARF reference.
+// This is used to allocate the memory for TPL stats for a GOP.
 #define MAX_ARF_GOP_SIZE (2 * MAX_LAG_BUFFERS)
 #define MAX_KMEANS_GROUPS 8
 

diff --git a/vp9/encoder/vp9_tpl_model.c b/vp9/encoder/vp9_tpl_model.c
index dbd7482..b62c66b 100644 (file)
--- a/vp9/encoder/vp9_tpl_model.c
+++ b/vp9/encoder/vp9_tpl_model.c
@@ -1134,6 +1134,8 @@ static void mc_flow_dispenser(VP9_COMP *cpi, GF_PICTURE *gf_picture,
   const int mi_height = num_8x8_blocks_high_lookup[bsize];
   const int mi_width = num_8x8_blocks_wide_lookup[bsize];
 
+  tpl_frame_stats_before_propagation->frame_width = cm->width;
+  tpl_frame_stats_before_propagation->frame_height = cm->height;
   // Setup scaling factor
 #if CONFIG_VP9_HIGHBITDEPTH
   vp9_setup_scale_factors_for_frame(

diff --git a/vp9/vp9_cx_iface.c b/vp9/vp9_cx_iface.c
index 4c7eaed..8bd880c 100644 (file)
--- a/vp9/vp9_cx_iface.c
+++ b/vp9/vp9_cx_iface.c
@@ -1788,6 +1788,23 @@ static vpx_codec_err_t ctrl_get_svc_ref_frame_config(vpx_codec_alg_priv_t *ctx,
   return VPX_CODEC_OK;
 }
 
+static vpx_codec_err_t ctrl_get_tpl_stats(vpx_codec_alg_priv_t *ctx,
+                                          va_list args) {
+  VP9_COMP *const cpi = ctx->cpi;
+  TplFrameStats *data = va_arg(args, TplFrameStats *);
+  int i;
+  if (data == NULL) {
+    return VPX_CODEC_INVALID_PARAM;
+  }
+  for (i = 0; i < MAX_ARF_GOP_SIZE; i++) {
+    data[i].frame_width = cpi->tpl_frame_stats[i].frame_width;
+    data[i].frame_height = cpi->tpl_frame_stats[i].frame_height;
+    data[i].block_stats_list = cpi->tpl_frame_stats[i].block_stats_list;
+  }
+
+  return VPX_CODEC_OK;
+}
+
 static vpx_codec_err_t ctrl_set_svc_ref_frame_config(vpx_codec_alg_priv_t *ctx,
                                                      va_list args) {
   VP9_COMP *const cpi = ctx->cpi;
@@ -2035,6 +2052,7 @@ static vpx_codec_ctrl_fn_map_t encoder_ctrl_maps[] = {
   { VP9E_GET_ACTIVEMAP, ctrl_get_active_map },
   { VP9E_GET_LEVEL, ctrl_get_level },
   { VP9E_GET_SVC_REF_FRAME_CONFIG, ctrl_get_svc_ref_frame_config },
+  { VP9E_GET_TPL_STATS, ctrl_get_tpl_stats },
 
   { -1, NULL },
 };

diff --git a/vpx/vp8cx.h b/vpx/vp8cx.h
index e0b679f..123a645 100644 (file)
--- a/vpx/vp8cx.h
+++ b/vpx/vp8cx.h
@@ -767,6 +767,18 @@ enum vp8e_enc_control_id {
    *
    */
   VP9E_SET_QUANTIZER_ONE_PASS,
+
+  /*!\brief Codec control to get TPL stats for the current GOP.
+   *
+   * Allocation and free of memory of size MAX_ARF_GOP_SIZE (50) *
+   * sizeof(TplFrameStats) should be done by applications.
+   *
+   * VPX_CODEC_INVALID_PARAM will be returned if the pointer passed in is NULL.
+   *
+   * Supported in codecs: VP9
+   *
+   */
+  VP9E_GET_TPL_STATS,
 };
 
 /*!\brief vpx 1-D scaling mode
@@ -1097,6 +1109,8 @@ VPX_CTRL_USE_TYPE(VP8E_SET_RTC_EXTERNAL_RATECTRL, int)
 #define VPX_CTRL_VP8E_SET_RTC_EXTERNAL_RATECTRL
 VPX_CTRL_USE_TYPE(VP9E_SET_QUANTIZER_ONE_PASS, int)
 #define VPX_CTRL_VP9E_SET_QUANTIZER_ONE_PASS
+VPX_CTRL_USE_TYPE(VP9E_GET_TPL_STATS, void *)
+#define VPX_CTRL_VP9E_GET_TPL_STATS
 
 /*!\endcond */
 /*! @} - end defgroup vp8_encoder */

diff --git a/vpx/vpx_encoder.h b/vpx/vpx_encoder.h
index 9247231..1910a19 100644 (file)
--- a/vpx/vpx_encoder.h
+++ b/vpx/vpx_encoder.h
@@ -58,7 +58,7 @@ extern "C" {
  * fields to structures
  */
 #define VPX_ENCODER_ABI_VERSION \
-  (15 + VPX_CODEC_ABI_VERSION + \
+  (16 + VPX_CODEC_ABI_VERSION + \
    VPX_EXT_RATECTRL_ABI_VERSION) /**<\hideinitializer*/
 
 /*! \brief Encoder capabilities bitfield