External partition: Add feature ids to indicate active features
The feature id is used to tell the external model which partition
pruning stage it is and what features to use to make partition
decisions.
Change-Id: I0aaff7e3cb54842bc35d6c6bff6d6a0d893139cb
diff --git a/aom/aom_external_partition.h b/aom/aom_external_partition.h
index d944554..13a95f0 100644
--- a/aom/aom_external_partition.h
+++ b/aom/aom_external_partition.h
@@ -50,9 +50,16 @@
/*!\brief Features pass to the external model to make partition decisions.
* Specifically, features collected before NONE partition.
+ * Features "f" are used to determine:
+ * partition_none_allowed, partition_horz_allowed, partition_vert_allowed,
+ * do_rectangular_split, do_square_split
+ * Features "f_part2" are used to determine:
+ * prune_horz, prune_vert.
*/
typedef struct aom_partition_features_before_none {
float f[17]; /**< features to determine whether partition types allowed */
+ float
+ f_part2[25]; /**< features to determine whether partition types allowed */
} aom_partition_features_before_none_t;
/*!\brief Features pass to the external model to make partition decisions.
@@ -89,6 +96,18 @@
float f[18]; /**< features to determine pruning 4-way partition */
} aom_partition_features_ab_t;
+/*!\brief Feature id to tell the external model the current stage in partition
+ * pruning and what features to use to make decisions accordingly.
+ */
+typedef enum {
+ FEATURE_BEFORE_PART_NONE,
+ FEATURE_BEFORE_PART_NONE_PART2,
+ FEATURE_AFTER_PART_NONE,
+ FEATURE_AFTER_PART_SPLIT,
+ FEATURE_AFTER_PART_RECT,
+ FEATURE_AFTER_PART_AB
+} PART_FEATURE_ID;
+
/*!\brief Features pass to the external model to make partition decisions.
*
* The encoder sends these features to the external model through
@@ -98,6 +117,7 @@
* Once new features are finalized, bump the major version of libaom.
*/
typedef struct aom_partition_features {
+ PART_FEATURE_ID id; /**< Feature ID to indicate active features */
aom_partition_features_before_none_t
before_part_none; /**< Features collected before NONE partition */
aom_partition_features_none_t
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index 351bc13..f4d0758 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -2022,6 +2022,7 @@
// Setup features.
aom_partition_features_t features;
+ features.id = FEATURE_BEFORE_PART_NONE;
for (int i = 0; i < FEATURE_SIZE_SMS_SPLIT; ++i) {
features.before_part_none.f[i] = features_from_motion[i];
}
@@ -2058,8 +2059,9 @@
// Setup features.
aom_partition_features_t features;
+ features.id = FEATURE_BEFORE_PART_NONE_PART2;
for (int i = 0; i < FEATURE_SIZE_SMS_PRUNE_PART; ++i) {
- features.before_part_none.f[i] = features_from_motion[i];
+ features.before_part_none.f_part2[i] = features_from_motion[i];
}
// Send necessary features to the external model.
@@ -2094,6 +2096,7 @@
if (!frame_is_intra_only(cm) && ext_part_controller->ready) {
// Setup features.
aom_partition_features_t features;
+ features.id = FEATURE_AFTER_PART_NONE;
prepare_features_after_part_none(cpi, cm, x, sms_tree, part_search_state,
pb_source_variance, &features);
@@ -2135,6 +2138,7 @@
if (!frame_is_intra_only(cm) && cpi->ext_part_controller.ready) {
// Setup features.
aom_partition_features_t features;
+ features.id = FEATURE_AFTER_PART_SPLIT;
prepare_features_after_split(cpi, x, sms_tree, part_search_state, best_rdc,
partition_none_rdcost, partition_split_rdcost,
split_block_rdcost, &features);
@@ -2189,6 +2193,7 @@
ab_partition_allowed) {
// Setup features.
aom_partition_features_t features;
+ features.id = FEATURE_AFTER_PART_RECT;
prepare_features_after_part_rect(cpi, x, pc_tree, bsize, best_rdcost,
rect_part_rd, split_rd, &features);
@@ -2230,6 +2235,7 @@
if (!frame_is_intra_only(cm) && ext_part_controller->ready) {
// Setup features.
aom_partition_features_t features;
+ features.id = FEATURE_AFTER_PART_AB;
prepare_features_after_part_ab(cpi, x, bsize, part_ctx, best_rd,
rect_part_rd, split_rd, pb_source_variance,
mi_row, mi_col, &features);
