| /* |
| * 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. |
| */ |
| #ifndef AOM_AOM_AOM_EXTERNAL_PARTITION_H_ |
| #define AOM_AOM_AOM_EXTERNAL_PARTITION_H_ |
| |
| /*!\defgroup aom_encoder AOMedia AOM/AV1 Encoder |
| * \ingroup aom |
| * |
| * @{ |
| */ |
| #include <stdint.h> |
| |
| /*!\file |
| * \brief Provides function pointer definitions for the external partition. |
| */ |
| |
| /*!\brief Current ABI version number |
| * |
| * \internal |
| * If this file is altered in any way that changes the ABI, this value |
| * must be bumped. Examples include, but are not limited to, changing |
| * types, removing or reassigning enums, adding/removing/rearranging |
| * fields to structures. |
| */ |
| #define AOM_EXT_PART_ABI_VERSION 3 |
| |
| #ifdef __cplusplus |
| extern "C" { |
| #endif |
| |
| /*!\brief Abstract external partition model handler |
| */ |
| typedef void *aom_ext_part_model_t; |
| |
| /*!\brief Number of features to determine whether to skip partition none and |
| * do partition split directly. The same as "FEATURE_SIZE_SMS_SPLIT". |
| */ |
| #define AOM_EXT_PART_SIZE_DIRECT_SPLIT 17 |
| |
| /*!\brief Number of features to use simple motion search to prune out |
| * rectangular partition in some direction. The same as |
| * "FEATURE_SIZE_SMS_PRUNE_PART". |
| */ |
| #define AOM_EXT_PART_SIZE_PRUNE_PART 25 |
| |
| /*!\brief Number of features to prune split and rectangular partition |
| * after PARTITION_NONE. |
| */ |
| #define AOM_EXT_PART_SIZE_PRUNE_NONE 4 |
| |
| /*!\brief Number of features to terminates partition after partition none using |
| * simple_motion_search features and the rate, distortion, and rdcost of |
| * PARTITION_NONE. The same as "FEATURE_SIZE_SMS_TERM_NONE". |
| */ |
| #define AOM_EXT_PART_SIZE_TERM_NONE 28 |
| |
| /*!\brief Number of features to terminates partition after partition split. |
| */ |
| #define AOM_EXT_PART_SIZE_TERM_SPLIT 31 |
| |
| /*!\brief Number of features to prune rectangular partition using stats |
| * collected after partition split. |
| */ |
| #define AOM_EXT_PART_SIZE_PRUNE_RECT 9 |
| |
| /*!\brief Number of features to prune AB partition using stats |
| * collected after rectangular partition.. |
| */ |
| #define AOM_EXT_PART_SIZE_PRUNE_AB 10 |
| |
| /*!\brief Number of features to prune 4-way partition using stats |
| * collected after AB partition. |
| */ |
| #define AOM_EXT_PART_SIZE_PRUNE_4_WAY 18 |
| |
| /*!\brief Decision mode of the external partition model. |
| * AOM_EXT_PART_WHOLE_TREE: the external partition model should provide the |
| * whole partition tree for the superblock. |
| * |
| * AOM_EXT_PART_RECURSIVE: 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 { |
| AOM_EXT_PART_WHOLE_TREE = 0, |
| AOM_EXT_PART_RECURSIVE = 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. |
| */ |
| typedef struct aom_ext_part_config { |
| int superblock_size; ///< super block size (either 64x64 or 128x128) |
| } aom_ext_part_config_t; |
| |
| /*!\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 { |
| /*! features to determine whether skip partition none and do split directly */ |
| float f[AOM_EXT_PART_SIZE_DIRECT_SPLIT]; |
| /*! features to determine whether to prune rectangular partition */ |
| float f_part2[AOM_EXT_PART_SIZE_PRUNE_PART]; |
| } 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 { |
| /*! features to prune split and rectangular partition */ |
| float f[AOM_EXT_PART_SIZE_PRUNE_NONE]; |
| /*! features to determine termination of partition */ |
| float f_terminate[AOM_EXT_PART_SIZE_TERM_NONE]; |
| } 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 { |
| /*! features to determine termination of partition */ |
| float f_terminate[AOM_EXT_PART_SIZE_TERM_SPLIT]; |
| /*! features to determine pruning rect partition */ |
| float f_prune_rect[AOM_EXT_PART_SIZE_PRUNE_RECT]; |
| } 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 { |
| /*! features to determine pruning AB partition */ |
| float f[AOM_EXT_PART_SIZE_PRUNE_AB]; |
| } 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 { |
| /*! features to determine pruning 4-way partition */ |
| float f[AOM_EXT_PART_SIZE_PRUNE_4_WAY]; |
| } 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 { |
| AOM_EXT_PART_FEATURE_BEFORE_NONE, |
| AOM_EXT_PART_FEATURE_BEFORE_NONE_PART2, |
| AOM_EXT_PART_FEATURE_AFTER_NONE, |
| AOM_EXT_PART_FEATURE_AFTER_NONE_PART2, |
| AOM_EXT_PART_FEATURE_AFTER_SPLIT, |
| AOM_EXT_PART_FEATURE_AFTER_SPLIT_PART2, |
| AOM_EXT_PART_FEATURE_AFTER_RECT, |
| AOM_EXT_PART_FEATURE_AFTER_AB |
| } AOM_EXT_PART_FEATURE_ID; |
| |
| /*!\brief Features collected from the tpl process. |
| * |
| * The tpl process collects information that help measure the inter-frame |
| * dependency. |
| * The tpl process is computed in the unit of tpl_bsize_1d (16x16). |
| * Therefore, the max number of units inside a superblock is |
| * 128x128 / (16x16) = 64. Change it if the tpl process changes. |
| */ |
| typedef struct aom_sb_tpl_features { |
| int available; ///< If tpl stats are available |
| int tpl_unit_length; ///< The block length of tpl process |
| int num_units; ///< The number of units inside the current superblock |
| int64_t intra_cost[64]; ///< The intra cost of each unit |
| int64_t inter_cost[64]; ///< The inter cost of each unit |
| int64_t mc_dep_cost[64]; ///< The motion compensated dependency cost |
| } aom_sb_tpl_features_t; |
| |
| /*!\brief Features collected from the simple motion process. |
| * |
| * The simple motion process collects information by applying motion compensated |
| * prediction on each block. |
| * The block size is 16x16, which could be changed. If it is changed, update |
| * comments and the array size here. |
| */ |
| typedef struct aom_sb_simple_motion_features { |
| int unit_length; ///< The block length of the simple motion process |
| int num_units; ///< The number of units inside the current superblock |
| int block_sse[64]; ///< Sum of squared error of each unit |
| int block_var[64]; ///< Variance of each unit |
| } aom_sb_simple_motion_features_t; |
| |
| /*!\brief Features of each super block. |
| * |
| * Features collected for each super block before partition search. |
| */ |
| typedef struct aom_sb_features { |
| /*! Features from motion search */ |
| aom_sb_simple_motion_features_t motion_features; |
| /*! Features from tpl process */ |
| aom_sb_tpl_features_t tpl_features; |
| } aom_sb_features_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 ..... |
| * |
| * NOTE: new member variables may be added to this structure in the future. |
| * Once new features are finalized, bump the major version of libaom. |
| */ |
| typedef struct aom_partition_features { |
| // Features for the current supervised multi-stage ML model. |
| /*! Feature ID to indicate active features */ |
| AOM_EXT_PART_FEATURE_ID id; |
| /*! Features collected before NONE partition */ |
| aom_partition_features_before_none_t before_part_none; |
| /*! Features collected after NONE partition */ |
| aom_partition_features_none_t after_part_none; |
| /*! Features collected after SPLIT partition */ |
| aom_partition_features_split_t after_part_split; |
| /*! Features collected after RECTANGULAR partition */ |
| aom_partition_features_rect_t after_part_rect; |
| /*! Features collected after AB partition */ |
| aom_partition_features_ab_t after_part_ab; |
| |
| // Features for a new ML model. |
| aom_sb_features_t sb_features; ///< Features collected for the super block |
| int mi_row; ///< Mi_row position of the block |
| int mi_col; ///< Mi_col position of the block |
| int frame_width; ///< Frame width |
| int frame_height; ///< Frame height |
| int block_size; ///< As "BLOCK_SIZE" in av1/common/enums.h |
| } aom_partition_features_t; |
| |
| /*!\brief Partition decisions received from the external model. |
| * |
| * The encoder receives partition decisions and encodes the superblock |
| * with the given partition type. |
| * The encoder receives it from "func()" define in .... |
| * |
| * NOTE: new member variables may be added to this structure in the future. |
| * 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's the final decision |
| int num_nodes; ///< The number of leaf nodes |
| int partition_decision[2048]; ///< Partition decisions |
| int current_decision; ///< Partition decision for the current block |
| |
| // 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. |
| * |
| * The encoding stats collected by encoding the superblock with the |
| * given partition types. |
| * The encoder sends the stats to the external model for training |
| * or inference though "func()" defined in .... |
| */ |
| typedef struct aom_partition_stats { |
| int rate; ///< Rate cost of the block |
| int64_t dist; ///< Distortion of the block |
| int64_t rdcost; ///< Rate-distortion cost of the block |
| } aom_partition_stats_t; |
| |
| /*!\brief Enum for return status. |
| */ |
| typedef enum aom_ext_part_status { |
| AOM_EXT_PART_OK = 0, ///< Status of success |
| AOM_EXT_PART_ERROR = 1, ///< Status of failure |
| AOM_EXT_PART_TEST = 2, ///< Status used for tests |
| } aom_ext_part_status_t; |
| |
| /*!\brief Callback of creating an external partition model. |
| * |
| * The callback is invoked by the encoder to create an external partition |
| * model. |
| * |
| * \param[in] priv Callback's private data |
| * \param[in] part_config Config information pointer for model creation |
| * \param[out] ext_part_model Pointer to the model |
| */ |
| typedef aom_ext_part_status_t (*aom_ext_part_create_model_fn_t)( |
| void *priv, const aom_ext_part_config_t *part_config, |
| aom_ext_part_model_t *ext_part_model); |
| |
| /*!\brief Callback of sending features to the external partition model. |
| * |
| * The callback is invoked by the encoder to send features to the external |
| * partition model. |
| * |
| * \param[in] ext_part_model The external model |
| * \param[in] part_features Pointer to the features |
| */ |
| typedef aom_ext_part_status_t (*aom_ext_part_send_features_fn_t)( |
| aom_ext_part_model_t ext_part_model, |
| const aom_partition_features_t *part_features); |
| |
| /*!\brief Callback of receiving partition decisions from the external |
| * partition model. |
| * |
| * The callback is invoked by the encoder to receive partition decisions from |
| * the external partition model. |
| * |
| * \param[in] ext_part_model The external model |
| * \param[in] ext_part_decision Pointer to the partition decisions |
| */ |
| typedef aom_ext_part_status_t (*aom_ext_part_get_decision_fn_t)( |
| aom_ext_part_model_t ext_part_model, |
| aom_partition_decision_t *ext_part_decision); |
| |
| /*!\brief Callback of sending stats to the external partition model. |
| * |
| * The callback is invoked by the encoder to send encoding stats to |
| * the external partition model. |
| * |
| * \param[in] ext_part_model The external model |
| * \param[in] ext_part_stats Pointer to the encoding stats |
| */ |
| typedef aom_ext_part_status_t (*aom_ext_part_send_partition_stats_fn_t)( |
| aom_ext_part_model_t ext_part_model, |
| const aom_partition_stats_t *ext_part_stats); |
| |
| /*!\brief Callback of deleting the external partition model. |
| * |
| * The callback is invoked by the encoder to delete the external partition |
| * model. |
| * |
| * \param[in] ext_part_model The external model |
| */ |
| typedef aom_ext_part_status_t (*aom_ext_part_delete_model_fn_t)( |
| aom_ext_part_model_t ext_part_model); |
| |
| /*!\brief Callback function set for external partition model. |
| * |
| * Uses can enable external partition model by registering a set of |
| * callback functions with the flag: AV1E_SET_EXTERNAL_PARTITION_MODEL |
| */ |
| typedef struct aom_ext_part_funcs { |
| /*! |
| * Create an external partition model. |
| */ |
| aom_ext_part_create_model_fn_t create_model; |
| |
| /*! |
| * Send features to the external partition model to make partition decisions. |
| */ |
| aom_ext_part_send_features_fn_t send_features; |
| |
| /*! |
| * Get partition decisions from the external partition model. |
| */ |
| aom_ext_part_get_decision_fn_t get_partition_decision; |
| |
| /*! |
| * Send stats of the current partition to the external model. |
| */ |
| aom_ext_part_send_partition_stats_fn_t send_partition_stats; |
| |
| /*! |
| * Delete the external partition model. |
| */ |
| 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; |
| } aom_ext_part_funcs_t; |
| |
| /*!@} - end defgroup aom_encoder*/ |
| #ifdef __cplusplus |
| } // extern "C" |
| #endif |
| |
| #endif // AOM_AOM_AOM_EXTERNAL_PARTITION_H_ |