| /* |
| * Copyright (c) 2020, 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_AV1_COMMON_CNN_TFLITE_H |
| #define AOM_AV1_COMMON_CNN_TFLITE_H |
| |
| #ifdef __cplusplus |
| extern "C" { |
| #endif |
| |
| #include "av1/common/av1_common_int.h" |
| #include "av1/common/resize.h" |
| |
| #include "av1/encoder/rd.h" |
| #include "av1/encoder/ratectrl.h" |
| #if CONFIG_CNN_GUIDED_QUADTREE |
| #include "av1/common/guided_quadtree.h" |
| #endif // CONFIG_CNN_GUIDED_QUADTREE |
| |
| // Minimum base_qindex needed to run cnn. |
| #define MIN_CNN_Q_INDEX 0 |
| // Number of options for Y channel CNN model, with chosen option signaled in the |
| // bitstream as an index. |
| #define CNN_INDICES_Y 3 |
| // Number of bits required to signal combined code for cnn_index |
| // (CNN_INDICES_Y) and whether CNN is used or not (1). |
| #define CNN_INDEX_Y_BITS 2 // log2(CNN_INDICES_Y + 1) |
| #if CONFIG_EXT_SUPERRES |
| #define SELECT_CNN_FOR_SUPERRES 1 |
| #endif // CONFIG_EXT_SUPERRES |
| |
| // Returns true if we are allowed to use CNN for restoration for the frame. |
| static INLINE bool av1_allow_cnn(const AV1_COMMON *cm) { |
| return (cm->quant_params.base_qindex > MIN_CNN_Q_INDEX) |
| #if !CONFIG_EXT_SUPERRES |
| && !av1_superres_scaled(cm) |
| #endif // !CONFIG_EXT_SUPERRES |
| ; |
| } |
| |
| // Returns true if we are allowed to use CNN for restoration for the plane. |
| static INLINE bool av1_allow_cnn_for_plane(const AV1_COMMON *cm, int plane) { |
| (void)plane; |
| if (!av1_allow_cnn(cm)) return false; |
| if (plane >= av1_num_planes(cm)) return false; |
| if (av1_superres_scaled(cm)) { |
| #if CONFIG_EXT_SUPERRES |
| return (plane == 0); |
| #else |
| return false; |
| #endif // CONFIG_EXT_SUPERRES |
| } else { |
| return (plane == 0); |
| } |
| } |
| |
| static INLINE int av1_num_cnn_indices_for_plane(const AV1_COMMON *cm, |
| int plane) { |
| if (!av1_allow_cnn_for_plane(cm, plane)) return 0; |
| #if CONFIG_EXT_SUPERRES && SELECT_CNN_FOR_SUPERRES |
| if (av1_superres_scaled(cm)) return plane == AOM_PLANE_Y ? CNN_INDICES_Y : 1; |
| #else // CONFIG_EXT_SUPERRES && SELECT_CNN_FOR_SUPERRES |
| if (av1_superres_scaled(cm)) return 1; |
| #endif // CONFIG_EXT_SUPERRES && SELECT_CNN_FOR_SUPERRES |
| if (plane == AOM_PLANE_Y) { |
| return CNN_INDICES_Y; |
| } else { |
| return 1; |
| } |
| } |
| |
| static INLINE int av1_num_cnn_combined_code_bits(int num_indices) { |
| assert(num_indices == CNN_INDICES_Y); |
| (void)num_indices; |
| return CNN_INDEX_Y_BITS; |
| } |
| |
| // Returns true if we are allowed to use CNN for restoration. |
| static INLINE bool av1_use_cnn_encode(const AV1_COMMON *cm, |
| FRAME_UPDATE_TYPE update_type) { |
| const bool is_overlay_update = |
| (update_type == OVERLAY_UPDATE || update_type == INTNL_OVERLAY_UPDATE); |
| |
| return av1_allow_cnn(cm) && !is_overlay_update; |
| } |
| |
| // Restores image in 'dgd' with a CNN model using TFlite and stores output in |
| // 'rst'. Returns true on success. |
| int av1_restore_cnn_img_tflite_highbd(int qindex, int superres_denom, |
| const uint16_t *dgd, int width, |
| int height, int dgd_stride, uint16_t *rst, |
| int rst_stride, int num_threads, |
| int bit_depth, int is_intra_only, |
| int is_luma, int cnn_index); |
| |
| struct AV1Common; |
| |
| // Restore current frame buffer in 'cm' in-place with a CNN model using TFlite. |
| // Apply CNN to plane 'p' if and only if apply_cnn[p] is true. |
| void av1_restore_cnn_tflite(const struct AV1Common *cm, int num_threads, |
| const int apply_cnn[MAX_MB_PLANE], |
| const int cnn_indices[MAX_MB_PLANE]); |
| |
| #if CONFIG_CNN_GUIDED_QUADTREE |
| |
| // Apply Guided CNN restoration on encoder side. |
| int av1_restore_cnn_quadtree_encode_tflite( |
| struct AV1Common *cm, YV12_BUFFER_CONFIG *source_frame, int RDMULT, |
| int *quad_split_costs, int *binary_split_costs, int (*norestorecosts)[2], |
| int num_threads, const int apply_cnn[MAX_MB_PLANE], |
| const int cnn_indices[MAX_MB_PLANE], QUADInfo *quad_info, double *rdcost); |
| |
| // Apply Guided CNN restoration on decoder side. |
| int av1_restore_cnn_quadtree_decode_tflite(struct AV1Common *cm, |
| int num_threads, int use_quadtree, |
| const int apply_cnn[MAX_MB_PLANE], |
| const int cnn_indices[MAX_MB_PLANE]); |
| #endif // CONFIG_CNN_GUIDED_QUADTREE |
| |
| #ifdef __cplusplus |
| } |
| #endif |
| |
| #endif // AOM_AV1_COMMON_CNN_TFLITE_H |