External partition: features for partition decisions
Add features in the external partition interface.
These features are the same as those used in the current partition
pruning ml model.
Change-Id: I889f602667c0728449ada44bdc735684b0db5360
diff --git a/aom/aom_external_partition.h b/aom/aom_external_partition.h
index 129ef3a..d944554 100644
--- a/aom/aom_external_partition.h
+++ b/aom/aom_external_partition.h
@@ -49,6 +49,47 @@
} aom_ext_part_config_t;
/*!\brief Features pass to the external model to make partition decisions.
+ * Specifically, features collected before NONE partition.
+ */
+typedef struct aom_partition_features_before_none {
+ float f[17]; /**< features to determine whether partition types allowed */
+} aom_partition_features_before_none_t;
+
+/*!\brief Features pass to the external model to make partition decisions.
+ * Specifically, features collected after NONE partition.
+ */
+typedef struct aom_partition_features_none {
+ float rate; /**< normalized rate cost of the partition */
+ float dist; /**< normalized distortion of the partition */
+ float q; /**< normalized quantization parameter */
+ float source_variance; /**< normalized variance of the source */
+ float f_terminate[28]; /**< features to determine termination of partition */
+} aom_partition_features_none_t;
+
+/*!\brief Features pass to the external model to make partition decisions.
+ * Specifically, features collected after SPLIT partition.
+ */
+typedef struct aom_partition_features_split {
+ float f_terminate[31]; /**< features to determine termination of partition */
+ float f_prune_rect[9]; /**< features to determine pruning rect partition */
+} aom_partition_features_split_t;
+
+/*!\brief Features pass to the external model to make partition decisions.
+ * Specifically, features collected after RECTANGULAR partition.
+ */
+typedef struct aom_partition_features_rect {
+ float f[10]; /**< features to determine pruning AB partition */
+} aom_partition_features_rect_t;
+
+/*!\brief Features pass to the external model to make partition decisions.
+ * Specifically, features collected after AB partition: HORZ_A, HORZ_B, VERT_A,
+ * VERT_B.
+ */
+typedef struct aom_partition_features_ab {
+ float f[18]; /**< features to determine pruning 4-way partition */
+} aom_partition_features_ab_t;
+
+/*!\brief Features pass to the external model to make partition decisions.
*
* The encoder sends these features to the external model through
* "func()" defined in .....
@@ -57,10 +98,16 @@
* Once new features are finalized, bump the major version of libaom.
*/
typedef struct aom_partition_features {
- // Features are unclear yet
- // Candidates include:
- // mv
- int64_t sse; /**< sum squared error */
+ aom_partition_features_before_none_t
+ before_part_none; /**< Features collected before NONE partition */
+ aom_partition_features_none_t
+ after_part_none; /**< Features collected after NONE partition */
+ aom_partition_features_split_t
+ after_part_split; /**< Features collected after SPLIT partition */
+ aom_partition_features_rect_t
+ after_part_rect; /**< Features collected after RECTANGULAR partition */
+ aom_partition_features_ab_t
+ after_part_ab; /**< Features collected after AB partition */
} aom_partition_features_t;
/*!\brief Partition decisions received from the external model.
@@ -73,8 +120,23 @@
* Once new features are finalized, bump the major version of libaom.
*/
typedef struct aom_partition_decision {
+ // Decisions for directly set partition types
int is_final_decision; /**< The flag whether it is the final decision */
int partition_decision[256]; /**< Partition decisions */
+
+ // Decisions for partition type pruning
+ int terminate_partition_search; /**< Terminate further partition search */
+ int partition_none_allowed; /**< Allow partition none type */
+ int partition_rect_allowed[2]; /**< Allow rectangular partitions */
+ int do_rectangular_split; /**< Try rectangular split partition */
+ int do_square_split; /**< Try square split partition */
+ int prune_rect_part[2]; /**< Prune rectangular partition */
+ int horza_partition_allowed; /**< Allow HORZ_A partitioin */
+ int horzb_partition_allowed; /**< Allow HORZ_B partitioin */
+ int verta_partition_allowed; /**< Allow VERT_A partitioin */
+ int vertb_partition_allowed; /**< Allow VERT_B partitioin */
+ int partition_horz4_allowed; /**< Allow HORZ4 partition */
+ int partition_vert4_allowed; /**< Allow VERT4 partition */
} aom_partition_decision_t;
/*!\brief Encoding stats for the given partition decision.
