| /* |
| * 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. |
| */ |
| |
| #include <vector> |
| |
| #include "av1/common/cnn_tflite.h" |
| #include "av1/common/onyxc_int.h" |
| #include "av1/tflite_models/op_registrations.h" |
| #include "av1/tflite_models/intra_frame_model/uv_qp22.h" |
| #include "av1/tflite_models/intra_frame_model/uv_qp32.h" |
| #include "av1/tflite_models/intra_frame_model/uv_qp43.h" |
| #include "av1/tflite_models/intra_frame_model/uv_qp53.h" |
| #include "av1/tflite_models/intra_frame_model/uv_qp63.h" |
| #include "av1/tflite_models/intra_frame_model/qp22.h" |
| #include "av1/tflite_models/intra_frame_model/qp32.h" |
| #include "av1/tflite_models/intra_frame_model/qp43.h" |
| #include "av1/tflite_models/intra_frame_model/qp53.h" |
| #include "av1/tflite_models/intra_frame_model/qp63.h" |
| #include "av1/tflite_models/inter_frame_model/uv_qp68_107.h" |
| #include "av1/tflite_models/inter_frame_model/uv_qp108_147.h" |
| #include "av1/tflite_models/inter_frame_model/uv_qp148_191.h" |
| #include "av1/tflite_models/inter_frame_model/uv_qp192_231.h" |
| #include "av1/tflite_models/inter_frame_model/uv_qp232_255.h" |
| #include "av1/tflite_models/inter_frame_model/qp68_107.h" |
| #include "av1/tflite_models/inter_frame_model/qp108_147.h" |
| #include "av1/tflite_models/inter_frame_model/qp148_191.h" |
| #include "av1/tflite_models/inter_frame_model/qp192_231.h" |
| #include "av1/tflite_models/inter_frame_model/qp232_255.h" |
| |
| #if CONFIG_CNN_CRLC_GUIDED |
| #include "av1/tflite_models/crlc_model/qp12_crlc.h" |
| #include "av1/tflite_models/crlc_model/qp22_crlc.h" |
| #include "av1/tflite_models/crlc_model/qp28_crlc.h" |
| #include "av1/tflite_models/crlc_model/qp33_crlc.h" |
| #include "av1/tflite_models/crlc_model/qp43_crlc.h" |
| #include "av1/tflite_models/crlc_model/qp53_crlc.h" |
| #include "av1/tflite_models/crlc_model/qp63_crlc.h" |
| #endif // CONFIG_CNN_CRLC_GUIDED |
| |
| #if CONFIG_NN_RECON |
| #include "av1/tflite_models/intra_txfm_recon_model/tx16x16.h" |
| #endif // CONFIG_NN_RECON |
| #include "common/tf_lite_includes.h" |
| |
| #if CONFIG_CNN_RESTORATION || CONFIG_LOOP_RESTORE_CNN |
| // Returns the TF-lite model based on the qindex. |
| static const unsigned char *get_intra_model_from_qindex(int qindex, |
| int is_luma) { |
| if (qindex <= MIN_CNN_Q_INDEX) { |
| assert(0); |
| return nullptr; |
| } |
| |
| if (is_luma) { |
| #if CONFIG_CNN_CRLC_GUIDED |
| int QP = qindex / 4; |
| if (QP < 17) { |
| return qp12_crlc_model_tflite_data; |
| } else if (QP < 27) { |
| return qp22_crlc_model_tflite_data; |
| } else if (QP < 31) { |
| return qp28_crlc_model_tflite_data; |
| } else if (QP < 37) { |
| return qp33_crlc_model_tflite_data; |
| } else if (QP < 47) { |
| return qp43_crlc_model_tflite_data; |
| } else if (QP < 57) { |
| return qp53_crlc_model_tflite_data; |
| } else { |
| return qp63_crlc_model_tflite_data; |
| } |
| #else |
| if (qindex < 108) { |
| return qp22_model_tflite_data; |
| } else if (qindex < 148) { |
| return qp32_model_tflite_data; |
| } else if (qindex < 192) { |
| return qp43_model_tflite_data; |
| } else if (qindex < 232) { |
| return qp53_model_tflite_data; |
| } else { |
| return qp63_model_tflite_data; |
| } |
| #endif // CONFIG_CNN_CRLC_GUIDED |
| } else { |
| if (qindex < 108) { |
| return uv_qp22_model_tflite_data; |
| } else if (qindex < 148) { |
| return uv_qp32_model_tflite_data; |
| } else if (qindex < 192) { |
| return uv_qp43_model_tflite_data; |
| } else if (qindex < 232) { |
| return uv_qp53_model_tflite_data; |
| } else { |
| return uv_qp63_model_tflite_data; |
| } |
| } |
| } |
| |
| // Returns the TF-lite model based on the qindex. |
| static const unsigned char *get_inter_model_from_qindex(int qindex, |
| int is_luma) { |
| if (qindex <= MIN_CNN_Q_INDEX) { |
| assert(0); |
| return nullptr; |
| } |
| |
| if (is_luma) { |
| #if CONFIG_CNN_CRLC_GUIDED |
| int QP = qindex / 4; |
| if (QP < 17) { |
| return qp12_crlc_model_tflite_data; |
| } else if (QP < 27) { |
| return qp22_crlc_model_tflite_data; |
| } else if (QP < 31) { |
| return qp28_crlc_model_tflite_data; |
| } else if (QP < 37) { |
| return qp33_crlc_model_tflite_data; |
| } else if (QP < 47) { |
| return qp43_crlc_model_tflite_data; |
| } else if (QP < 57) { |
| return qp53_crlc_model_tflite_data; |
| } else { |
| return qp63_crlc_model_tflite_data; |
| } |
| #else |
| if (qindex < 108) { |
| return qp68_107_inter_model_tflite_data; |
| } else if (qindex < 148) { |
| return qp108_147_inter_model_tflite_data; |
| } else if (qindex < 192) { |
| return qp148_191_inter_model_tflite_data; |
| } else if (qindex < 232) { |
| return qp192_231_inter_model_tflite_data; |
| } else { |
| return qp232_255_inter_model_tflite_data; |
| } |
| #endif // CONFIG_CNN_CRLC_GUIDED |
| } else { |
| if (qindex < 108) { |
| return uv_qp68_107_inter_model_tflite_data; |
| } else if (qindex < 148) { |
| return uv_qp108_147_inter_model_tflite_data; |
| } else if (qindex < 192) { |
| return uv_qp148_191_inter_model_tflite_data; |
| } else if (qindex < 232) { |
| return uv_qp192_231_inter_model_tflite_data; |
| } else { |
| return uv_qp232_255_inter_model_tflite_data; |
| } |
| } |
| } |
| |
| static TfLiteDelegate *get_tflite_xnnpack_delegate(int num_threads) { |
| TfLiteXNNPackDelegateOptions xnnpack_options = |
| TfLiteXNNPackDelegateOptionsDefault(); |
| xnnpack_options.num_threads = AOMMAX(num_threads, 1); |
| return TfLiteXNNPackDelegateCreate(&xnnpack_options); |
| } |
| |
| // Builds and returns the TFlite interpreter. |
| static std::unique_ptr<tflite::Interpreter> get_tflite_interpreter( |
| int qindex, int width, int height, int num_threads, int is_intra_only, |
| int is_luma, TfLiteDelegate *xnnpack_delegate) { |
| const unsigned char *const model_tflite_data = |
| is_intra_only ? get_intra_model_from_qindex(qindex, is_luma) |
| : get_inter_model_from_qindex(qindex, is_luma); |
| auto model = tflite::GetModel(model_tflite_data); |
| tflite::MutableOpResolver resolver; |
| RegisterSelectedOpsAllQps(&resolver); |
| tflite::InterpreterBuilder builder(model, resolver); |
| // TODO(urvang): Investigate if caching the interpreter object provides |
| // further speed-up. May still have to re-build the interpreter if qindex |
| // changes. |
| std::unique_ptr<tflite::Interpreter> interpreter; |
| builder(&interpreter); |
| interpreter->SetNumThreads(AOMMAX(num_threads, 1)); |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| |
| // Dimension order: batch_size, height, width, num_channels. |
| // Note: height comes before width here! |
| const std::vector<int> in_out_dims = { 1, height, width, 1 }; |
| // We only need to resize the input tensor. All other tensors (including |
| // output tensor) will be resized automatically. |
| if (interpreter->ResizeInputTensor(interpreter->inputs()[0], in_out_dims) != |
| kTfLiteOk) { |
| reporter->Report("Failed at input tensor resize"); |
| return nullptr; |
| } |
| |
| if (interpreter->AllocateTensors() != kTfLiteOk) { |
| reporter->Report("Failed at tensor allocation"); |
| return nullptr; |
| } |
| |
| if (interpreter->ModifyGraphWithDelegate(xnnpack_delegate) != kTfLiteOk) { |
| reporter->Report("Failed at modifying graph with XNNPack delegate"); |
| return nullptr; |
| } |
| |
| return interpreter; |
| } |
| |
| extern "C" int av1_restore_cnn_img_tflite(int qindex, const uint8_t *dgd, |
| int width, int height, int dgd_stride, |
| uint8_t *rst, int rst_stride, |
| int num_threads, int is_intra_only, |
| int is_luma) { |
| TfLiteDelegate *xnnpack_delegate = get_tflite_xnnpack_delegate(num_threads); |
| std::unique_ptr<tflite::Interpreter> interpreter = |
| get_tflite_interpreter(qindex, width, height, num_threads, is_intra_only, |
| is_luma, xnnpack_delegate); |
| |
| // Prepare input. |
| const float max_val = 255.0f; |
| const int in_stride = width; |
| auto input = interpreter->typed_input_tensor<float>(0); |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| input[r * in_stride + c] = |
| static_cast<float>(dgd[r * dgd_stride + c]) / max_val; |
| assert(input[r * in_stride + c] >= 0.0f); |
| assert(input[r * in_stride + c] <= 1.0f); |
| } |
| } |
| |
| // Invoke TFlite inference. |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| auto status = interpreter->Invoke(); |
| if (status != kTfLiteOk) { |
| reporter->Report("Failed at interpreter invocation"); |
| return 0; |
| } |
| |
| // Use the output to restore 'dgd' and store in 'rst'. |
| const auto output = interpreter->typed_output_tensor<float>(0); |
| const int out_stride = width; |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| const int residue = |
| static_cast<int>(output[r * out_stride + c] * max_val + 0.5); |
| rst[r * rst_stride + c] = clip_pixel(dgd[r * dgd_stride + c] + residue); |
| } |
| } |
| |
| // IMPORTANT: release the interpreter before destroying the delegate. |
| interpreter.reset(); |
| TfLiteXNNPackDelegateDelete(xnnpack_delegate); |
| |
| return 1; |
| } |
| |
| extern "C" int av1_restore_cnn_img_tflite_highbd( |
| int qindex, 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) { |
| TfLiteDelegate *xnnpack_delegate = get_tflite_xnnpack_delegate(num_threads); |
| std::unique_ptr<tflite::Interpreter> interpreter = |
| get_tflite_interpreter(qindex, width, height, num_threads, is_intra_only, |
| is_luma, xnnpack_delegate); |
| |
| // Prepare input. |
| const auto max_val = static_cast<float>((1 << bit_depth) - 1); |
| const int in_stride = width; |
| auto input = interpreter->typed_input_tensor<float>(0); |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| input[r * in_stride + c] = |
| static_cast<float>(dgd[r * dgd_stride + c]) / max_val; |
| assert(input[r * in_stride + c] >= 0.0f); |
| assert(input[r * in_stride + c] <= 1.0f); |
| } |
| } |
| |
| // Invoke TFlite inference. |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| auto status = interpreter->Invoke(); |
| if (status != kTfLiteOk) { |
| reporter->Report("Failed at interpreter invocation"); |
| return 0; |
| } |
| |
| // Use the output to restore 'dgd' and store in 'rst'. |
| const auto output = interpreter->typed_output_tensor<float>(0); |
| const int out_stride = width; |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| const int residue = |
| static_cast<int>(output[r * out_stride + c] * max_val + 0.5); |
| rst[r * rst_stride + c] = |
| clip_pixel_highbd(dgd[r * dgd_stride + c] + residue, bit_depth); |
| } |
| } |
| |
| // IMPORTANT: release the interpreter before destroying the delegate. |
| interpreter.reset(); |
| TfLiteXNNPackDelegateDelete(xnnpack_delegate); |
| |
| return 1; |
| } |
| |
| extern "C" void av1_restore_cnn_tflite(const AV1_COMMON *cm, int num_threads, |
| int plane_from, int plane_to) { |
| YV12_BUFFER_CONFIG *buf = &cm->cur_frame->buf; |
| const int is_intra_only = frame_is_intra_only(cm); |
| for (int plane = plane_from; plane <= plane_to; ++plane) { |
| const int is_luma = (plane == AOM_PLANE_Y); |
| if (cm->seq_params.use_highbitdepth) { |
| switch (plane) { |
| case AOM_PLANE_Y: |
| av1_restore_cnn_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->y_buffer), |
| buf->y_crop_width, buf->y_crop_height, buf->y_stride, |
| CONVERT_TO_SHORTPTR(buf->y_buffer), buf->y_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma); |
| break; |
| case AOM_PLANE_U: |
| av1_restore_cnn_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->u_buffer), |
| buf->uv_crop_width, buf->uv_crop_height, buf->uv_stride, |
| CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma); |
| break; |
| case AOM_PLANE_V: |
| av1_restore_cnn_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->v_buffer), |
| buf->uv_crop_width, buf->uv_crop_height, buf->uv_stride, |
| CONVERT_TO_SHORTPTR(buf->v_buffer), buf->uv_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma); |
| break; |
| default: assert(0 && "Invalid plane index"); |
| } |
| } else { |
| assert(cm->seq_params.bit_depth == 8); |
| switch (plane) { |
| case AOM_PLANE_Y: |
| av1_restore_cnn_img_tflite( |
| cm->base_qindex, buf->y_buffer, buf->y_crop_width, |
| buf->y_crop_height, buf->y_stride, buf->y_buffer, buf->y_stride, |
| num_threads, is_intra_only, is_luma); |
| break; |
| case AOM_PLANE_U: |
| av1_restore_cnn_img_tflite( |
| cm->base_qindex, buf->u_buffer, buf->uv_crop_width, |
| buf->uv_crop_height, buf->uv_stride, buf->u_buffer, |
| buf->uv_stride, num_threads, is_intra_only, is_luma); |
| break; |
| case AOM_PLANE_V: |
| av1_restore_cnn_img_tflite( |
| cm->base_qindex, buf->v_buffer, buf->uv_crop_width, |
| buf->uv_crop_height, buf->uv_stride, buf->v_buffer, |
| buf->uv_stride, num_threads, is_intra_only, is_luma); |
| break; |
| default: assert(0 && "Invalid plane index"); |
| } |
| } |
| } |
| } |
| #endif // CONFIG_CNN_RESTORATION || CONFIG_LOOP_RESTORE_CNN |
| |
| #if CONFIG_CNN_CRLC_GUIDED |
| |
| extern "C" int av1_restore_cnn_guided_img_tflite( |
| int qindex, const uint8_t *dgd, int width, int height, int dgd_stride, |
| uint8_t *rst, int rst_stride, int num_threads, int is_intra_only, |
| int is_luma, const uint8_t *src, int src_stride, CRLCInfo *ci, |
| int frameType) { |
| TfLiteDelegate *xnnpack_delegate = get_tflite_xnnpack_delegate(num_threads); |
| std::unique_ptr<tflite::Interpreter> interpreter = |
| get_tflite_interpreter(qindex, width, height, num_threads, is_intra_only, |
| is_luma, xnnpack_delegate); |
| |
| // Prepare input. |
| const float max_val = 255.0f; |
| const int in_stride = width; |
| auto input = interpreter->typed_input_tensor<float>(0); |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| input[r * in_stride + c] = |
| static_cast<float>(dgd[r * dgd_stride + c]) / max_val; |
| assert(input[r * in_stride + c] >= 0.0f); |
| assert(input[r * in_stride + c] <= 1.0f); |
| } |
| } |
| |
| // Invoke TFlite inference. |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| auto status = interpreter->Invoke(); |
| if (status != kTfLiteOk) { |
| reporter->Report("Failed at interpreter invocation"); |
| return 0; |
| } |
| |
| // Use the output to restore 'dgd' and store in 'rst'. |
| const auto output = interpreter->typed_output_tensor<float>(0); |
| const int out_stride = width; |
| |
| uint8_t **sub_dgr = new uint8_t *[height]; |
| for (int i = 0; i < height; i++) { |
| sub_dgr[i] = new uint8_t[width]; |
| } |
| uint8_t **sub_src = new uint8_t *[height]; |
| for (int i = 0; i < height; i++) { |
| sub_src[i] = new uint8_t[width]; |
| } |
| int **sub_r = new int *[height]; |
| for (int i = 0; i < height; i++) { |
| sub_r[i] = new int[width]; |
| } |
| // channel 0 |
| double **r0 = new double *[height]; |
| for (int i = 0; i < height; i++) { |
| r0[i] = new double[width]; |
| } |
| // channel 1 |
| double **r1 = new double *[height]; |
| for (int i = 0; i < height; i++) { |
| r1[i] = new double[width]; |
| } |
| |
| uint8_t **repic = new uint8_t *[height]; |
| for (int i = 0; i < height; i++) { |
| repic[i] = new uint8_t[width]; |
| } |
| |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| // reconstruct image |
| sub_dgr[r][c] = dgd[r * dgd_stride + c]; |
| // src img |
| sub_src[r][c] = src[r * src_stride + c]; |
| // src img-reconstruct image |
| sub_r[r][c] = sub_src[r][c] - sub_dgr[r][c]; |
| // from tflite get channel 0 |
| r0[r][c] = output[r * 2 * out_stride + c * 2] * max_val; |
| // from tflite get channel 1 |
| r1[r][c] = output[r * 2 * out_stride + c * 2 + 1] * max_val; |
| } |
| } |
| int scale, A0_min, A1_min; |
| int qp = qindex / 4; |
| if (qp < 17) { |
| scale = 16384; |
| A0_min = -7; |
| A1_min = -5; |
| } else if (17 <= qp && qp < 27) { |
| scale = 16384; |
| A0_min = -12; |
| A1_min = -7; |
| } else if (27 <= qp && qp < 31) { |
| scale = 8192; |
| A0_min = -12; |
| A1_min = -3; |
| } else if (31 <= qp && qp < 37) { |
| scale = 8192; |
| A0_min = -13; |
| A1_min = -10; |
| } else if (37 <= qp && qp < 47) { |
| scale = 4192; |
| A0_min = -13; |
| A1_min = -10; |
| } else if (47 <= qp && qp < 57) { |
| scale = 2046; |
| A0_min = -13; |
| A1_min = -10; |
| } else if (qp > 56) { |
| scale = 2046; |
| A0_min = -15; |
| A1_min = -6; |
| } |
| |
| int blockSize = frameType; |
| int cols = int(ceil(double(height) / blockSize)); |
| int rows = int(ceil(double(width) / blockSize)); |
| int number_crlc = cols * rows; |
| int *A = new int[(int)number_crlc * 2]; |
| int index_A = 0; |
| int start_row = 0; |
| int end_row = 0; |
| int start_clow = 0; |
| int end_clow = 0; |
| int testnum = 10; |
| for (int i = 0; i < cols; i++) { |
| for (int j = 0; j < rows; j++) { |
| if (i == cols - 1) { |
| start_clow = height - blockSize; |
| end_clow = height; |
| } else { |
| start_clow = i * blockSize; |
| end_clow = (i + 1) * blockSize; |
| } |
| if (j == rows - 1) { |
| start_row = width - blockSize; |
| end_row = width; |
| } else { |
| start_row = j * blockSize; |
| end_row = (j + 1) * blockSize; |
| } |
| if (width < blockSize) { |
| start_row = 0; |
| end_row = width; |
| } |
| if (height < blockSize) { |
| start_clow = 0; |
| end_clow = height; |
| } |
| int lenth_clows = end_clow - start_clow; |
| int lenth_rows = end_row - start_row; |
| |
| int lenth = lenth_clows * lenth_rows; |
| int *sub_r_flatten = new int[lenth]; |
| int k = 0; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| sub_r_flatten[k] = sub_r[i][j]; |
| k = k + 1; |
| } |
| } |
| |
| double *sub_r0 = new double[lenth]; |
| |
| int k_r0 = 0; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| sub_r0[k_r0] = r0[i][j]; |
| k_r0++; |
| } |
| } |
| |
| double *sub_r1 = new double[lenth]; |
| int k_r1 = 0; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| sub_r1[k_r1] = r1[i][j]; |
| k_r1++; |
| } |
| } |
| |
| double **R = new double *[lenth]; |
| for (int i = 0; i < lenth; i++) { |
| R[i] = new double[2]; |
| } |
| |
| for (int i = 0; i < lenth; i++) { |
| for (int j = 0; j < 2; j++) { |
| if (j == 0) { |
| R[i][j] = sub_r0[i]; |
| } |
| if (j == 1) { |
| R[i][j] = sub_r1[i]; |
| } |
| } |
| } |
| |
| double **R_T = new double *[2]; |
| for (int i = 0; i < 2; i++) { |
| R_T[i] = new double[lenth]; |
| } |
| |
| for (int i = 0; i < 2; i++) { |
| for (int j = 0; j < lenth; j++) { |
| if (i == 0) { |
| R_T[i][j] = sub_r0[j]; |
| } |
| if (i == 1) { |
| R_T[i][j] = sub_r1[j]; |
| } |
| } |
| } |
| |
| double **R_TDotR = new double *[2]; |
| for (int i = 0; i < 2; i++) { |
| R_TDotR[i] = new double[2]; |
| } |
| |
| R_TDotR[0][0] = 0; |
| for (int i = 0; i < lenth; i++) { |
| R_TDotR[0][0] += R_T[0][i] * R[i][0]; |
| } |
| R_TDotR[0][1] = 0; |
| for (int i = 0; i < lenth; i++) { |
| R_TDotR[0][1] += R_T[0][i] * R[i][1]; |
| } |
| R_TDotR[1][0] = 0; |
| for (int i = 0; i < lenth; i++) { |
| R_TDotR[1][0] += R_T[1][i] * R[i][0]; |
| } |
| R_TDotR[1][1] = 0; |
| for (int i = 0; i < lenth; i++) { |
| R_TDotR[1][1] += R_T[1][i] * R[i][1]; |
| } |
| |
| double value_R_TDotR = |
| R_TDotR[0][0] * R_TDotR[1][1] - R_TDotR[0][1] * R_TDotR[1][0]; |
| double a00 = R_TDotR[1][1] / value_R_TDotR; |
| double a01 = -1 * R_TDotR[0][1] / value_R_TDotR; |
| double a10 = -1 * R_TDotR[1][0] / value_R_TDotR; |
| double a11 = R_TDotR[0][0] / value_R_TDotR; |
| |
| double **R_TDotR_inver = new double *[2]; |
| for (int i = 0; i < 2; i++) { |
| R_TDotR_inver[i] = new double[2]; |
| } |
| R_TDotR_inver[0][0] = a00; |
| R_TDotR_inver[0][1] = a01; |
| R_TDotR_inver[1][0] = a10; |
| R_TDotR_inver[1][1] = a11; |
| |
| double **mid = new double *[2]; |
| for (int i = 0; i < 2; i++) { |
| mid[i] = new double[lenth]; |
| } |
| for (int i = 0; i < 2; i++) { |
| for (int j = 0; j < lenth; j++) { |
| if (i == 0) { |
| mid[i][j] = R_TDotR_inver[0][0] * R_T[0][j] + |
| R_TDotR_inver[0][1] * R_T[1][j]; |
| } |
| if (i == 1) { |
| mid[i][j] = R_TDotR_inver[1][0] * R_T[0][j] + |
| R_TDotR_inver[1][1] * R_T[1][j]; |
| } |
| } |
| } |
| |
| double A0 = 0; |
| double A1 = 0; |
| for (int i = 0; i < lenth; i++) { |
| A0 += mid[0][i] * sub_r_flatten[i]; |
| A1 += mid[1][i] * sub_r_flatten[i]; |
| } |
| A0 = A0 * scale; |
| A1 = A1 * scale; |
| |
| A0 = int(round(A0)); |
| A1 = int(round(A1)); |
| if (A0 < A0_min) { |
| A0 = A0_min; |
| } |
| if (A0 > A0_min + 15) { |
| A0 = A0_min + 15; |
| } |
| A[index_A] = int(A0); |
| index_A = index_A + 1; |
| if (A1 < A1_min) { |
| A1 = A1_min; |
| } |
| if (A1 > A1_min + 15) { |
| A1 = A1_min + 15; |
| } |
| A[index_A] = int(A1); |
| index_A = index_A + 1; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| repic[i][j] = int(round(sub_dgr[i][j] + A0 * r0[i][j] / scale + |
| A1 * r1[i][j] / scale)); |
| repic[i][j] = clip_pixel(repic[i][j]); |
| } |
| } |
| } |
| } |
| ci->num_crlc_unit = (int)number_crlc; |
| for (int i = 0; i < number_crlc * 2; i++) { |
| if (i % 2 == 0) { |
| if (A[i] < A0_min) { |
| A[i] = A0_min; |
| } |
| if (A[i] > A0_min + 15) { |
| A[i] = A0_min + 15; |
| } |
| } else { |
| if (A[i] < A1_min) { |
| A[i] = A1_min; |
| } |
| if (A[i] > A1_min + 15) { |
| A[i] = A1_min + 15; |
| } |
| } |
| ci->unit_info[i / 2].xqd[i % 2] = (int)A[i]; |
| } |
| |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| rst[r * rst_stride + c] = clip_pixel(repic[r][c]); |
| } |
| } |
| return 1; |
| } |
| |
| extern "C" int av1_restore_cnn_guided_img_tflite_highbd( |
| int qindex, 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, const uint16_t *src, int src_stride, |
| CRLCInfo *ci, int frameType) { |
| TfLiteDelegate *xnnpack_delegate = get_tflite_xnnpack_delegate(num_threads); |
| std::unique_ptr<tflite::Interpreter> interpreter = |
| get_tflite_interpreter(qindex, width, height, num_threads, is_intra_only, |
| is_luma, xnnpack_delegate); |
| // Prepare input. |
| const auto max_val = static_cast<float>((1 << bit_depth) - 1); |
| const int in_stride = width; |
| auto input = interpreter->typed_input_tensor<float>(0); |
| |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| input[r * in_stride + c] = |
| static_cast<float>(dgd[r * dgd_stride + c]) / max_val; |
| assert(input[r * in_stride + c] >= 0.0f); |
| assert(input[r * in_stride + c] <= 1.0f); |
| } |
| } |
| |
| // Invoke TFlite inference. |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| auto status = interpreter->Invoke(); |
| if (status != kTfLiteOk) { |
| reporter->Report("Failed at interpreter invocation"); |
| return 0; |
| } |
| |
| // Use the output to restore 'dgd' and store in 'rst'. |
| const auto output = interpreter->typed_output_tensor<float>(0); |
| const int out_stride = width; |
| |
| uint16_t **sub_dgr = new uint16_t *[height]; |
| for (int i = 0; i < height; i++) { |
| sub_dgr[i] = new uint16_t[width]; |
| } |
| uint16_t **sub_src = new uint16_t *[height]; |
| for (int i = 0; i < height; i++) { |
| sub_src[i] = new uint16_t[width]; |
| } |
| int **sub_r = new int *[height]; |
| for (int i = 0; i < height; i++) { |
| sub_r[i] = new int[width]; |
| } |
| // channel 0 |
| double **r0 = new double *[height]; |
| for (int i = 0; i < height; i++) { |
| r0[i] = new double[width]; |
| } |
| // channel 1 |
| double **r1 = new double *[height]; |
| for (int i = 0; i < height; i++) { |
| r1[i] = new double[width]; |
| } |
| |
| uint16_t **repic = new uint16_t *[height]; |
| for (int i = 0; i < height; i++) { |
| repic[i] = new uint16_t[width]; |
| } |
| |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| // reconstruct image |
| sub_dgr[r][c] = dgd[r * dgd_stride + c]; |
| // src img |
| sub_src[r][c] = src[r * src_stride + c]; |
| // src img-reconstruct image |
| sub_r[r][c] = sub_src[r][c] - sub_dgr[r][c]; |
| // from tflite get channel 0 |
| r0[r][c] = output[r * 2 * out_stride + c * 2] * max_val; |
| // from tflite get channel 1 |
| r1[r][c] = output[r * 2 * out_stride + c * 2 + 1] * max_val; |
| } |
| } |
| int scale, A0_min, A1_min; |
| int qp = qindex / 4; |
| if (qp < 17) { |
| scale = 16384; |
| A0_min = -7; |
| A1_min = -5; |
| } else if (17 <= qp && qp < 27) { |
| scale = 16384; |
| A0_min = -12; |
| A1_min = -7; |
| } else if (27 <= qp && qp < 31) { |
| scale = 8192; |
| A0_min = -12; |
| A1_min = -3; |
| } else if (31 <= qp && qp < 37) { |
| scale = 8192; |
| A0_min = -13; |
| A1_min = -10; |
| } else if (37 <= qp && qp < 47) { |
| scale = 4192; |
| A0_min = -13; |
| A1_min = -10; |
| } else if (47 <= qp && qp < 57) { |
| scale = 2046; |
| A0_min = -13; |
| A1_min = -10; |
| } else if (qp > 56) { |
| scale = 2046; |
| A0_min = -15; |
| A1_min = -6; |
| } |
| int blockSize = frameType; |
| int cols = int(ceil(double(height) / blockSize)); |
| int rows = int(ceil(double(width) / blockSize)); |
| int number_crlc = cols * rows; |
| int *A = new int[(int)number_crlc * 2]; |
| int index_A = 0; |
| int start_row = 0; |
| int end_row = 0; |
| int start_clow = 0; |
| int end_clow = 0; |
| int testnum = 10; |
| for (int i = 0; i < cols; i++) { |
| for (int j = 0; j < rows; j++) { |
| if (i == cols - 1) { |
| start_clow = height - blockSize; |
| end_clow = height; |
| } else { |
| start_clow = i * blockSize; |
| end_clow = (i + 1) * blockSize; |
| } |
| if (j == rows - 1) { |
| start_row = width - blockSize; |
| end_row = width; |
| } else { |
| start_row = j * blockSize; |
| end_row = (j + 1) * blockSize; |
| } |
| if (width < blockSize) { |
| start_row = 0; |
| end_row = width; |
| } |
| if (height < blockSize) { |
| start_clow = 0; |
| end_clow = height; |
| } |
| int lenth_clows = end_clow - start_clow; |
| int lenth_rows = end_row - start_row; |
| |
| int lenth = lenth_clows * lenth_rows; |
| int *sub_r_flatten = new int[lenth]; |
| int k = 0; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| sub_r_flatten[k] = sub_r[i][j]; |
| k = k + 1; |
| } |
| } |
| |
| double *sub_r0 = new double[lenth]; |
| |
| int k_r0 = 0; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| sub_r0[k_r0] = r0[i][j]; |
| k_r0++; |
| } |
| } |
| |
| double *sub_r1 = new double[lenth]; |
| int k_r1 = 0; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| sub_r1[k_r1] = r1[i][j]; |
| k_r1++; |
| } |
| } |
| |
| double **R = new double *[lenth]; |
| for (int i = 0; i < lenth; i++) { |
| R[i] = new double[2]; |
| } |
| |
| for (int i = 0; i < lenth; i++) { |
| for (int j = 0; j < 2; j++) { |
| if (j == 0) { |
| R[i][j] = sub_r0[i]; |
| } |
| if (j == 1) { |
| R[i][j] = sub_r1[i]; |
| } |
| } |
| } |
| |
| double **R_T = new double *[2]; |
| for (int i = 0; i < 2; i++) { |
| R_T[i] = new double[lenth]; |
| } |
| |
| for (int i = 0; i < 2; i++) { |
| for (int j = 0; j < lenth; j++) { |
| if (i == 0) { |
| R_T[i][j] = sub_r0[j]; |
| } |
| if (i == 1) { |
| R_T[i][j] = sub_r1[j]; |
| } |
| } |
| } |
| |
| double **R_TDotR = new double *[2]; |
| for (int i = 0; i < 2; i++) { |
| R_TDotR[i] = new double[2]; |
| } |
| |
| R_TDotR[0][0] = 0; |
| for (int i = 0; i < lenth; i++) { |
| R_TDotR[0][0] += R_T[0][i] * R[i][0]; |
| } |
| R_TDotR[0][1] = 0; |
| for (int i = 0; i < lenth; i++) { |
| R_TDotR[0][1] += R_T[0][i] * R[i][1]; |
| } |
| R_TDotR[1][0] = 0; |
| for (int i = 0; i < lenth; i++) { |
| R_TDotR[1][0] += R_T[1][i] * R[i][0]; |
| } |
| R_TDotR[1][1] = 0; |
| for (int i = 0; i < lenth; i++) { |
| R_TDotR[1][1] += R_T[1][i] * R[i][1]; |
| } |
| |
| double value_R_TDotR = |
| R_TDotR[0][0] * R_TDotR[1][1] - R_TDotR[0][1] * R_TDotR[1][0]; |
| double a00 = R_TDotR[1][1] / value_R_TDotR; |
| double a01 = -1 * R_TDotR[0][1] / value_R_TDotR; |
| double a10 = -1 * R_TDotR[1][0] / value_R_TDotR; |
| double a11 = R_TDotR[0][0] / value_R_TDotR; |
| |
| double **R_TDotR_inver = new double *[2]; |
| for (int i = 0; i < 2; i++) { |
| R_TDotR_inver[i] = new double[2]; |
| } |
| R_TDotR_inver[0][0] = a00; |
| R_TDotR_inver[0][1] = a01; |
| R_TDotR_inver[1][0] = a10; |
| R_TDotR_inver[1][1] = a11; |
| |
| double **mid = new double *[2]; |
| for (int i = 0; i < 2; i++) { |
| mid[i] = new double[lenth]; |
| } |
| for (int i = 0; i < 2; i++) { |
| for (int j = 0; j < lenth; j++) { |
| if (i == 0) { |
| mid[i][j] = R_TDotR_inver[0][0] * R_T[0][j] + |
| R_TDotR_inver[0][1] * R_T[1][j]; |
| } |
| if (i == 1) { |
| mid[i][j] = R_TDotR_inver[1][0] * R_T[0][j] + |
| R_TDotR_inver[1][1] * R_T[1][j]; |
| } |
| } |
| } |
| |
| double A0 = 0; |
| double A1 = 0; |
| for (int i = 0; i < lenth; i++) { |
| A0 += mid[0][i] * sub_r_flatten[i]; |
| A1 += mid[1][i] * sub_r_flatten[i]; |
| } |
| A0 = A0 * scale; |
| A1 = A1 * scale; |
| A0 = int(round(A0)); |
| A1 = int(round(A1)); |
| if (A0 < A0_min) { |
| A0 = A0_min; |
| } |
| if (A0 > A0_min + 15) { |
| A0 = A0_min + 15; |
| } |
| A[index_A] = int(A0); |
| index_A = index_A + 1; |
| if (A1 < A1_min) { |
| A1 = A1_min; |
| } |
| if (A1 > A1_min + 15) { |
| A1 = A1_min + 15; |
| } |
| A[index_A] = int(A1); |
| index_A = index_A + 1; |
| |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| repic[i][j] = int(round(sub_dgr[i][j] + A0 * r0[i][j] / scale + |
| A1 * r1[i][j] / scale)); |
| repic[i][j] = clip_pixel_highbd(repic[i][j], bit_depth); |
| } |
| } |
| } |
| } |
| ci->num_crlc_unit = (int)number_crlc; |
| for (int i = 0; i < number_crlc * 2; i++) { |
| if (i % 2 == 0) { |
| if (A[i] < A0_min) { |
| A[i] = A0_min; |
| } |
| if (A[i] > A0_min + 15) { |
| A[i] = A0_min + 15; |
| } |
| } else { |
| if (A[i] < A1_min) { |
| A[i] = A1_min; |
| } |
| if (A[i] > A1_min + 15) { |
| A[i] = A1_min + 15; |
| } |
| } |
| ci->unit_info[i / 2].xqd[i % 2] = (int)A[i]; |
| } |
| |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| rst[r * rst_stride + c] = clip_pixel_highbd(repic[r][c], bit_depth); |
| } |
| } |
| return 1; |
| } |
| |
| extern "C" int av1_restore_cnn_guided_decode_img_tflite( |
| int qindex, const uint8_t *dgd, int width, int height, int dgd_stride, |
| uint8_t *rst, int rst_stride, int num_threads, int is_intra_only, |
| int is_luma, CRLCInfo *ci, int frameType) { |
| TfLiteDelegate *xnnpack_delegate = get_tflite_xnnpack_delegate(num_threads); |
| std::unique_ptr<tflite::Interpreter> interpreter = |
| get_tflite_interpreter(qindex, width, height, num_threads, is_intra_only, |
| is_luma, xnnpack_delegate); |
| |
| // Prepare input. |
| const float max_val = 255.0f; |
| const int in_stride = width; |
| auto input = interpreter->typed_input_tensor<float>(0); |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| input[r * in_stride + c] = |
| static_cast<float>(dgd[r * dgd_stride + c]) / max_val; |
| assert(input[r * in_stride + c] >= 0.0f); |
| assert(input[r * in_stride + c] <= 1.0f); |
| } |
| } |
| |
| // Invoke TFlite inference. |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| auto status = interpreter->Invoke(); |
| if (status != kTfLiteOk) { |
| reporter->Report("Failed at interpreter invocation"); |
| return 0; |
| } |
| |
| // Use the output to restore 'dgd' and store in 'rst'. |
| const auto output = interpreter->typed_output_tensor<float>(0); |
| const int out_stride = width; |
| |
| uint8_t **sub_dgr = new uint8_t *[height]; |
| for (int i = 0; i < height; i++) { |
| sub_dgr[i] = new uint8_t[width]; |
| } |
| |
| double **r0 = new double *[height]; |
| for (int i = 0; i < height; i++) { |
| r0[i] = new double[width]; |
| } |
| |
| double **r1 = new double *[height]; |
| for (int i = 0; i < height; i++) { |
| r1[i] = new double[width]; |
| } |
| |
| uint8_t **repic = new uint8_t *[height]; |
| for (int i = 0; i < height; i++) { |
| repic[i] = new uint8_t[width]; |
| } |
| |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| // sub_dgr[r][c] = dgd[r * in_stride + c]; |
| sub_dgr[r][c] = dgd[r * dgd_stride + c]; |
| r0[r][c] = output[r * 2 * out_stride + c * 2] * max_val; |
| r1[r][c] = output[r * 2 * out_stride + c * 2 + 1] * max_val; |
| } |
| } |
| int scale; |
| int qp = qindex / 4; |
| if (qp < 17) { |
| scale = 16384; |
| } else if (17 <= qp && qp < 27) { |
| scale = 16384; |
| } else if (27 <= qp && qp < 31) { |
| scale = 8192; |
| } else if (31 <= qp && qp < 37) { |
| scale = 8192; |
| } else if (37 <= qp && qp < 47) { |
| scale = 4192; |
| } else if (47 <= qp && qp < 57) { |
| scale = 2046; |
| } else { |
| scale = 2046; |
| } |
| int blockSize = frameType; |
| double cols = ceil(double(height) / blockSize); |
| double rows = ceil(double(width) / blockSize); |
| double number_crlc = cols * rows; |
| int *A = new int[(int)number_crlc * 2]; |
| int index_A = 0; |
| int start_row = 0; |
| int end_row = 0; |
| int start_clow = 0; |
| int end_clow = 0; |
| int num_block = 0; |
| int testnum = 10; |
| |
| for (int i = 0; i < cols; i++) { |
| for (int j = 0; j < rows; j++) { |
| if (i == cols - 1) { |
| start_clow = height - blockSize; |
| end_clow = height; |
| } else { |
| start_clow = i * blockSize; |
| end_clow = (i + 1) * blockSize; |
| } |
| |
| if (j == rows - 1) { |
| start_row = width - blockSize; |
| end_row = width; |
| } else { |
| start_row = j * blockSize; |
| end_row = (j + 1) * blockSize; |
| } |
| |
| if (width < blockSize) { |
| start_row = 0; |
| end_row = width; |
| } |
| |
| if (height < blockSize) { |
| start_clow = 0; |
| end_clow = height; |
| } |
| |
| int lenth_clows = end_clow - start_clow; |
| int lenth_rows = end_row - start_row; |
| |
| int lenth = lenth_clows * lenth_rows; |
| int *sub_r_flatten = new int[lenth]; |
| int k = 0; |
| |
| double A0 = ci->unit_info[num_block].xqd[0]; |
| double A1 = ci->unit_info[num_block].xqd[1]; |
| |
| num_block++; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| repic[i][j] = int(round(sub_dgr[i][j] + A0 * r0[i][j] / scale + |
| A1 * r1[i][j] / scale)); |
| repic[i][j] = clip_pixel(repic[i][j]); |
| } |
| } |
| } |
| } |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| rst[r * rst_stride + c] = clip_pixel(repic[r][c]); |
| } |
| } |
| return 1; |
| } |
| |
| extern "C" int av1_restore_cnn_guided_decode_img_tflite_highbd( |
| int qindex, 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, CRLCInfo *ci, int frameType) { |
| TfLiteDelegate *xnnpack_delegate = get_tflite_xnnpack_delegate(num_threads); |
| std::unique_ptr<tflite::Interpreter> interpreter = |
| get_tflite_interpreter(qindex, width, height, num_threads, is_intra_only, |
| is_luma, xnnpack_delegate); |
| |
| // Prepare input. |
| const auto max_val = static_cast<float>((1 << bit_depth) - 1); |
| const int in_stride = width; |
| auto input = interpreter->typed_input_tensor<float>(0); |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| input[r * in_stride + c] = |
| static_cast<float>(dgd[r * dgd_stride + c]) / max_val; |
| assert(input[r * in_stride + c] >= 0.0f); |
| assert(input[r * in_stride + c] <= 1.0f); |
| } |
| } |
| |
| // Invoke TFlite inference. |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| auto status = interpreter->Invoke(); |
| if (status != kTfLiteOk) { |
| reporter->Report("Failed at interpreter invocation"); |
| return 0; |
| } |
| |
| // Use the output to restore 'dgd' and store in 'rst'. |
| const auto output = interpreter->typed_output_tensor<float>(0); |
| const int out_stride = width; |
| |
| uint16_t **sub_dgr = new uint16_t *[height]; |
| for (int i = 0; i < height; i++) { |
| sub_dgr[i] = new uint16_t[width]; |
| } |
| |
| double **r0 = new double *[height]; |
| for (int i = 0; i < height; i++) { |
| r0[i] = new double[width]; |
| } |
| |
| double **r1 = new double *[height]; |
| for (int i = 0; i < height; i++) { |
| r1[i] = new double[width]; |
| } |
| |
| uint16_t **repic = new uint16_t *[height]; |
| for (int i = 0; i < height; i++) { |
| repic[i] = new uint16_t[width]; |
| } |
| |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| // sub_dgr[r][c] = dgd[r * in_stride + c]; |
| sub_dgr[r][c] = dgd[r * dgd_stride + c]; |
| r0[r][c] = output[r * 2 * out_stride + c * 2] * max_val; |
| r1[r][c] = output[r * 2 * out_stride + c * 2 + 1] * max_val; |
| } |
| } |
| int scale; |
| int qp = qindex / 4; |
| if (qp < 17) { |
| scale = 16384; |
| } else if (17 <= qp && qp < 27) { |
| scale = 16384; |
| } else if (27 <= qp && qp < 31) { |
| scale = 8192; |
| } else if (31 <= qp && qp < 37) { |
| scale = 8192; |
| } else if (37 <= qp && qp < 47) { |
| scale = 4192; |
| } else if (47 <= qp && qp < 57) { |
| scale = 2046; |
| } else { |
| scale = 2046; |
| } |
| int blockSize = frameType; |
| double cols = ceil(double(height) / blockSize); |
| double rows = ceil(double(width) / blockSize); |
| double number_crlc = cols * rows; |
| int *A = new int[(int)number_crlc * 2]; |
| int index_A = 0; |
| int start_row = 0; |
| int end_row = 0; |
| int start_clow = 0; |
| int end_clow = 0; |
| int num_block = 0; |
| int testnum = 10; |
| |
| for (int i = 0; i < cols; i++) { |
| for (int j = 0; j < rows; j++) { |
| if (i == cols - 1) { |
| start_clow = height - blockSize; |
| end_clow = height; |
| } else { |
| start_clow = i * blockSize; |
| end_clow = (i + 1) * blockSize; |
| } |
| |
| if (j == rows - 1) { |
| start_row = width - blockSize; |
| end_row = width; |
| } else { |
| start_row = j * blockSize; |
| end_row = (j + 1) * blockSize; |
| } |
| |
| if (width < blockSize) { |
| start_row = 0; |
| end_row = width; |
| } |
| |
| if (height < blockSize) { |
| start_clow = 0; |
| end_clow = height; |
| } |
| |
| int lenth_clows = end_clow - start_clow; |
| int lenth_rows = end_row - start_row; |
| |
| int lenth = lenth_clows * lenth_rows; |
| int *sub_r_flatten = new int[lenth]; |
| int k = 0; |
| |
| double A0 = ci->unit_info[num_block].xqd[0]; |
| double A1 = ci->unit_info[num_block].xqd[1]; |
| |
| num_block++; |
| for (int i = start_clow; i < end_clow; i++) { |
| for (int j = start_row; j < end_row; j++) { |
| repic[i][j] = int(round(sub_dgr[i][j] + A0 * r0[i][j] / scale + |
| A1 * r1[i][j] / scale)); |
| repic[i][j] = clip_pixel_highbd(repic[i][j], bit_depth); |
| } |
| } |
| } |
| } |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| rst[r * rst_stride + c] = clip_pixel_highbd(repic[r][c], bit_depth); |
| } |
| } |
| return 1; |
| } |
| |
| extern "C" void av1_restore_cnn_guided_tflite(AV1_COMMON *cm, int num_threads, |
| YV12_BUFFER_CONFIG *source_frame, |
| int plane_from, int plane_to) { |
| YV12_BUFFER_CONFIG *buf = &cm->cur_frame->buf; |
| const int is_intra_only = frame_is_intra_only(cm); |
| for (int plane = plane_from; plane <= plane_to; ++plane) { |
| const int is_luma = (plane == AOM_PLANE_Y); |
| if (cm->seq_params.use_highbitdepth) { |
| switch (plane) { |
| case AOM_PLANE_Y: |
| av1_restore_cnn_guided_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->y_buffer), |
| buf->y_crop_width, buf->y_crop_height, buf->y_stride, |
| CONVERT_TO_SHORTPTR(buf->y_buffer), buf->y_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma, |
| CONVERT_TO_SHORTPTR(source_frame->y_buffer), |
| source_frame->y_stride, &cm->crlc_info[0], |
| cm->crlc_info->crlc_unit_size); |
| break; |
| case AOM_PLANE_U: |
| av1_restore_cnn_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->u_buffer), |
| buf->uv_crop_width, buf->uv_crop_height, buf->uv_stride, |
| CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma); |
| break; |
| case AOM_PLANE_V: |
| av1_restore_cnn_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->v_buffer), |
| buf->uv_crop_width, buf->uv_crop_height, buf->uv_stride, |
| CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma); |
| break; |
| default: assert(0 && "Invalid plane index"); |
| } |
| } else { |
| assert(cm->seq_params.bit_depth == 8); |
| switch (plane) { |
| case AOM_PLANE_Y: |
| av1_restore_cnn_guided_img_tflite( |
| cm->base_qindex, buf->y_buffer, buf->y_crop_width, |
| buf->y_crop_height, buf->y_stride, buf->y_buffer, buf->y_stride, |
| num_threads, is_intra_only, is_luma, source_frame->y_buffer, |
| source_frame->y_stride, &cm->crlc_info[0], |
| cm->crlc_info->crlc_unit_size); |
| break; |
| case AOM_PLANE_U: |
| av1_restore_cnn_img_tflite( |
| cm->base_qindex, buf->u_buffer, buf->uv_crop_width, |
| buf->uv_crop_height, buf->uv_stride, buf->u_buffer, |
| buf->uv_stride, num_threads, is_intra_only, is_luma); |
| break; |
| case AOM_PLANE_V: |
| av1_restore_cnn_img_tflite( |
| cm->base_qindex, buf->v_buffer, buf->uv_crop_width, |
| buf->uv_crop_height, buf->uv_stride, buf->v_buffer, |
| buf->uv_stride, num_threads, is_intra_only, is_luma); |
| break; |
| default: assert(0 && "Invalid plane index"); |
| } |
| } |
| } |
| } |
| extern "C" void av1_restore_cnn_guided_decode_tflite(AV1_COMMON *cm, |
| int num_threads, |
| int plane_from, |
| int plane_to) { |
| YV12_BUFFER_CONFIG *buf = &cm->cur_frame->buf; |
| const int is_intra_only = frame_is_intra_only(cm); |
| for (int plane = plane_from; plane <= plane_to; ++plane) { |
| const int is_luma = (plane == AOM_PLANE_Y); |
| if (cm->seq_params.use_highbitdepth) { |
| switch (plane) { |
| case AOM_PLANE_Y: |
| av1_restore_cnn_guided_decode_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->y_buffer), |
| buf->y_crop_width, buf->y_crop_height, buf->y_stride, |
| CONVERT_TO_SHORTPTR(buf->y_buffer), buf->y_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma, |
| &cm->crlc_info[0], cm->crlc_info->crlc_unit_size); |
| break; |
| case AOM_PLANE_U: |
| av1_restore_cnn_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->u_buffer), |
| buf->uv_crop_width, buf->uv_crop_height, buf->uv_stride, |
| CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma); |
| break; |
| case AOM_PLANE_V: |
| av1_restore_cnn_img_tflite_highbd( |
| cm->base_qindex, CONVERT_TO_SHORTPTR(buf->v_buffer), |
| buf->uv_crop_width, buf->uv_crop_height, buf->uv_stride, |
| CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_stride, num_threads, |
| cm->seq_params.bit_depth, is_intra_only, is_luma); |
| break; |
| default: assert(0 && "Invalid plane index"); |
| } |
| } else { |
| assert(cm->seq_params.bit_depth == 8); |
| switch (plane) { |
| case AOM_PLANE_Y: |
| av1_restore_cnn_guided_decode_img_tflite( |
| cm->base_qindex, buf->y_buffer, buf->y_crop_width, |
| buf->y_crop_height, buf->y_stride, buf->y_buffer, buf->y_stride, |
| num_threads, is_intra_only, is_luma, &cm->crlc_info[0], |
| cm->crlc_info->crlc_unit_size); |
| break; |
| case AOM_PLANE_U: |
| av1_restore_cnn_img_tflite( |
| cm->base_qindex, buf->u_buffer, buf->uv_crop_width, |
| buf->uv_crop_height, buf->uv_stride, buf->u_buffer, |
| buf->uv_stride, num_threads, is_intra_only, is_luma); |
| break; |
| case AOM_PLANE_V: |
| av1_restore_cnn_img_tflite( |
| cm->base_qindex, buf->v_buffer, buf->uv_crop_width, |
| buf->uv_crop_height, buf->uv_stride, buf->v_buffer, |
| buf->uv_stride, num_threads, is_intra_only, is_luma); |
| break; |
| default: assert(0 && "Invalid plane index"); |
| } |
| } |
| } |
| } |
| |
| #endif // CONFIG_CNN_CRLC_GUIDED |
| |
| #if CONFIG_NN_RECON |
| // Builds and returns the TFlite interpreter. |
| static std::unique_ptr<tflite::Interpreter> get_nn_recon_tflite_interpreter( |
| int width, int height, int num_threads) { |
| auto model = tflite::GetModel(tx16x16_tflite); |
| tflite::MutableOpResolver resolver; |
| RegisterSelectedOpsAllQps(&resolver); |
| tflite::InterpreterBuilder builder(model, resolver); |
| |
| std::unique_ptr<tflite::Interpreter> interpreter; |
| builder(&interpreter); |
| interpreter->SetNumThreads(AOMMAX(num_threads, 1)); |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| |
| // Dimension order: batch_size, height, width, num_channels. |
| // Note: height comes before width here! |
| const std::vector<int> in_out_dims = { 1, height, width, 1 }; |
| // We only need to resize the input tensor. All other tensors (including |
| // output tensor) will be resized automatically. |
| if (interpreter->ResizeInputTensor(interpreter->inputs()[0], in_out_dims) != |
| kTfLiteOk) { |
| reporter->Report("Failed at input tensor resize"); |
| return nullptr; |
| } |
| |
| if (interpreter->AllocateTensors() != kTfLiteOk) { |
| reporter->Report("Failed at tensor allocation"); |
| return nullptr; |
| } |
| return interpreter; |
| } |
| |
| extern "C" int av1_cnn_recon_tflite(uint8_t *dst, int dst_stride, int height, |
| int width) { |
| const int num_threads = 1; |
| std::unique_ptr<tflite::Interpreter> interpreter = |
| get_nn_recon_tflite_interpreter(width, height, num_threads); |
| |
| // Prepare input. |
| const int in_stride = width; |
| auto input = interpreter->typed_input_tensor<float>(0); |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| input[r * in_stride + c] = static_cast<float>(dst[r * dst_stride + c]); |
| } |
| } |
| |
| // Invoke TFlite inference. |
| tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter(); |
| auto status = interpreter->Invoke(); |
| if (status != kTfLiteOk) { |
| reporter->Report("Failed at interpreter invocation"); |
| return 0; |
| } |
| |
| // Use the output to restore 'dgd' and store in 'rst'. |
| const auto output = interpreter->typed_output_tensor<float>(0); |
| const int out_stride = width; |
| for (int r = 0; r < height; ++r) { |
| for (int c = 0; c < width; ++c) { |
| const int residue = static_cast<int>(output[r * out_stride + c] + 0.5); |
| dst[r * dst_stride + c] = clip_pixel(dst[r * dst_stride + c] + residue); |
| } |
| } |
| return 1; |
| } |
| #endif // CONFIG_NN_RECON |
