aom_external_partition.h

Go to the documentation of this file.

/*
 * 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_
 
#include <stdint.h>
 
#define AOM_EXT_PART_ABI_VERSION 8
 
#ifdef __cplusplus
extern "C" {
#endif
 
typedef void *aom_ext_part_model_t;
 
#define AOM_EXT_PART_SIZE_DIRECT_SPLIT 17
 
#define AOM_EXT_PART_SIZE_PRUNE_PART 25
 
#define AOM_EXT_PART_SIZE_PRUNE_NONE 4
 
#define AOM_EXT_PART_SIZE_TERM_NONE 28
 
#define AOM_EXT_PART_SIZE_TERM_SPLIT 31
 
#define AOM_EXT_PART_SIZE_PRUNE_RECT 9
 
#define AOM_EXT_PART_SIZE_PRUNE_AB 10
 
#define AOM_EXT_PART_SIZE_PRUNE_4_WAY 18
 
typedef enum aom_ext_part_decision_mode {
  AOM_EXT_PART_WHOLE_TREE = 0,
  AOM_EXT_PART_RECURSIVE = 1,
} aom_ext_part_decision_mode_t;
 
typedef struct aom_ext_part_config {
  int superblock_size;  
} aom_ext_part_config_t;
 
typedef struct aom_partition_features_before_none {
  float f[AOM_EXT_PART_SIZE_DIRECT_SPLIT];
  float f_part2[AOM_EXT_PART_SIZE_PRUNE_PART];
} aom_partition_features_before_none_t;
 
typedef struct aom_partition_features_none {
  float f[AOM_EXT_PART_SIZE_PRUNE_NONE];
  float f_terminate[AOM_EXT_PART_SIZE_TERM_NONE];
} aom_partition_features_none_t;
 
typedef struct aom_partition_features_split {
  float f_terminate[AOM_EXT_PART_SIZE_TERM_SPLIT];
  float f_prune_rect[AOM_EXT_PART_SIZE_PRUNE_RECT];
} aom_partition_features_split_t;
 
typedef struct aom_partition_features_rect {
  float f[AOM_EXT_PART_SIZE_PRUNE_AB];
} aom_partition_features_rect_t;
 
typedef struct aom_partition_features_ab {
  float f[AOM_EXT_PART_SIZE_PRUNE_4_WAY];
} aom_partition_features_ab_t;
 
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;
 
typedef struct aom_sb_tpl_features {
  int available;        
  int tpl_unit_length;  
  int num_units;        
  int64_t intra_cost[64];   
  int64_t inter_cost[64];   
  int64_t mc_dep_cost[64];  
} aom_sb_tpl_features_t;
 
typedef struct aom_sb_simple_motion_features {
  int unit_length;    
  int num_units;      
  int block_sse[64];  
  int block_var[64];  
} aom_sb_simple_motion_features_t;
 
typedef struct aom_sb_features {
  aom_sb_simple_motion_features_t motion_features;
  aom_sb_tpl_features_t tpl_features;
} aom_sb_features_t;
 
typedef struct aom_partition_features {
  // Features for the current supervised multi-stage ML model.
  AOM_EXT_PART_FEATURE_ID id;
  aom_partition_features_before_none_t before_part_none;
  aom_partition_features_none_t after_part_none;
  aom_partition_features_split_t after_part_split;
  aom_partition_features_rect_t after_part_rect;
  aom_partition_features_ab_t after_part_ab;
 
  // Features for a new ML model.
  aom_sb_features_t sb_features;  
  int mi_row;                     
  int mi_col;                     
  int frame_width;                
  int frame_height;               
  int block_size;                 
  int valid_partition_types;
  int update_type;    
  int qindex;         
  int rdmult;         
  int pyramid_level;  
  int has_above_block;     
  int above_block_width;   
  int above_block_height;  
  int has_left_block;      
  int left_block_width;    
  int left_block_height;   
  unsigned int block_sse;          
  unsigned int block_var;          
  unsigned int sub_block_sse[4];   
  unsigned int sub_block_var[4];   
  unsigned int horz_block_sse[2];  
  unsigned int horz_block_var[2];  
  unsigned int vert_block_sse[2];  
  unsigned int vert_block_var[2];  
  int64_t tpl_intra_cost;   
  int64_t tpl_inter_cost;   
  int64_t tpl_mc_dep_cost;  
} aom_partition_features_t;
 
typedef struct aom_partition_decision {
  // Decisions for directly set partition types
  int is_final_decision;         
  int num_nodes;                 
  int partition_decision[2048];  
  int current_decision;          
 
  // Decisions for partition type pruning
  int terminate_partition_search;  
  int partition_none_allowed;      
  int partition_rect_allowed[2];   
  int do_rectangular_split;        
  int do_square_split;             
  int prune_rect_part[2];          
  int horza_partition_allowed;     
  int horzb_partition_allowed;     
  int verta_partition_allowed;     
  int vertb_partition_allowed;     
  int partition_horz4_allowed;     
  int partition_vert4_allowed;     
} aom_partition_decision_t;
 
typedef struct aom_partition_stats {
  int rate;        
  int64_t dist;    
  int64_t rdcost;  
} aom_partition_stats_t;
 
typedef enum aom_ext_part_status {
  AOM_EXT_PART_OK = 0,     
  AOM_EXT_PART_ERROR = 1,  
  AOM_EXT_PART_TEST = 2,   
} aom_ext_part_status_t;
 
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);
 
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);
 
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);
 
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);
 
typedef aom_ext_part_status_t (*aom_ext_part_delete_model_fn_t)(
    aom_ext_part_model_t ext_part_model);
 
typedef struct aom_ext_part_funcs {
  aom_ext_part_create_model_fn_t create_model;
 
  aom_ext_part_send_features_fn_t send_features;
 
  aom_ext_part_get_decision_fn_t get_partition_decision;
 
  aom_ext_part_send_partition_stats_fn_t send_partition_stats;
 
  aom_ext_part_delete_model_fn_t delete_model;
 
  aom_ext_part_decision_mode_t decision_mode;
 
  void *priv;
} aom_ext_part_funcs_t;
 
#ifdef __cplusplus
}  // extern "C"
#endif
 
#endif  // AOM_AOM_AOM_EXTERNAL_PARTITION_H_