ExtPart: Add a variable to distinguish deicion mode
Change-Id: Ia84f517ee2ba85db79e2689c2745cc26eedc81c4
diff --git a/aom/aom_external_partition.h b/aom/aom_external_partition.h
index b6ea96e..daf8fe4 100644
--- a/aom/aom_external_partition.h
+++ b/aom/aom_external_partition.h
@@ -81,6 +81,19 @@
*/
#define SIZE_PRUNE_4_WAY 18
+/*!\brief Decision mode of the external partition model.
+ * WHOLE_TREE_DECISION: the external partition model should provide the
+ * whole partition tree for the superblock.
+ *
+ * RECURSIVE_DECISION: the external partition model provides the partition
+ * decision of the current block only. The decision process starts from
+ * the superblock size, down to the smallest block size (4x4) recursively.
+ */
+typedef enum aom_ext_part_decision_mode {
+ WHOLE_TREE_DECISION = 0,
+ RECURSIVE_DECISION = 1,
+} aom_ext_part_decision_mode_t;
+
/*!\brief Config information sent to the external partition model.
*
* For example, the maximum superblock size determined by the sequence header.
@@ -368,6 +381,11 @@
aom_ext_part_delete_model_fn_t delete_model;
/*!
+ * The decision mode of the model.
+ */
+ aom_ext_part_decision_mode_t decision_mode;
+
+ /*!
* Private data for the external partition model.
*/
void *priv;
diff --git a/av1/encoder/external_partition.c b/av1/encoder/external_partition.c
index 542b2bb..79f8b4c 100644
--- a/av1/encoder/external_partition.c
+++ b/av1/encoder/external_partition.c
@@ -91,3 +91,8 @@
if (status != AOM_EXT_PART_OK) return false;
return true;
}
+
+aom_ext_part_decision_mode_t av1_get_ext_part_decision_mode(
+ const ExtPartController *ext_part_controller) {
+ return ext_part_controller->funcs.decision_mode;
+}
diff --git a/av1/encoder/external_partition.h b/av1/encoder/external_partition.h
index 20f03ed..f74973e 100644
--- a/av1/encoder/external_partition.h
+++ b/av1/encoder/external_partition.h
@@ -47,6 +47,9 @@
bool av1_ext_part_send_features(ExtPartController *ext_part_controller,
const aom_partition_features_t *features);
+aom_ext_part_decision_mode_t av1_get_ext_part_decision_mode(
+ const ExtPartController *ext_part_controller);
+
/*!\endcond */
#ifdef __cplusplus
} // extern "C"
diff --git a/av1/encoder/partition_search.c b/av1/encoder/partition_search.c
index ccdd8e5..8a1e2aa 100644
--- a/av1/encoder/partition_search.c
+++ b/av1/encoder/partition_search.c
@@ -4320,12 +4320,17 @@
RD_STATS *best_rd_cost) {
if (cpi->ext_part_controller.ready) {
bool valid_search = true;
- if (ML_PARTITION_WHOLE_TREE_DECISION) {
+ const aom_ext_part_decision_mode_t decision_mode =
+ av1_get_ext_part_decision_mode(&cpi->ext_part_controller);
+ if (decision_mode == WHOLE_TREE_DECISION) {
valid_search = ml_partition_search_whole_tree(
cpi, td, tile_data, tp, sms_root, mi_row, mi_col, bsize);
- } else {
+ } else if (decision_mode == RECURSIVE_DECISION) {
valid_search = ml_partition_search_partial(
cpi, td, tile_data, tp, sms_root, mi_row, mi_col, bsize);
+ } else {
+ assert(0 && "Unknown decision mode.");
+ return false;
}
if (!valid_search) {
assert(0 && "Invalid search from ML model, partition search failed.");
diff --git a/test/av1_external_partition_test.cc b/test/av1_external_partition_test.cc
index 45ca847a..20329de 100644
--- a/test/av1_external_partition_test.cc
+++ b/test/av1_external_partition_test.cc
@@ -245,6 +245,7 @@
if (video->frame() == 0) {
aom_ext_part_funcs_t ext_part_funcs;
ext_part_funcs.priv = reinterpret_cast<void *>(&test_data_);
+ ext_part_funcs.decision_mode = WHOLE_TREE_DECISION;
ext_part_funcs.create_model = ext_part_create_model;
ext_part_funcs.send_features = ext_part_send_features;
ext_part_funcs.get_partition_decision =
