ExtPart: Add features: qindex, frame type, depth
Add features and associate them with each block.
Change-Id: Ibaf797778535e77429382f6dfb235015c0d4f30a
diff --git a/aom/aom_external_partition.h b/aom/aom_external_partition.h
index e39921d..34503b5 100644
--- a/aom/aom_external_partition.h
+++ b/aom/aom_external_partition.h
@@ -30,7 +30,7 @@
* types, removing or reassigning enums, adding/removing/rearranging
* fields to structures.
*/
-#define AOM_EXT_PART_ABI_VERSION 4
+#define AOM_EXT_PART_ABI_VERSION 5
#ifdef __cplusplus
extern "C" {
@@ -248,6 +248,10 @@
* 0x09 (00...001001) stands for PARTITION_NONE and PARTITION_SPLIT are valid.
*/
int valid_partition_types;
+ int update_type; ///< Frame update type, defined in ratectrl.h
+ int qindex; ///< Quantization index, range: [0, 255]
+ int rdmult; ///< Rate-distortion multiplier
+ int pyramid_level; ///< The level of this frame in the hierarchical structure
} aom_partition_features_t;
/*!\brief Partition decisions received from the external model.
diff --git a/av1/encoder/partition_search.c b/av1/encoder/partition_search.c
index e4d5671..001c4c1 100644
--- a/av1/encoder/partition_search.c
+++ b/av1/encoder/partition_search.c
@@ -4273,6 +4273,7 @@
const AV1_COMMON *const cm = &cpi->common;
ExtPartController *const ext_part_controller = &cpi->ext_part_controller;
MACROBLOCK *const x = &td->mb;
+ MACROBLOCKD *const xd = &x->e_mbd;
if (mi_row >= cm->mi_params.mi_rows || mi_col >= cm->mi_params.mi_cols) {
return false;
}
@@ -4289,8 +4290,20 @@
PartitionBlkParams blk_params = part_search_state.part_blk_params;
if (!av1_blk_has_rows_and_cols(&blk_params))
set_partition_cost_for_edge_blk(cm, &part_search_state);
+ const int orig_rdmult = x->rdmult;
+ setup_block_rdmult(cpi, x, mi_row, mi_col, bsize, NO_AQ, NULL);
const int valid_partition_types =
get_valid_partition_types(cpi, &part_search_state, bsize);
+ const FRAME_UPDATE_TYPE update_type =
+ get_frame_update_type(&cpi->ppi->gf_group, cpi->gf_frame_index);
+ const int qindex = av1_get_qindex(&cm->seg, xd->mi[0]->segment_id,
+ cm->quant_params.base_qindex);
+ // RD multiplier
+ const int rdmult = x->rdmult;
+ // pyramid level
+ const int pyramid_level =
+ cpi->ppi->gf_group.layer_depth[cpi->gf_frame_index];
+ x->rdmult = orig_rdmult;
aom_partition_features_t features;
features.mi_row = mi_row;
@@ -4299,6 +4312,10 @@
features.frame_height = cpi->frame_info.frame_height;
features.block_size = bsize;
features.valid_partition_types = valid_partition_types;
+ features.update_type = update_type;
+ features.qindex = qindex;
+ features.rdmult = rdmult;
+ features.pyramid_level = pyramid_level;
av1_ext_part_send_features(ext_part_controller, &features);
const bool valid_decision = av1_ext_part_get_partition_decision(
ext_part_controller, &partition_decision);
@@ -4316,9 +4333,8 @@
pc_tree->split[i] = av1_alloc_pc_tree_node(subsize);
pc_tree->split[i]->index = i;
}
- const int orig_rdmult = x->rdmult;
+ const int orig_rdmult_tmp = x->rdmult;
setup_block_rdmult(cpi, x, mi_row, mi_col, bsize, NO_AQ, NULL);
- (void)orig_rdmult;
// TODO(chengchen): check boundary conditions
// top-left
recursive_partition(cpi, td, tile_data, tp, sms_root, pc_tree->split[0],
@@ -4343,7 +4359,7 @@
this_rdcost->dist += split_rdc[i].dist;
av1_rd_cost_update(x->rdmult, this_rdcost);
}
- x->rdmult = orig_rdmult;
+ x->rdmult = orig_rdmult_tmp;
} else {
*this_rdcost = rd_search_for_fixed_partition(
cpi, td, tile_data, tp, sms_root, mi_row, mi_col, bsize, pc_tree);
