ExtPart: unit tests for whole partition tree
Add a unit test that compares psnr values to two encoding runs.
The first run is the baseline encoder, where we restrict the max
and min coding blocks to 4x4.
In the second run, we read partition decisions from the a toy model,
which provides the whole partition decision tree for the superblock.
And the decision tree is exactly all splits down to 4x4 block.
We expect these two runs to return the same psnr value.
Change-Id: Iec81f83164ad6327aa41278dc9822b91f11610a9
diff --git a/aom/aom_external_partition.h b/aom/aom_external_partition.h
index 0a14260..4d50833 100644
--- a/aom/aom_external_partition.h
+++ b/aom/aom_external_partition.h
@@ -218,6 +218,10 @@
// Features for a new ML model.
aom_sb_features_t sb_features; /**< Features collected for the super block */
+ int mi_row; /**< Mi_row position of the block */
+ int mi_col; /**< Mi_col position of the block */
+ int frame_width; /**< Frame width */
+ int frame_height; /**< Frame height */
} aom_partition_features_t;
/*!\brief Partition decisions received from the external model.
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 60f8e50..0d530b2 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -683,10 +683,17 @@
if (num_passes == 1) {
#if CONFIG_PARTITION_SEARCH_ORDER
- av1_reset_part_sf(&cpi->sf.part_sf);
- RD_STATS this_rdc;
- av1_rd_partition_search(cpi, td, tile_data, tp, sms_root, mi_row, mi_col,
- sb_size, &this_rdc);
+ if (cpi->ext_part_controller.ready && !frame_is_intra_only(cm)) {
+ av1_reset_part_sf(&cpi->sf.part_sf);
+ RD_STATS this_rdc;
+ av1_rd_partition_search(cpi, td, tile_data, tp, sms_root, mi_row,
+ mi_col, sb_size, &this_rdc);
+ } else {
+ PC_TREE *const pc_root = av1_alloc_pc_tree_node(sb_size);
+ av1_rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, sb_size,
+ &dummy_rdc, dummy_rdc, pc_root, sms_root, NULL,
+ SB_SINGLE_PASS, NULL);
+ }
#else
PC_TREE *const pc_root = av1_alloc_pc_tree_node(sb_size);
av1_rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, sb_size,
diff --git a/av1/encoder/partition_search.c b/av1/encoder/partition_search.c
index 9958f81..28a40f4 100644
--- a/av1/encoder/partition_search.c
+++ b/av1/encoder/partition_search.c
@@ -4077,6 +4077,11 @@
static void prepare_sb_features_before_search(
AV1_COMP *const cpi, ThreadData *td, int mi_row, int mi_col,
const BLOCK_SIZE bsize, aom_partition_features_t *features) {
+ // TODO(chengchen): properly handle feature collection for unit tests.
+ // Also take care of cases where tpl stats are not available.
+ // Now in unit test, this function causes failures, due to tpl stats
+ // not ready.
+ return;
av1_collect_motion_search_features_sb(cpi, td, mi_row, mi_col, bsize,
features);
collect_tpl_stats_sb(cpi, bsize, mi_row, mi_col, features);
@@ -4130,6 +4135,10 @@
ExtPartController *const ext_part_controller = &cpi->ext_part_controller;
aom_partition_features_t features;
prepare_sb_features_before_search(cpi, td, mi_row, mi_col, bsize, &features);
+ features.mi_row = mi_row;
+ features.mi_col = mi_col;
+ features.frame_width = cpi->frame_info.frame_width;
+ features.frame_height = cpi->frame_info.frame_height;
av1_ext_part_send_features(ext_part_controller, &features);
PC_TREE *pc_tree;
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index 146d4dd..090a9c3 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -2296,6 +2296,8 @@
const BLOCK_SIZE bsize,
aom_partition_features_t *features) {
const AV1_COMMON *const cm = &cpi->common;
+ if (frame_is_intra_only(cm)) return;
+
MACROBLOCK *const x = &td->mb;
const BLOCK_SIZE fixed_block_size = BLOCK_16X16;
const int col_step = mi_size_wide[fixed_block_size];
diff --git a/test/av1_external_partition_test.cc b/test/av1_external_partition_test.cc
index 96e5c66..45ca847a 100644
--- a/test/av1_external_partition_test.cc
+++ b/test/av1_external_partition_test.cc
@@ -38,8 +38,281 @@
TestData *data;
aom_ext_part_config_t config;
aom_ext_part_funcs_t funcs;
+ int mi_row;
+ int mi_col;
+ int frame_width;
+ int frame_height;
} ToyModel;
+// Note:
+// if CONFIG_PARTITION_SEARCH_ORDER = 0, we test APIs designed for the baseline
+// encoder's DFS partition search workflow.
+// if CONFIG_PARTITION_SEARCH_ORDER = 1, we test APIs designed for the new
+// ML model's partition search workflow.
+#if CONFIG_PARTITION_SEARCH_ORDER
+aom_ext_part_status_t ext_part_create_model(
+ void *priv, const aom_ext_part_config_t *part_config,
+ aom_ext_part_model_t *ext_part_model) {
+ TestData *received_data = reinterpret_cast<TestData *>(priv);
+ EXPECT_EQ(received_data->version, kVersion);
+ ToyModel *toy_model = new (std::nothrow) ToyModel;
+ EXPECT_NE(toy_model, nullptr);
+ toy_model->data = received_data;
+ *ext_part_model = toy_model;
+ EXPECT_EQ(part_config->superblock_size, BLOCK_64X64);
+ return AOM_EXT_PART_OK;
+}
+
+aom_ext_part_status_t ext_part_send_features(
+ aom_ext_part_model_t ext_part_model,
+ const aom_partition_features_t *part_features) {
+ ToyModel *toy_model = static_cast<ToyModel *>(ext_part_model);
+ toy_model->mi_row = part_features->mi_row;
+ toy_model->mi_col = part_features->mi_col;
+ toy_model->frame_width = part_features->frame_width;
+ toy_model->frame_height = part_features->frame_height;
+ return AOM_EXT_PART_OK;
+}
+
+// The model provide the whole decision tree to the encoder.
+aom_ext_part_status_t ext_part_get_partition_decision_whole_tree(
+ aom_ext_part_model_t ext_part_model,
+ aom_partition_decision_t *ext_part_decision) {
+ ToyModel *toy_model = static_cast<ToyModel *>(ext_part_model);
+ // A toy model that always asks the encoder to encode with
+ // 4x4 blocks (the smallest).
+ ext_part_decision->is_final_decision = 1;
+ // Note: super block size is fixed to BLOCK_64X64 for the
+ // input video. It is determined inside the encoder, see the
+ // check in "ext_part_create_model".
+ const int is_last_sb_col =
+ toy_model->mi_col * 4 + 64 > toy_model->frame_width;
+ const int is_last_sb_row =
+ toy_model->mi_row * 4 + 64 > toy_model->frame_height;
+ if (is_last_sb_row && is_last_sb_col) {
+ // 64x64: 1 node
+ // 32x32: 4 nodes (only the first one will further split)
+ // 16x16: 4 nodes
+ // 8x8: 4 * 4 nodes
+ // 4x4: 4 * 4 * 4 nodes
+ const int num_blocks = 1 + 4 + 4 + 4 * 4 + 4 * 4 * 4;
+ const int num_4x4_blocks = 4 * 4 * 4;
+ ext_part_decision->num_nodes = num_blocks;
+ // 64x64
+ ext_part_decision->partition_decision[0] = PARTITION_SPLIT;
+ // 32x32, only the first one will split, the other three are
+ // out of frame boundary.
+ ext_part_decision->partition_decision[1] = PARTITION_SPLIT;
+ ext_part_decision->partition_decision[2] = PARTITION_NONE;
+ ext_part_decision->partition_decision[3] = PARTITION_NONE;
+ ext_part_decision->partition_decision[4] = PARTITION_NONE;
+ // The rest blocks inside the top-left 32x32 block.
+ for (int i = 5; i < num_blocks - num_4x4_blocks; ++i) {
+ ext_part_decision->partition_decision[0] = PARTITION_SPLIT;
+ }
+ for (int i = num_blocks - num_4x4_blocks; i < num_blocks; ++i) {
+ ext_part_decision->partition_decision[i] = PARTITION_NONE;
+ }
+ } else if (is_last_sb_row) {
+ // 64x64: 1 node
+ // 32x32: 4 nodes (only the first two will further split)
+ // 16x16: 2 * 4 nodes
+ // 8x8: 2 * 4 * 4 nodes
+ // 4x4: 2 * 4 * 4 * 4 nodes
+ const int num_blocks = 1 + 4 + 2 * 4 + 2 * 4 * 4 + 2 * 4 * 4 * 4;
+ const int num_4x4_blocks = 2 * 4 * 4 * 4;
+ ext_part_decision->num_nodes = num_blocks;
+ // 64x64
+ ext_part_decision->partition_decision[0] = PARTITION_SPLIT;
+ // 32x32, only the first two will split, the other two are out
+ // of frame boundary.
+ ext_part_decision->partition_decision[1] = PARTITION_SPLIT;
+ ext_part_decision->partition_decision[2] = PARTITION_SPLIT;
+ ext_part_decision->partition_decision[3] = PARTITION_NONE;
+ ext_part_decision->partition_decision[4] = PARTITION_NONE;
+ // The rest blocks.
+ for (int i = 5; i < num_blocks - num_4x4_blocks; ++i) {
+ ext_part_decision->partition_decision[0] = PARTITION_SPLIT;
+ }
+ for (int i = num_blocks - num_4x4_blocks; i < num_blocks; ++i) {
+ ext_part_decision->partition_decision[i] = PARTITION_NONE;
+ }
+ } else if (is_last_sb_col) {
+ // 64x64: 1 node
+ // 32x32: 4 nodes (only the top-left and bottom-left will further split)
+ // 16x16: 2 * 4 nodes
+ // 8x8: 2 * 4 * 4 nodes
+ // 4x4: 2 * 4 * 4 * 4 nodes
+ const int num_blocks = 1 + 4 + 2 * 4 + 2 * 4 * 4 + 2 * 4 * 4 * 4;
+ const int num_4x4_blocks = 2 * 4 * 4 * 4;
+ ext_part_decision->num_nodes = num_blocks;
+ // 64x64
+ ext_part_decision->partition_decision[0] = PARTITION_SPLIT;
+ // 32x32, only the top-left and bottom-left will split, the other two are
+ // out of frame boundary.
+ ext_part_decision->partition_decision[1] = PARTITION_SPLIT;
+ ext_part_decision->partition_decision[2] = PARTITION_NONE;
+ ext_part_decision->partition_decision[3] = PARTITION_SPLIT;
+ ext_part_decision->partition_decision[4] = PARTITION_NONE;
+ // The rest blocks.
+ for (int i = 5; i < num_blocks - num_4x4_blocks; ++i) {
+ ext_part_decision->partition_decision[0] = PARTITION_SPLIT;
+ }
+ for (int i = num_blocks - num_4x4_blocks; i < num_blocks; ++i) {
+ ext_part_decision->partition_decision[i] = PARTITION_NONE;
+ }
+ } else {
+ // 64x64: 1 node
+ // 32x32: 4 nodes
+ // 16x16: 4 * 4 nodes
+ // 8x8: 4 * 4 * 4 nodes
+ // 4x4: 4 * 4 * 4 * 4 nodes
+ const int num_blocks = 1 + 4 + 4 * 4 + 4 * 4 * 4 + 4 * 4 * 4 * 4;
+ const int num_4x4_blocks = 4 * 4 * 4 * 4;
+ ext_part_decision->num_nodes = num_blocks;
+ for (int i = 0; i < num_blocks - num_4x4_blocks; ++i) {
+ ext_part_decision->partition_decision[i] = PARTITION_SPLIT;
+ }
+ for (int i = num_blocks - num_4x4_blocks; i < num_blocks; ++i) {
+ ext_part_decision->partition_decision[i] = PARTITION_NONE;
+ }
+ }
+
+ return AOM_EXT_PART_OK;
+}
+
+aom_ext_part_status_t ext_part_send_partition_stats(
+ aom_ext_part_model_t ext_part_model,
+ const aom_partition_stats_t *ext_part_stats) {
+ (void)ext_part_model;
+ (void)ext_part_stats;
+ return AOM_EXT_PART_OK;
+}
+
+aom_ext_part_status_t ext_part_delete_model(
+ aom_ext_part_model_t ext_part_model) {
+ ToyModel *toy_model = static_cast<ToyModel *>(ext_part_model);
+ EXPECT_EQ(toy_model->data->version, kVersion);
+ delete toy_model;
+ return AOM_EXT_PART_OK;
+}
+
+class ExternalPartitionTestAPI
+ : public ::libaom_test::CodecTestWith2Params<libaom_test::TestMode, int>,
+ public ::libaom_test::EncoderTest {
+ protected:
+ ExternalPartitionTestAPI()
+ : EncoderTest(GET_PARAM(0)), encoding_mode_(GET_PARAM(1)),
+ cpu_used_(GET_PARAM(2)), psnr_(0.0), nframes_(0) {}
+ virtual ~ExternalPartitionTestAPI() {}
+
+ virtual void SetUp() {
+ InitializeConfig(encoding_mode_);
+ const aom_rational timebase = { 1, 30 };
+ cfg_.g_timebase = timebase;
+ cfg_.rc_end_usage = AOM_VBR;
+ cfg_.g_threads = 1;
+ cfg_.g_lag_in_frames = 4;
+ cfg_.rc_target_bitrate = 400;
+ init_flags_ = AOM_CODEC_USE_PSNR;
+ }
+
+ virtual bool DoDecode() const { return false; }
+
+ virtual void BeginPassHook(unsigned int) {
+ psnr_ = 0.0;
+ nframes_ = 0;
+ }
+
+ virtual void PSNRPktHook(const aom_codec_cx_pkt_t *pkt) {
+ psnr_ += pkt->data.psnr.psnr[0];
+ nframes_++;
+ }
+
+ double GetAveragePsnr() const {
+ if (nframes_) return psnr_ / nframes_;
+ return 0.0;
+ }
+
+ void SetExternalPartition(bool use_external_partition) {
+ use_external_partition_ = use_external_partition;
+ }
+
+ void SetPartitionControlMode(int mode) { partition_control_mode_ = mode; }
+
+ virtual void PreEncodeFrameHook(::libaom_test::VideoSource *video,
+ ::libaom_test::Encoder *encoder) {
+ if (video->frame() == 0) {
+ aom_ext_part_funcs_t ext_part_funcs;
+ ext_part_funcs.priv = reinterpret_cast<void *>(&test_data_);
+ ext_part_funcs.create_model = ext_part_create_model;
+ ext_part_funcs.send_features = ext_part_send_features;
+ ext_part_funcs.get_partition_decision =
+ ext_part_get_partition_decision_whole_tree;
+ ext_part_funcs.send_partition_stats = ext_part_send_partition_stats;
+ ext_part_funcs.delete_model = ext_part_delete_model;
+
+ encoder->Control(AOME_SET_CPUUSED, cpu_used_);
+ encoder->Control(AOME_SET_ENABLEAUTOALTREF, 1);
+ if (use_external_partition_) {
+ encoder->Control(AV1E_SET_EXTERNAL_PARTITION, &ext_part_funcs);
+ }
+ if (partition_control_mode_ == -1) {
+ encoder->Control(AV1E_SET_MAX_PARTITION_SIZE, 128);
+ encoder->Control(AV1E_SET_MIN_PARTITION_SIZE, 4);
+ } else {
+ switch (partition_control_mode_) {
+ case 1:
+ encoder->Control(AV1E_SET_MAX_PARTITION_SIZE, 64);
+ encoder->Control(AV1E_SET_MIN_PARTITION_SIZE, 64);
+ break;
+ case 2:
+ encoder->Control(AV1E_SET_MAX_PARTITION_SIZE, 4);
+ encoder->Control(AV1E_SET_MIN_PARTITION_SIZE, 4);
+ break;
+ default: assert(0 && "Invalid partition control mode."); break;
+ }
+ }
+ }
+ }
+
+ private:
+ libaom_test::TestMode encoding_mode_;
+ int cpu_used_;
+ double psnr_;
+ unsigned int nframes_;
+ bool use_external_partition_ = false;
+ TestData test_data_;
+ int partition_control_mode_ = -1;
+};
+
+// Encode twice and expect the same psnr value.
+// The first run is a normal encoding run with restricted partition types,
+// i.e., we use control flags to force the encoder to encode with the
+// 4x4 block size.
+// The second run is to get partition decisions from a toy model that we
+// built, which will asks the encoder to encode with the 4x4 blocks.
+// We expect the encoding results are the same.
+TEST_P(ExternalPartitionTestAPI, WholePartitionTree4x4Block) {
+ ::libaom_test::Y4mVideoSource video("paris_352_288_30.y4m", 0, kFrameNum);
+ SetExternalPartition(false);
+ SetPartitionControlMode(2);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ const double psnr = GetAveragePsnr();
+
+ SetExternalPartition(true);
+ SetPartitionControlMode(2);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ const double psnr2 = GetAveragePsnr();
+
+ EXPECT_DOUBLE_EQ(psnr, psnr2);
+}
+
+AV1_INSTANTIATE_TEST_SUITE(ExternalPartitionTestAPI,
+ ::testing::Values(::libaom_test::kTwoPassGood),
+ ::testing::Values(4)); // cpu_used
+
+#else // !CONFIG_PARTITION_SEARCH_ORDER
// Feature files written during encoding, as defined in partition_strategy.c.
std::string feature_file_names[] = {
"feature_before_partition_none",
@@ -167,14 +440,14 @@
return AOM_EXT_PART_OK;
}
-class ExternalPartitionTest
+class ExternalPartitionTestDfsAPI
: public ::libaom_test::CodecTestWith2Params<libaom_test::TestMode, int>,
public ::libaom_test::EncoderTest {
protected:
- ExternalPartitionTest()
+ ExternalPartitionTestDfsAPI()
: EncoderTest(GET_PARAM(0)), encoding_mode_(GET_PARAM(1)),
cpu_used_(GET_PARAM(2)), psnr_(0.0), nframes_(0) {}
- virtual ~ExternalPartitionTest() {}
+ virtual ~ExternalPartitionTestDfsAPI() {}
virtual void SetUp() {
InitializeConfig(encoding_mode_);
@@ -256,7 +529,7 @@
// Here, we let the partition decision return invalid for all stages.
// In this case, the external partition doesn't alter the original encoder
// behavior. So we expect the same encoding results.
-TEST_P(ExternalPartitionTest, EncodeMatch) {
+TEST_P(ExternalPartitionTestDfsAPI, EncodeMatch) {
::libaom_test::Y4mVideoSource video("paris_352_288_30.y4m", 0, kFrameNum);
SetExternalPartition(false);
ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
@@ -274,7 +547,7 @@
// The second run calls send partition features function to send features to
// the external model, and we write them to file.
// The generated files should match each other.
-TEST_P(ExternalPartitionTest, SendFeatures) {
+TEST_P(ExternalPartitionTestDfsAPI, SendFeatures) {
::libaom_test::Y4mVideoSource video("paris_352_288_30.y4m", 0, kFrameNum);
SetExternalPartition(true);
SetTestSendFeatures(0);
@@ -304,9 +577,10 @@
system(command.c_str());
}
-AV1_INSTANTIATE_TEST_SUITE(ExternalPartitionTest,
+AV1_INSTANTIATE_TEST_SUITE(ExternalPartitionTestDfsAPI,
::testing::Values(::libaom_test::kTwoPassGood),
::testing::Values(4)); // cpu_used
+#endif // CONFIG_PARTITION_SEARCH_ORDER
} // namespace
#endif // !CONFIG_REALTIME_ONLY
