av1/encoder/part_split_prune_tflite.h

/*
 * Copyright (c) 2024, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 3-Clause Clear License
 * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear
 * License was not distributed with this source code in the LICENSE file, you
 * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/.  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
 * aomedia.org/license/patent-license/.
 */

#ifndef AV1_ENCODER_SIMPLE_INTRAPRED_TFLITE_H_
#define AV1_ENCODER_SIMPLE_INTRAPRED_TFLITE_H_

#ifdef __cplusplus
extern "C" {
#endif

#include "av1/common/av1_common_int.h"

#define DEFAULT_QP_LOW 85
#define DEFAULT_QP_HIGH 85

typedef enum {
  MODEL_OTHER = 0,
  MODEL_128X128,
  MODEL_64X64,
  MODEL_32X32,
  MODEL_16X16,
  MODEL_INTER_NONE_64X64_110,
  MODEL_INTER_NONE_64X64_135,
  MODEL_INTER_NONE_32X32_110,
  MODEL_INTER_NONE_32X32_135,
  MODEL_INTER_NONE_16X16_110,
  MODEL_INTER_NONE_16X16_135,
  MODEL_INTER_NONE_BS11_110,
  MODEL_INTER_NONE_BS11_135,
  MODEL_INTER_NONE_BS10_110,
  MODEL_INTER_NONE_BS10_135,
  MODEL_INTER_NONE_BS8_110,
  MODEL_INTER_NONE_BS8_135,
  MODEL_INTER_NONE_BS7_110,
  MODEL_INTER_NONE_BS7_135,
  MODEL_INTER_NONE_BS5_110,
  MODEL_INTER_NONE_BS5_135,
  MODEL_INTER_NONE_BS4_110,
  MODEL_INTER_NONE_BS4_135,
  MODEL_INTER_SPLIT_64X64,
  MODEL_INTER_SPLIT_32X32,
  MODEL_INTER_SPLIT_16X16,
  MODEL_INTER_SPLIT_8X8,
  MODEL_COUNT

} MODEL_TYPE;

enum {
  FEATURE_INTER_RD_MULT = 0,
  FEATURE_INTER_FULL_PSNR,
  FEATURE_INTER_FULL_Q_COEFF_MAX,
  FEATURE_INTER_FULL_Q_COEFF_NONZ,
  FEATURE_INTER_SQ_0_PSNR,
  FEATURE_INTER_SQ_0_Q_COEFF_MAX,
  FEATURE_INTER_SQ_0_Q_COEFF_NONZ,
  FEATURE_INTER_SQ_1_PSNR,
  FEATURE_INTER_SQ_1_Q_COEFF_MAX,
  FEATURE_INTER_SQ_1_Q_COEFF_NONZ,
  FEATURE_INTER_SQ_2_PSNR,
  FEATURE_INTER_SQ_2_Q_COEFF_MAX,
  FEATURE_INTER_SQ_2_Q_COEFF_NONZ,
  FEATURE_INTER_SQ_3_PSNR,
  FEATURE_INTER_SQ_3_Q_COEFF_MAX,
  FEATURE_INTER_SQ_3_Q_COEFF_NONZ,
  FEATURE_INTER_FULL_LOG_MAG,
  FEATURE_INTER_FULL_ANGLE_RAD,
  FEATURE_INTER_SQ_0_LOG_MAG,
  FEATURE_INTER_SQ_0_ANGLE_RAD,
  FEATURE_INTER_SQ_1_LOG_MAG,
  FEATURE_INTER_SQ_1_ANGLE_RAD,
  FEATURE_INTER_SQ_2_LOG_MAG,
  FEATURE_INTER_SQ_2_ANGLE_RAD,
  FEATURE_INTER_SQ_3_LOG_MAG,
  FEATURE_INTER_SQ_3_ANGLE_RAD,
  FEATURE_INTER_FULL_LOG_SATDQ,
  FEATURE_INTER_SQ_0_LOG_SATDQ,
  FEATURE_INTER_SQ_1_LOG_SATDQ,
  FEATURE_INTER_SQ_2_LOG_SATDQ,
  FEATURE_INTER_SQ_3_LOG_SATDQ,
  FEATURE_INTER_FULL_LOG_SATD,
  FEATURE_INTER_SQ_0_LOG_SATD,
  FEATURE_INTER_SQ_1_LOG_SATD,
  FEATURE_INTER_SQ_2_LOG_SATD,
  FEATURE_INTER_SQ_3_LOG_SATD,
  FEATURE_INTER_HOR_0_PSNR,
  FEATURE_INTER_HOR_0_Q_COEFF_MAX,
  FEATURE_INTER_HOR_0_Q_COEFF_NONZ,
  FEATURE_INTER_HOR_0_LOG_MAG,
  FEATURE_INTER_HOR_0_ANGLE_RAD,
  FEATURE_INTER_HOR_0_LOG_SATDQ,
  FEATURE_INTER_HOR_0_LOG_SATD,
  FEATURE_INTER_HOR_1_PSNR,
  FEATURE_INTER_HOR_1_Q_COEFF_MAX,
  FEATURE_INTER_HOR_1_Q_COEFF_NONZ,
  FEATURE_INTER_HOR_1_LOG_MAG,
  FEATURE_INTER_HOR_1_ANGLE_RAD,
  FEATURE_INTER_HOR_1_LOG_SATDQ,
  FEATURE_INTER_HOR_1_LOG_SATD,
  FEATURE_INTER_VER_0_PSNR,
  FEATURE_INTER_VER_0_Q_COEFF_MAX,
  FEATURE_INTER_VER_0_Q_COEFF_NONZ,
  FEATURE_INTER_VER_0_LOG_MAG,
  FEATURE_INTER_VER_0_ANGLE_RAD,
  FEATURE_INTER_VER_0_LOG_SATDQ,
  FEATURE_INTER_VER_0_LOG_SATD,
  FEATURE_INTER_VER_1_PSNR,
  FEATURE_INTER_VER_1_Q_COEFF_MAX,
  FEATURE_INTER_VER_1_Q_COEFF_NONZ,
  FEATURE_INTER_VER_1_LOG_MAG,
  FEATURE_INTER_VER_1_ANGLE_RAD,
  FEATURE_INTER_VER_1_LOG_SATDQ,
  FEATURE_INTER_VER_1_LOG_SATD,

  // V1 features
  // NONE w RECT switch, search_none_after_rect
  FEATURE_INTER_SWITCH,
  // partition type: mixed, luma, chroma
  FEATURE_INTER_PART_T,

  FEATURE_INTER_MAX
};

enum { PT_INVAL = -1, PT_NONE = 0, PT_SPLIT };

struct InputNorm {
  bool valid;
  double *mean;
  double *std;
  double *invstd;
};

int av2_part_prune_tflite_exec(void **context, const float *ml_input,
                               float *ml_output, MODEL_TYPE model_type);
void av2_part_prune_tflite_close(void **context);
int av2_model_input_norm(MODEL_TYPE model_type, struct InputNorm *norm);
int get_model_part_type(MODEL_TYPE type);

#ifdef __cplusplus
}
#endif

#endif  // AV1_ENCODER_SIMPLE_INTRAPRED_TFLITE_H_