summaryrefslogtreecommitdiffstats
path: root/third_party/aom/test/av1_external_partition_test.cc
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--third_party/aom/test/av1_external_partition_test.cc702
1 files changed, 702 insertions, 0 deletions
diff --git a/third_party/aom/test/av1_external_partition_test.cc b/third_party/aom/test/av1_external_partition_test.cc
new file mode 100644
index 0000000000..88f6216fa5
--- /dev/null
+++ b/third_party/aom/test/av1_external_partition_test.cc
@@ -0,0 +1,702 @@
+/*
+ * Copyright (c) 2021, Alliance for Open Media. All rights reserved
+ *
+ * This source code is subject to the terms of the BSD 2 Clause License and
+ * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
+ * was not distributed with this source code in the LICENSE file, you can
+ * obtain it at www.aomedia.org/license/software. If the Alliance for Open
+ * Media Patent License 1.0 was not distributed with this source code in the
+ * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
+ */
+
+#include <fstream>
+#include <new>
+#include <sstream>
+#include <string>
+
+#include "aom/aom_codec.h"
+#include "aom/aom_external_partition.h"
+#include "av1/common/blockd.h"
+#include "av1/encoder/encodeframe_utils.h"
+#include "third_party/googletest/src/googletest/include/gtest/gtest.h"
+#include "test/codec_factory.h"
+#include "test/encode_test_driver.h"
+#include "test/y4m_video_source.h"
+#include "test/util.h"
+
+#if CONFIG_AV1_ENCODER
+#if !CONFIG_REALTIME_ONLY
+namespace {
+
+constexpr int kFrameNum = 8;
+constexpr int kVersion = 1;
+
+typedef struct TestData {
+ int version = kVersion;
+} TestData;
+
+typedef struct ToyModel {
+ 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;
+ BLOCK_SIZE block_size;
+} 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;
+ if (toy_model == nullptr) {
+ EXPECT_NE(toy_model, nullptr);
+ return AOM_EXT_PART_ERROR;
+ }
+ 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;
+ toy_model->block_size = static_cast<BLOCK_SIZE>(part_features->block_size);
+ 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_get_partition_decision_recursive(
+ aom_ext_part_model_t ext_part_model,
+ aom_partition_decision_t *ext_part_decision) {
+ ext_part_decision->current_decision = PARTITION_NONE;
+ ext_part_decision->is_final_decision = 1;
+ ToyModel *toy_model = static_cast<ToyModel *>(ext_part_model);
+ // 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) {
+ if (block_size_wide[toy_model->block_size] == 64) {
+ ext_part_decision->current_decision = PARTITION_SPLIT;
+ } else {
+ ext_part_decision->current_decision = PARTITION_NONE;
+ }
+ } else if (is_last_sb_row) {
+ if (block_size_wide[toy_model->block_size] == 64) {
+ ext_part_decision->current_decision = PARTITION_SPLIT;
+ } else {
+ ext_part_decision->current_decision = PARTITION_NONE;
+ }
+ } else if (is_last_sb_col) {
+ if (block_size_wide[toy_model->block_size] == 64) {
+ ext_part_decision->current_decision = PARTITION_SPLIT;
+ } else {
+ ext_part_decision->current_decision = PARTITION_NONE;
+ }
+ } else {
+ ext_part_decision->current_decision = 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) {}
+ ~ExternalPartitionTestAPI() override {}
+
+ void SetUp() override {
+ 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;
+ }
+
+ bool DoDecode() const override { return false; }
+
+ void BeginPassHook(unsigned int) override {
+ psnr_ = 0.0;
+ nframes_ = 0;
+ }
+
+ void PSNRPktHook(const aom_codec_cx_pkt_t *pkt) override {
+ 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; }
+
+ void SetDecisionMode(aom_ext_part_decision_mode_t mode) {
+ decision_mode_ = mode;
+ }
+
+ void PreEncodeFrameHook(::libaom_test::VideoSource *video,
+ ::libaom_test::Encoder *encoder) override {
+ if (video->frame() == 0) {
+ if (decision_mode_ == AOM_EXT_PART_WHOLE_TREE) {
+ aom_ext_part_funcs_t ext_part_funcs;
+ ext_part_funcs.priv = reinterpret_cast<void *>(&test_data_);
+ ext_part_funcs.decision_mode = AOM_EXT_PART_WHOLE_TREE;
+ 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;
+ }
+ }
+ } else if (decision_mode_ == AOM_EXT_PART_RECURSIVE) {
+ aom_ext_part_funcs_t ext_part_funcs;
+ ext_part_funcs.priv = reinterpret_cast<void *>(&test_data_);
+ ext_part_funcs.decision_mode = AOM_EXT_PART_RECURSIVE;
+ 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_recursive;
+ 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;
+ }
+ }
+ } else {
+ assert(0 && "Invalid decision mode.");
+ }
+ }
+ }
+
+ 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;
+ aom_ext_part_decision_mode_t decision_mode_;
+};
+
+// 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);
+ SetDecisionMode(AOM_EXT_PART_WHOLE_TREE);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ const double psnr = GetAveragePsnr();
+
+ SetExternalPartition(true);
+ SetPartitionControlMode(2);
+ SetDecisionMode(AOM_EXT_PART_WHOLE_TREE);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ const double psnr2 = GetAveragePsnr();
+
+ EXPECT_DOUBLE_EQ(psnr, psnr2);
+}
+
+TEST_P(ExternalPartitionTestAPI, RecursivePartition) {
+ ::libaom_test::Y4mVideoSource video("paris_352_288_30.y4m", 0, kFrameNum);
+ SetExternalPartition(false);
+ SetPartitionControlMode(1);
+ SetDecisionMode(AOM_EXT_PART_RECURSIVE);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ const double psnr = GetAveragePsnr();
+
+ SetExternalPartition(true);
+ SetPartitionControlMode(1);
+ SetDecisionMode(AOM_EXT_PART_RECURSIVE);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ const double psnr2 = GetAveragePsnr();
+
+ const double psnr_thresh = 0.02;
+ EXPECT_NEAR(psnr, psnr2, psnr_thresh);
+}
+
+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",
+ "feature_before_partition_none_prune_rect",
+ "feature_after_partition_none_prune",
+ "feature_after_partition_none_terminate",
+ "feature_after_partition_split_terminate",
+ "feature_after_partition_split_prune_rect",
+ "feature_after_partition_rect",
+ "feature_after_partition_ab",
+};
+
+// Files written here in the test, where the feature data is received
+// from the API.
+std::string test_feature_file_names[] = {
+ "test_feature_before_partition_none",
+ "test_feature_before_partition_none_prune_rect",
+ "test_feature_after_partition_none_prune",
+ "test_feature_after_partition_none_terminate",
+ "test_feature_after_partition_split_terminate",
+ "test_feature_after_partition_split_prune_rect",
+ "test_feature_after_partition_rect",
+ "test_feature_after_partition_ab",
+};
+
+static void write_features_to_file(const float *features,
+ const int feature_size, const int id) {
+ if (!WRITE_FEATURE_TO_FILE) return;
+ char filename[256];
+ snprintf(filename, sizeof(filename), "%s",
+ test_feature_file_names[id].c_str());
+ FILE *pfile = fopen(filename, "a");
+ ASSERT_NE(pfile, nullptr);
+ for (int i = 0; i < feature_size; ++i) {
+ fprintf(pfile, "%.6f", features[i]);
+ if (i < feature_size - 1) fprintf(pfile, ",");
+ }
+ fprintf(pfile, "\n");
+ fclose(pfile);
+}
+
+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;
+ if (toy_model == nullptr) {
+ EXPECT_NE(toy_model, nullptr);
+ return AOM_EXT_PART_ERROR;
+ }
+ 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_create_model_test(
+ void *priv, const aom_ext_part_config_t *part_config,
+ aom_ext_part_model_t *ext_part_model) {
+ (void)priv;
+ (void)ext_part_model;
+ EXPECT_EQ(part_config->superblock_size, BLOCK_64X64);
+ // Return status indicates it's a encoder test. It lets the encoder
+ // set a flag and write partition features to text files.
+ return AOM_EXT_PART_TEST;
+}
+
+aom_ext_part_status_t ext_part_send_features(
+ aom_ext_part_model_t ext_part_model,
+ const aom_partition_features_t *part_features) {
+ (void)ext_part_model;
+ (void)part_features;
+ return AOM_EXT_PART_OK;
+}
+
+aom_ext_part_status_t ext_part_send_features_test(
+ aom_ext_part_model_t ext_part_model,
+ const aom_partition_features_t *part_features) {
+ (void)ext_part_model;
+ if (part_features->id == AOM_EXT_PART_FEATURE_BEFORE_NONE) {
+ write_features_to_file(part_features->before_part_none.f,
+ AOM_EXT_PART_SIZE_DIRECT_SPLIT, 0);
+ } else if (part_features->id == AOM_EXT_PART_FEATURE_BEFORE_NONE_PART2) {
+ write_features_to_file(part_features->before_part_none.f_part2,
+ AOM_EXT_PART_SIZE_PRUNE_PART, 1);
+ } else if (part_features->id == AOM_EXT_PART_FEATURE_AFTER_NONE) {
+ write_features_to_file(part_features->after_part_none.f,
+ AOM_EXT_PART_SIZE_PRUNE_NONE, 2);
+ } else if (part_features->id == AOM_EXT_PART_FEATURE_AFTER_NONE_PART2) {
+ write_features_to_file(part_features->after_part_none.f_terminate,
+ AOM_EXT_PART_SIZE_TERM_NONE, 3);
+ } else if (part_features->id == AOM_EXT_PART_FEATURE_AFTER_SPLIT) {
+ write_features_to_file(part_features->after_part_split.f_terminate,
+ AOM_EXT_PART_SIZE_TERM_SPLIT, 4);
+ } else if (part_features->id == AOM_EXT_PART_FEATURE_AFTER_SPLIT_PART2) {
+ write_features_to_file(part_features->after_part_split.f_prune_rect,
+ AOM_EXT_PART_SIZE_PRUNE_RECT, 5);
+ } else if (part_features->id == AOM_EXT_PART_FEATURE_AFTER_RECT) {
+ write_features_to_file(part_features->after_part_rect.f,
+ AOM_EXT_PART_SIZE_PRUNE_AB, 6);
+ } else if (part_features->id == AOM_EXT_PART_FEATURE_AFTER_AB) {
+ write_features_to_file(part_features->after_part_ab.f,
+ AOM_EXT_PART_SIZE_PRUNE_4_WAY, 7);
+ }
+ return AOM_EXT_PART_TEST;
+}
+
+aom_ext_part_status_t ext_part_get_partition_decision(
+ aom_ext_part_model_t ext_part_model,
+ aom_partition_decision_t *ext_part_decision) {
+ (void)ext_part_model;
+ (void)ext_part_decision;
+ // Return an invalid decision such that the encoder doesn't take any
+ // partition decision from the ml model.
+ return AOM_EXT_PART_ERROR;
+}
+
+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 ExternalPartitionTestDfsAPI
+ : public ::libaom_test::CodecTestWith2Params<libaom_test::TestMode, int>,
+ public ::libaom_test::EncoderTest {
+ protected:
+ ExternalPartitionTestDfsAPI()
+ : EncoderTest(GET_PARAM(0)), encoding_mode_(GET_PARAM(1)),
+ cpu_used_(GET_PARAM(2)), psnr_(0.0), nframes_(0) {}
+ ~ExternalPartitionTestDfsAPI() override = default;
+
+ void SetUp() override {
+ 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;
+ }
+
+ bool DoDecode() const override { return false; }
+
+ void BeginPassHook(unsigned int) override {
+ psnr_ = 0.0;
+ nframes_ = 0;
+ }
+
+ void PSNRPktHook(const aom_codec_cx_pkt_t *pkt) override {
+ 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 SetTestSendFeatures(int test_send_features) {
+ test_send_features_ = test_send_features;
+ }
+
+ void PreEncodeFrameHook(::libaom_test::VideoSource *video,
+ ::libaom_test::Encoder *encoder) override {
+ if (video->frame() == 0) {
+ aom_ext_part_funcs_t ext_part_funcs;
+ ext_part_funcs.priv = reinterpret_cast<void *>(&test_data_);
+ if (use_external_partition_) {
+ ext_part_funcs.create_model = ext_part_create_model;
+ ext_part_funcs.send_features = ext_part_send_features;
+ }
+ if (test_send_features_ == 1) {
+ ext_part_funcs.create_model = ext_part_create_model;
+ ext_part_funcs.send_features = ext_part_send_features_test;
+ } else if (test_send_features_ == 0) {
+ ext_part_funcs.create_model = ext_part_create_model_test;
+ ext_part_funcs.send_features = ext_part_send_features;
+ }
+ ext_part_funcs.get_partition_decision = ext_part_get_partition_decision;
+ 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);
+ }
+ }
+ }
+
+ private:
+ libaom_test::TestMode encoding_mode_;
+ int cpu_used_;
+ double psnr_;
+ unsigned int nframes_;
+ bool use_external_partition_ = false;
+ int test_send_features_ = -1;
+ TestData test_data_;
+};
+
+// Encode twice and expect the same psnr value.
+// The first run is the baseline without external partition.
+// The second run is to get partition decisions from the toy model we defined.
+// 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(ExternalPartitionTestDfsAPI, EncodeMatch) {
+ ::libaom_test::Y4mVideoSource video("paris_352_288_30.y4m", 0, kFrameNum);
+ SetExternalPartition(false);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ const double psnr = GetAveragePsnr();
+
+ SetExternalPartition(true);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ const double psnr2 = GetAveragePsnr();
+
+ EXPECT_DOUBLE_EQ(psnr, psnr2);
+}
+
+// Encode twice to compare generated feature files.
+// The first run let the encoder write partition features to file.
+// 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(ExternalPartitionTestDfsAPI, SendFeatures) {
+ ::libaom_test::Y4mVideoSource video("paris_352_288_30.y4m", 0, kFrameNum);
+ SetExternalPartition(true);
+ SetTestSendFeatures(0);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+
+ SetExternalPartition(true);
+ SetTestSendFeatures(1);
+ ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+ if (!WRITE_FEATURE_TO_FILE) return;
+
+ // Compare feature files by reading them into strings.
+ for (int i = 0; i < 8; ++i) {
+ std::ifstream base_file(feature_file_names[i]);
+ ASSERT_TRUE(base_file.good());
+ std::stringstream base_stream;
+ base_stream << base_file.rdbuf();
+ std::string base_string = base_stream.str();
+
+ std::ifstream test_file(test_feature_file_names[i]);
+ ASSERT_TRUE(test_file.good());
+ std::stringstream test_stream;
+ test_stream << test_file.rdbuf();
+ std::string test_string = test_stream.str();
+
+ EXPECT_STREQ(base_string.c_str(), test_string.c_str());
+ }
+
+ // Remove files.
+ std::string command("rm -f feature_* test_feature_*");
+ system(command.c_str());
+}
+
+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
+#endif // CONFIG_AV1_ENCODER
generated by cgit v1.2.3 (git 2.39.1) at 2024-08-27 22:00:55 +0000