Guided CNN Restoration: further improvements to unit sizes
Includes following changes:
* Similar to traditional loop restoration framework: If remaining rows/cols
near boundary are half-unit-height/width, extend the previous unit to cover those rows/cols.
* For narrow units near boundary, allow splitting either bigger dimension only (to HORZ/VERT type).
* Signal one of 4 unit sizes at frame level: largest allowed size is
min(2048, next_power_of_2 of largest dimension). And next 3 sizes are max_size/2, max_size/4 and max_size/8.
Also some code cleanups related to this.
diff --git a/av1/common/cnn_tflite.cc b/av1/common/cnn_tflite.cc
index b94792b..6ee12cc 100644
--- a/av1/common/cnn_tflite.cc
+++ b/av1/common/cnn_tflite.cc
@@ -618,30 +618,19 @@
return 1;
}
-typedef enum {
- GUIDED_QT_NONE,
- GUIDED_QT_SPLIT,
- GUIDED_QT_HORZ,
- GUIDED_QT_VERT,
- GUIDED_QT_TYPES,
- GUIDED_QT_INVALID = -1
-} GuidedQuadTreePartitionType;
-
// Get unit width and height based on max size and partition type.
-static void get_unit_size(int quadtree_max_size,
+static void get_unit_size(int max_unit_width, int max_unit_height,
GuidedQuadTreePartitionType partition_type,
int *unit_width, int *unit_height) {
assert(partition_type >= 0 && partition_type < GUIDED_QT_TYPES);
- const int full_size = quadtree_max_size;
- const int half_size = quadtree_max_size >> 1;
*unit_width =
(partition_type == GUIDED_QT_NONE || partition_type == GUIDED_QT_HORZ)
- ? full_size
- : half_size;
+ ? max_unit_width
+ : max_unit_width >> 1;
*unit_height =
(partition_type == GUIDED_QT_NONE || partition_type == GUIDED_QT_VERT)
- ? full_size
- : half_size;
+ ? max_unit_height
+ : max_unit_height >> 1;
}
// ------------------- Guided Quadtree: Encoder ------------------------------//
@@ -968,17 +957,24 @@
const std::vector<std::vector<std::vector<double>>> &interm,
GuidedQuadTreePartitionType partition_type, const uint16_t *src,
int src_stride, const uint16_t *dgd, int dgd_stride, int start_row,
- int end_row, int start_col, int end_col, int quadtree_max_size,
- const int *quadtset, int rdmult, const std::pair<int, int> &prev_A,
- const int *splitcosts, const int *norestorecosts, int bit_depth,
- bool is_partial_unit, double *this_rdcost,
- std::vector<std::vector<uint16_t>> &out,
+ int end_row, int start_col, int end_col, int max_unit_width,
+ int max_unit_height, const int *quadtset, int rdmult,
+ const std::pair<int, int> &prev_A, const int *quad_split_costs,
+ const int *binary_split_costs, const int *norestorecosts, int bit_depth,
+ bool is_horz_partitioning_allowed, int is_vert_partitioning_allowed,
+ double *this_rdcost, std::vector<std::vector<uint16_t>> &out,
std::vector<std::pair<int, int>> &A) {
- assert(IMPLIES(is_partial_unit, partition_type == GUIDED_QT_NONE));
+ assert(IMPLIES(
+ !is_horz_partitioning_allowed,
+ partition_type == GUIDED_QT_NONE || partition_type == GUIDED_QT_VERT));
+ assert(IMPLIES(
+ !is_vert_partitioning_allowed,
+ partition_type == GUIDED_QT_NONE || partition_type == GUIDED_QT_HORZ));
// Get unit width and height based on partition type.
int unit_width;
int unit_height;
- get_unit_size(quadtree_max_size, partition_type, &unit_width, &unit_height);
+ get_unit_size(max_unit_width, max_unit_height, partition_type, &unit_width,
+ &unit_height);
// Compute restored unit, a0 and a1.
generate_linear_combination(interm, src, src_stride, dgd, dgd_stride,
@@ -995,11 +991,16 @@
compute_sse(out, src, src_stride, start_row, end_row, start_col, end_col);
// Compute Rate.
- const int num_bits_for_a = compute_rate(A, prev_A, quadtset, norestorecosts);
- // Partition is implied to be NONE in case of partial unit.
+ const int a_signaling_cost =
+ compute_rate(A, prev_A, quadtset, norestorecosts);
+ // Partition signaling cost depending on 1, 2 or 4 possible partition types.
const int partition_signaling_cost =
- is_partial_unit ? 0 : splitcosts[partition_type];
- const int bitrate = num_bits_for_a + partition_signaling_cost;
+ is_horz_partitioning_allowed && is_vert_partitioning_allowed
+ ? quad_split_costs[partition_type]
+ : (is_horz_partitioning_allowed || is_vert_partitioning_allowed)
+ ? binary_split_costs[partition_type]
+ : 0;
+ const int bitrate = a_signaling_cost + partition_signaling_cost;
// Compute RDCost.
*this_rdcost =
@@ -1013,39 +1014,54 @@
static void select_quadtree_partitioning(
const std::vector<std::vector<std::vector<double>>> &interm,
const uint16_t *src, int src_stride, int start_row, int start_col,
- int width, int height, int quadtree_max_size, const int *quadtset,
- int rdmult, const std::pair<int, int> &prev_A, const int *splitcosts,
- const int norestorecosts[2], int bit_depth, const uint16_t *dgd,
- int dgd_stride, std::vector<int> &split,
+ int width, int height, int quadtree_max_size, int max_unit_width,
+ int max_unit_height, const int *quadtset, int rdmult,
+ const std::pair<int, int> &prev_A, const int *quad_split_costs,
+ const int *binary_split_costs, const int norestorecosts[2], int bit_depth,
+ const uint16_t *dgd, int dgd_stride, std::vector<int> &split,
std::vector<std::pair<int, int>> &A, double *rdcost) {
- const int end_row = AOMMIN(start_row + quadtree_max_size, height);
- const int end_col = AOMMIN(start_col + quadtree_max_size, width);
- const bool is_partial_unit = (start_row + quadtree_max_size > height) ||
- (start_col + quadtree_max_size > width);
+ const int end_row = AOMMIN(start_row + max_unit_height, height);
+ const int end_col = AOMMIN(start_col + max_unit_width, width);
+ // Check for special cases near boundary.
+ const bool is_horz_partitioning_allowed =
+ (max_unit_height >= quadtree_max_size);
+ const bool is_vert_partitioning_allowed =
+ (max_unit_width >= quadtree_max_size);
+ const bool is_split_partitioning_allowed =
+ is_horz_partitioning_allowed && is_vert_partitioning_allowed;
auto best_rdcost = DBL_MAX;
std::vector<std::pair<int, int>> best_A;
std::vector<std::vector<uint16_t>> best_out(
- quadtree_max_size, std::vector<uint16_t>(quadtree_max_size));
+ max_unit_height, std::vector<uint16_t>(max_unit_width));
GuidedQuadTreePartitionType best_partition_type = GUIDED_QT_INVALID;
for (int type = 0; type < GUIDED_QT_TYPES; ++type) {
const auto this_partition_type = (GuidedQuadTreePartitionType)type;
- // Special case: if only partial unit is within boundary, we implicitly
- // use NONE partitioning and do not try the splitting options.
- if (is_partial_unit && (this_partition_type != GUIDED_QT_NONE)) {
+ // Check for special cases near boundary.
+ if (!is_horz_partitioning_allowed &&
+ (this_partition_type == GUIDED_QT_HORZ)) {
continue;
}
-
+ if (!is_vert_partitioning_allowed &&
+ (this_partition_type == GUIDED_QT_VERT)) {
+ continue;
+ }
+ if (!is_split_partitioning_allowed &&
+ (this_partition_type == GUIDED_QT_SPLIT)) {
+ continue;
+ }
+ // Try this partition type.
double this_rdcost;
std::vector<std::pair<int, int>> this_A;
std::vector<std::vector<uint16_t>> this_out(
- quadtree_max_size, std::vector<uint16_t>(quadtree_max_size));
- try_one_partition(interm, this_partition_type, src, src_stride, dgd,
- dgd_stride, start_row, end_row, start_col, end_col,
- quadtree_max_size, quadtset, rdmult, prev_A, splitcosts,
- norestorecosts, bit_depth, is_partial_unit, &this_rdcost,
- this_out, this_A);
+ max_unit_height, std::vector<uint16_t>(max_unit_width));
+ try_one_partition(
+ interm, this_partition_type, src, src_stride, dgd, dgd_stride,
+ start_row, end_row, start_col, end_col, max_unit_width, max_unit_height,
+ quadtset, rdmult, prev_A, quad_split_costs, binary_split_costs,
+ norestorecosts, bit_depth, is_horz_partitioning_allowed,
+ is_vert_partitioning_allowed, &this_rdcost, this_out, this_A);
if (this_rdcost < best_rdcost) {
best_rdcost = this_rdcost;
best_A = this_A;
@@ -1063,45 +1079,29 @@
}
// Save split decision.
- if (is_partial_unit) {
+ if (!is_horz_partitioning_allowed && !is_vert_partitioning_allowed) {
// Nothing should be added to 'split' array.
assert(best_partition_type == GUIDED_QT_NONE);
return;
}
- switch (best_partition_type) {
- case GUIDED_QT_NONE:
- split.push_back(0);
- split.push_back(0);
- break;
- case GUIDED_QT_SPLIT:
- split.push_back(0);
- split.push_back(1);
- break;
- case GUIDED_QT_HORZ:
- split.push_back(1);
- split.push_back(1);
- break;
- case GUIDED_QT_VERT:
- split.push_back(1);
- split.push_back(0);
- break;
- default: assert(0 && "Wrong partition type"); break;
- }
+ assert(best_partition_type >= 0 && best_partition_type < GUIDED_QT_TYPES);
+ split.push_back(best_partition_type);
}
static void apply_quadtree_partitioning(
const std::vector<std::vector<std::vector<double>>> &interm, int start_row,
int start_col, int width, int height, int quadtree_max_size,
- const int *quadtset, int bit_depth, const std::vector<int> &split,
- size_t &split_index, const std::vector<std::pair<int, int>> &A,
- size_t &A_index, uint16_t *dgd, int dgd_stride);
+ int max_unit_width, int max_unit_height, const int *quadtset, int bit_depth,
+ const std::vector<int> &split, size_t &split_index,
+ const std::vector<std::pair<int, int>> &A, size_t &A_index, uint16_t *dgd,
+ int dgd_stride);
// Top-level function to apply guided restoration on encoder side.
static int restore_cnn_quadtree_encode_img_tflite_highbd(
YV12_BUFFER_CONFIG *source_frame, AV1_COMMON *cm, int superres_denom,
- int rdmult, const int *splitcosts, int (*norestorecosts)[2],
- int num_threads, int bit_depth, int is_intra_only, int is_luma,
- int cnn_index, QUADInfo *quad_info, double *rdcost) {
+ int rdmult, const int *quad_split_costs, const int *binary_split_costs,
+ int (*norestorecosts)[2], int num_threads, int bit_depth, int is_intra_only,
+ int is_luma, int cnn_index, QUADInfo *quad_info, double *rdcost) {
YV12_BUFFER_CONFIG *dgd_buf = &cm->cur_frame->buf;
uint16_t *dgd = CONVERT_TO_SHORTPTR(dgd_buf->y_buffer);
const int dgd_stride = dgd_buf->y_stride;
@@ -1136,7 +1136,8 @@
std::vector<int> best_split; // selected partitioning options.
std::vector<std::pair<int, int>> best_A; // selected a0, a1 weight pairs.
double best_rdcost_total = DBL_MAX;
- for (int this_unit_index = 0; this_unit_index <= 1; ++this_unit_index) {
+ for (int this_unit_index = 0; this_unit_index < GUIDED_QT_UNIT_SIZES;
+ ++this_unit_index) {
const int quadtree_max_size =
quad_tree_get_unit_size(width, height, this_unit_index);
// For each quadtree unit, compute the best partitioning out of
@@ -1147,20 +1148,27 @@
// Previous a0, a1 pair is mid-point of the range by default.
std::pair<int, int> prev_A =
std::make_pair(GUIDED_A_MID + A0_min, GUIDED_A_MID + A1_min);
- // TODO(urvang): Include padded area in a unit if it's < unit size / 2?
- // If so, need to modify / replace quad_tree_get_unit_info_length().
- // Also double check: quad_tree_get_split_info_length().
- for (int row = 0; row < height; row += quadtree_max_size) {
- for (int col = 0; col < width; col += quadtree_max_size) {
+ const int ext_size = quadtree_max_size * 3 / 2;
+ for (int row = 0; row < height;) {
+ const int remaining_height = height - row;
+ const int this_unit_height =
+ (remaining_height < ext_size) ? remaining_height : quadtree_max_size;
+ for (int col = 0; col < width;) {
+ const int remaining_width = width - col;
+ const int this_unit_width =
+ (remaining_width < ext_size) ? remaining_width : quadtree_max_size;
double this_rdcost;
select_quadtree_partitioning(
interm, src, src_stride, row, col, width, height, quadtree_max_size,
- quadtset, rdmult, prev_A, splitcosts, this_norestorecosts,
+ this_unit_width, this_unit_height, quadtset, rdmult, prev_A,
+ quad_split_costs, binary_split_costs, this_norestorecosts,
bit_depth, dgd, dgd_stride, this_split, this_A, &this_rdcost);
// updates.
this_rdcost_total += this_rdcost;
prev_A = this_A.back();
+ col += this_unit_width;
}
+ row += this_unit_height;
}
// Update best options.
if (this_rdcost_total < best_rdcost_total) {
@@ -1188,22 +1196,33 @@
// Apply guided restoration to 'dgd' using best options above.
size_t split_index = 0;
size_t A_index = 0;
- for (int row = 0; row < height; row += quad_info->unit_size) {
- for (int col = 0; col < width; col += quad_info->unit_size) {
+ const int quadtree_max_size = quad_info->unit_size;
+ const int ext_size = quadtree_max_size * 3 / 2;
+ for (int row = 0; row < height;) {
+ const int remaining_height = height - row;
+ const int this_unit_height =
+ (remaining_height < ext_size) ? remaining_height : quadtree_max_size;
+ for (int col = 0; col < width;) {
+ const int remaining_width = width - col;
+ const int this_unit_width =
+ (remaining_width < ext_size) ? remaining_width : quadtree_max_size;
apply_quadtree_partitioning(
- interm, row, col, width, height, quad_info->unit_size, quadtset,
- bit_depth, best_split, split_index, best_A, A_index, dgd, dgd_stride);
+ interm, row, col, width, height, quadtree_max_size, this_unit_width,
+ this_unit_height, quadtset, bit_depth, best_split, split_index,
+ best_A, A_index, dgd, dgd_stride);
+ col += this_unit_width;
}
+ row += this_unit_height;
}
return 1;
}
extern "C" int av1_restore_cnn_quadtree_encode_tflite(
- AV1_COMMON *cm, YV12_BUFFER_CONFIG *source_frame, int RDMULT,
- int *splitcosts, 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) {
+ 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) {
YV12_BUFFER_CONFIG *buf = &cm->cur_frame->buf;
const int is_intra_only = frame_is_intra_only(cm);
for (int plane = 0; plane < av1_num_planes(cm); ++plane) {
@@ -1217,8 +1236,9 @@
case AOM_PLANE_Y:
ret = restore_cnn_quadtree_encode_img_tflite_highbd(
source_frame, cm, cm->superres_scale_denominator, RDMULT,
- splitcosts, norestorecosts, num_threads, cm->seq_params.bit_depth,
- is_intra_only, is_luma, cnn_index, quad_info, rdcost);
+ quad_split_costs, binary_split_costs, norestorecosts, num_threads,
+ cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index,
+ quad_info, rdcost);
if (ret == 0) return ret;
break;
case AOM_PLANE_U:
@@ -1290,42 +1310,30 @@
static void apply_quadtree_partitioning(
const std::vector<std::vector<std::vector<double>>> &interm, int start_row,
int start_col, int width, int height, int quadtree_max_size,
- const int *quadtset, int bit_depth, const std::vector<int> &split,
- size_t &split_index, const std::vector<std::pair<int, int>> &A,
- size_t &A_index, uint16_t *dgd, int dgd_stride) {
- const int end_row = AOMMIN(start_row + quadtree_max_size, height);
- const int end_col = AOMMIN(start_col + quadtree_max_size, width);
- const bool is_partial_unit = (start_row + quadtree_max_size > height) ||
- (start_col + quadtree_max_size > width);
+ int max_unit_width, int max_unit_height, const int *quadtset, int bit_depth,
+ const std::vector<int> &split, size_t &split_index,
+ const std::vector<std::pair<int, int>> &A, size_t &A_index, uint16_t *dgd,
+ int dgd_stride) {
+ const int end_row = AOMMIN(start_row + max_unit_height, height);
+ const int end_col = AOMMIN(start_col + max_unit_width, width);
+ // Check for special cases near boundary.
+ const bool is_horz_partitioning_allowed =
+ (max_unit_height >= quadtree_max_size);
+ const bool is_vert_partitioning_allowed =
+ (max_unit_width >= quadtree_max_size);
// Get partition type.
GuidedQuadTreePartitionType partition_type = GUIDED_QT_NONE;
- if (!is_partial_unit) {
- const int spl1 = split[split_index++];
- const int spl2 = split[split_index++];
- if (spl1 == 0) {
- if (spl2 == 0) {
- partition_type = GUIDED_QT_NONE; // (0, 0)
- } else {
- assert(spl2 == 1);
- partition_type = GUIDED_QT_SPLIT; // (0, 1)
- }
- } else {
- assert(spl1 == 1);
- if (spl2 == 1) {
- partition_type = GUIDED_QT_HORZ; // (1, 1)
- } else {
- assert(spl2 == 0);
- partition_type = GUIDED_QT_VERT; // (1, 0)
- }
- }
+ if (is_horz_partitioning_allowed || is_vert_partitioning_allowed) {
+ partition_type = (GuidedQuadTreePartitionType)split[split_index++];
}
assert(partition_type >= 0 && partition_type < GUIDED_QT_TYPES);
// Get unit width and height based on partition type.
int unit_width;
int unit_height;
- get_unit_size(quadtree_max_size, partition_type, &unit_width, &unit_height);
+ get_unit_size(max_unit_width, max_unit_height, partition_type, &unit_width,
+ &unit_height);
// Compute restored unit, a0 and a1 with given A parameters.
apply_linear_combination(interm, start_row, end_row, start_col, end_col,
@@ -1377,12 +1385,22 @@
// For each quadtree unit, apply given quadtree partitioning.
size_t split_index = 0;
size_t A_index = 0;
- for (int row = 0; row < height; row += quadtree_max_size) {
- for (int col = 0; col < width; col += quadtree_max_size) {
+ const int ext_size = quadtree_max_size * 3 / 2;
+ for (int row = 0; row < height;) {
+ const int remaining_height = height - row;
+ const int this_unit_height =
+ (remaining_height < ext_size) ? remaining_height : quadtree_max_size;
+ for (int col = 0; col < width;) {
+ const int remaining_width = width - col;
+ const int this_unit_width =
+ (remaining_width < ext_size) ? remaining_width : quadtree_max_size;
apply_quadtree_partitioning(interm, row, col, width, height,
- quadtree_max_size, quadtset, bit_depth, split,
+ quadtree_max_size, this_unit_width,
+ this_unit_height, quadtset, bit_depth, split,
split_index, A, A_index, dgd, dgd_stride);
+ col += this_unit_width;
}
+ row += this_unit_height;
}
assert(split_index == split.size());
assert(A_index == A.size());
diff --git a/av1/common/cnn_tflite.h b/av1/common/cnn_tflite.h
index f74165c..3ed3682 100644
--- a/av1/common/cnn_tflite.h
+++ b/av1/common/cnn_tflite.h
@@ -114,9 +114,9 @@
// 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 *splitcosts, 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);
+ 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,
diff --git a/av1/common/entropy.c b/av1/common/entropy.c
index 89a47ff..1e63d77 100644
--- a/av1/common/entropy.c
+++ b/av1/common/entropy.c
@@ -241,6 +241,7 @@
#endif // CONFIG_RST_MERGECOEFFS
#if CONFIG_CNN_GUIDED_QUADTREE
RESET_CDF_COUNTER(fc->cnn_guided_quad_cdf, 4);
+ RESET_CDF_COUNTER(fc->cnn_guided_binary_cdf, 2);
RESET_CDF_COUNTER(fc->cnn_guided_norestore_cdf, 2);
#endif // CONFIG_CNN_GUIDED_QUADTREE
#if CONFIG_AIMC
diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index 0d15acb..def65a8 100644
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -1560,6 +1560,9 @@
static const aom_cdf_prob default_cnn_guided_quad_cdf[CDF_SIZE(4)] = {
AOM_CDF4(23552, 24576, 28672),
};
+static const aom_cdf_prob default_cnn_guided_binary_cdf[CDF_SIZE(2)] = {
+ AOM_CDF2(23552),
+};
static const aom_cdf_prob
default_cnn_guided_norestore_cdf[GUIDED_NORESTORE_CONTEXTS][CDF_SIZE(2)] = {
{ AOM_CDF2(16384) }, { AOM_CDF2(24576) }
@@ -1967,6 +1970,7 @@
#endif // CONFIG_PC_WIENER
#if CONFIG_CNN_GUIDED_QUADTREE
av1_copy(fc->cnn_guided_quad_cdf, default_cnn_guided_quad_cdf);
+ av1_copy(fc->cnn_guided_binary_cdf, default_cnn_guided_binary_cdf);
av1_copy(fc->cnn_guided_norestore_cdf, default_cnn_guided_norestore_cdf);
#endif // CONFIG_CNN_GUIDED_QUADTREE
#if CONFIG_AIMC
diff --git a/av1/common/entropymode.h b/av1/common/entropymode.h
index 3b4e443..e277f52 100644
--- a/av1/common/entropymode.h
+++ b/av1/common/entropymode.h
@@ -240,7 +240,12 @@
aom_cdf_prob merged_param_cdf[CDF_SIZE(2)];
#endif // CONFIG_RST_MERGECOEFFS
#if CONFIG_CNN_GUIDED_QUADTREE
+ // CDF to signal one of GUIDED_QT_TYPES partition types for each unit, when
+ // all types are allowed.
aom_cdf_prob cnn_guided_quad_cdf[CDF_SIZE(4)];
+ // CDF to signal one of 2 partition types for each unit: GUIDED_QT_NONE or
+ // GUIDED_QT_HORZ/VERT, when only 2 types are allowed.
+ aom_cdf_prob cnn_guided_binary_cdf[CDF_SIZE(2)];
aom_cdf_prob cnn_guided_norestore_cdf[GUIDED_NORESTORE_CONTEXTS][CDF_SIZE(2)];
#endif // CONFIG_CNN_GUIDED_QUADTREE
#if !CONFIG_AIMC
diff --git a/av1/common/guided_quadtree.c b/av1/common/guided_quadtree.c
index e296c0b..d9c4400 100644
--- a/av1/common/guided_quadtree.c
+++ b/av1/common/guided_quadtree.c
@@ -275,52 +275,63 @@
dst->signaled = src->signaled;
}
-// Returns (int)floor(x / y),
-#define DIVIDE_WITH_FLOOR(x, y) ((x) / (y))
-// Returns (int)ceil(x / y),
-#define DIVIDE_WITH_CEILING(x, y) (((x) + (y)-1) / (y))
-
-int quad_tree_get_unit_info_length(int width, int height, int unit_length,
- const QUADSplitInfo *split_info,
- int split_info_length) {
- // We can compute total units as follows:
- // (1) regular units: they may / may not be split. So, compute length of
- // regular unit info by going through the split_info array. (2) unregular
- // units (blocks near boundary that are NOT unit_length in size): they are
- // never split. So, length of unregular unit info is same as number of
- // unregular units.
- const int regular_units = DIVIDE_WITH_FLOOR(width, unit_length) *
- DIVIDE_WITH_FLOOR(height, unit_length);
- assert(regular_units * 2 == split_info_length);
- const int total_units = DIVIDE_WITH_CEILING(width, unit_length) *
- DIVIDE_WITH_CEILING(height, unit_length);
- const int unregular_unit_info_len = total_units - regular_units;
-
- int regular_unit_info_len = 0;
- for (int i = 0; i < split_info_length; i += 2) {
- if (split_info == NULL ||
- (split_info[i].split == 0 && split_info[i + 1].split == 1)) {
- regular_unit_info_len += 4; // Split
- } else if (split_info[i].split == 1 && split_info[i + 1].split == 1) {
- regular_unit_info_len += 2; // Horz
- } else if (split_info[i].split == 1 && split_info[i + 1].split == 0) {
- regular_unit_info_len += 2; // Vert
- } else {
- assert(split_info[i].split == 0 && split_info[i + 1].split == 0);
- regular_unit_info_len += 1; // No split
+int quad_tree_get_max_unit_info_length(int width, int height, int unit_length) {
+ int unit_info_length = 0;
+ const int ext_size = unit_length * 3 / 2;
+ for (int row = 0; row < height;) {
+ const int remaining_height = height - row;
+ const int this_unit_height =
+ (remaining_height < ext_size) ? remaining_height : unit_length;
+ for (int col = 0; col < width;) {
+ const int remaining_width = width - col;
+ const int this_unit_width =
+ (remaining_width < ext_size) ? remaining_width : unit_length;
+ // Check for special cases near boundary.
+ const bool is_horz_partitioning_allowed =
+ (this_unit_height >= unit_length);
+ const bool is_vert_partitioning_allowed =
+ (this_unit_width >= unit_length);
+ if (!is_horz_partitioning_allowed && !is_vert_partitioning_allowed) {
+ // Implicitly no split, so single unit info will be signaled.
+ ++unit_info_length;
+ } else {
+ // Assume maximum possible sub-units.
+ const int max_sub_units =
+ is_horz_partitioning_allowed && is_vert_partitioning_allowed ? 4
+ : 2;
+ unit_info_length += max_sub_units;
+ }
+ col += this_unit_width;
}
+ row += this_unit_height;
}
-
- return regular_unit_info_len + unregular_unit_info_len;
+ return unit_info_length;
}
int quad_tree_get_split_info_length(int width, int height, int unit_length) {
- // Split info only signaled for units of full size. Blocks near boundaries are
- // never split, so no info is signaled for those.
- const int num_split_info_wide = DIVIDE_WITH_FLOOR(width, unit_length);
- const int num_split_info_high = DIVIDE_WITH_FLOOR(height, unit_length);
- // 2 bits signaled for each split info.
- return num_split_info_wide * num_split_info_high * 2;
+ int split_info_len = 0;
+ const int ext_size = unit_length * 3 / 2;
+ for (int row = 0; row < height;) {
+ const int remaining_height = height - row;
+ const int this_unit_height =
+ (remaining_height < ext_size) ? remaining_height : unit_length;
+ for (int col = 0; col < width;) {
+ const int remaining_width = width - col;
+ const int this_unit_width =
+ (remaining_width < ext_size) ? remaining_width : unit_length;
+ // Check for special cases near boundary.
+ const bool is_horz_partitioning_allowed =
+ (this_unit_height >= unit_length);
+ const bool is_vert_partitioning_allowed =
+ (this_unit_width >= unit_length);
+ if (is_horz_partitioning_allowed || is_vert_partitioning_allowed) {
+ ++split_info_len;
+ }
+ col += this_unit_width;
+ }
+ row += this_unit_height;
+ }
+ return split_info_len;
}
void av1_alloc_quadtree_struct(struct AV1Common *cm, QUADInfo *quad_info) {
diff --git a/av1/common/guided_quadtree.h b/av1/common/guided_quadtree.h
index 4fef1e6..f51d9fe 100644
--- a/av1/common/guided_quadtree.h
+++ b/av1/common/guided_quadtree.h
@@ -30,17 +30,27 @@
#define GUIDED_A_MID (GUIDED_A_NUM_VALUES >> 1)
#define GUIDED_A_RANGE (GUIDED_A_NUM_VALUES - 1)
#define GUIDED_A_PAIR_BITS (GUIDED_A_BITS * 2 - 1)
+#define GUIDED_QT_UNIT_SIZES_LOG2 2
+#define GUIDED_QT_UNIT_SIZES (1 << GUIDED_QT_UNIT_SIZES_LOG2)
+
+typedef enum {
+ GUIDED_QT_NONE,
+ GUIDED_QT_SPLIT,
+ GUIDED_QT_HORZ,
+ GUIDED_QT_VERT,
+ GUIDED_QT_TYPES,
+ GUIDED_QT_INVALID = -1
+} GuidedQuadTreePartitionType;
int *get_quadparm_from_qindex(int qindex, int superres_denom, int is_intra_only,
int is_luma, int cnn_index);
void quad_copy(const QUADInfo *src, QUADInfo *dst, struct AV1Common *cm);
-// Get the length of unit info array based on dimensions and split info.
-// If split_info == NULL, assumes each block uses split, thereby returning
-// longest possible unit info length.
-int quad_tree_get_unit_info_length(int width, int height, int unit_length,
- const QUADSplitInfo *split_info,
- int split_info_length);
+
+// Get the maximum possible length of unit info array based on dimensions,
+// assuming each block uses split.
+int quad_tree_get_max_unit_info_length(int width, int height, int unit_length);
+
// Get the length of split info array based on dimensions.
int quad_tree_get_split_info_length(int width, int height, int unit_length);
@@ -56,10 +66,12 @@
// Get quad tree unit size.
static INLINE int quad_tree_get_unit_size(int width, int height,
int unit_index) {
- const bool is_720p_or_smaller = (width * height <= 1280 * 720);
- const int min_unit_size = is_720p_or_smaller ? 256 : 512;
- assert(unit_index >= 0 && unit_index <= 1);
- return min_unit_size << unit_index;
+ const int max_dim = AOMMAX(width, height);
+ const int max_dim_pow_2_bits = 1 + get_msb(max_dim);
+ const int max_dim_pow_2 = 1 << max_dim_pow_2_bits;
+ const int max_unit_size = AOMMAX(AOMMIN(max_dim_pow_2, 2048), 256);
+ assert(unit_index >= 0 && unit_index < GUIDED_QT_UNIT_SIZES);
+ return max_unit_size >> unit_index;
}
// Allocates buffers in 'quad_info' assuming 'quad_info->unit_index',
diff --git a/av1/decoder/decodeframe.c b/av1/decoder/decodeframe.c
index a71b6a6..a736131 100644
--- a/av1/decoder/decodeframe.c
+++ b/av1/decoder/decodeframe.c
@@ -137,7 +137,8 @@
#if CONFIG_CNN_GUIDED_QUADTREE
static AOM_INLINE void read_filter_quadtree(FRAME_CONTEXT *ctx, int QP,
- int cnn_index, int superres_denom,
+ int cnn_index, int width,
+ int height, int superres_denom,
int is_intra_only, QUADInfo *qi,
aom_reader *rb);
#endif // CONFIG_CNN_GUIDED_QUADTREE
@@ -1822,19 +1823,81 @@
}
}
#if CONFIG_CNN_GUIDED_QUADTREE
+static void read_quadtree_split_info(FRAME_CONTEXT *ctx, int width, int height,
+ QUADInfo *qi, aom_reader *rb) {
+ int unit_info_length = 0;
+ int split_info_index = 0;
+ const int ext_size = qi->unit_size * 3 / 2;
+ for (int row = 0; row < height;) {
+ const int remaining_height = height - row;
+ const int this_unit_height =
+ (remaining_height < ext_size) ? remaining_height : qi->unit_size;
+ for (int col = 0; col < width;) {
+ const int remaining_width = width - col;
+ const int this_unit_width =
+ (remaining_width < ext_size) ? remaining_width : qi->unit_size;
+ // Check for special cases near boundary.
+ const bool is_horz_partitioning_allowed =
+ (this_unit_height >= qi->unit_size);
+ const bool is_vert_partitioning_allowed =
+ (this_unit_width >= qi->unit_size);
+ if (!is_horz_partitioning_allowed && !is_vert_partitioning_allowed) {
+ // Implicitly no split, and single unit info will be signaled.
+ ++unit_info_length;
+ } else {
+ assert(split_info_index < qi->split_info_length);
+ // Read split info, depending on how many partition types are allowed.
+ const bool all_partition_types_allowed =
+ (is_horz_partitioning_allowed && is_vert_partitioning_allowed);
+ if (all_partition_types_allowed) {
+ qi->split_info[split_info_index].split =
+ aom_read_symbol(rb, ctx->cnn_guided_quad_cdf, 4, ACCT_STR);
+ } else {
+ const int do_split =
+ aom_read_symbol(rb, ctx->cnn_guided_binary_cdf, 2, ACCT_STR);
+ if (do_split) {
+ qi->split_info[split_info_index].split =
+ is_horz_partitioning_allowed ? GUIDED_QT_HORZ : GUIDED_QT_VERT;
+ } else {
+ qi->split_info[split_info_index].split = GUIDED_QT_NONE;
+ }
+ }
+ // Look at the split info to determine number of (sub)units.
+ const GuidedQuadTreePartitionType partition_type =
+ qi->split_info[split_info_index].split;
+ switch (partition_type) {
+ case GUIDED_QT_NONE: unit_info_length += 1; break;
+ case GUIDED_QT_HORZ:
+ case GUIDED_QT_VERT: unit_info_length += 2; break;
+ case GUIDED_QT_SPLIT: unit_info_length += 4; break;
+ default: assert(0 && "Wrong guided quadtree split type."); break;
+ }
+ ++split_info_index;
+ }
+ col += this_unit_width;
+ }
+ row += this_unit_height;
+ }
+ assert(qi->split_info_length == split_info_index);
+ qi->unit_info_length = unit_info_length;
+}
+
static AOM_INLINE void read_filter_quadtree(FRAME_CONTEXT *ctx, int QP,
- int cnn_index, int superres_denom,
+ int cnn_index, int width,
+ int height, int superres_denom,
int is_intra_only, QUADInfo *qi,
aom_reader *rb) {
- int A0_min, A1_min;
- int *quadtset;
- quadtset =
+ // Read partitioning info.
+ read_quadtree_split_info(ctx, width, height, qi, rb);
+
+ // Read weight parameters 'a'.
+ const int *quadtset =
get_quadparm_from_qindex(QP, superres_denom, is_intra_only, 1, cnn_index);
const int norestore_ctx =
get_guided_norestore_ctx(QP, superres_denom, is_intra_only);
- A0_min = quadtset[2];
- A1_min = quadtset[3];
+ const int A0_min = quadtset[2];
+ const int A1_min = quadtset[3];
int ref_0 = GUIDED_A_MID;
int ref_1 = GUIDED_A_MID;
@@ -1861,8 +1924,6 @@
ref_0 = qi->unit_info[i].xqd[0] - A0_min;
ref_1 = qi->unit_info[i].xqd[1] - A1_min;
}
- // printf("a0:%d a1:%d\n", qi->unit_info[i].xqd[0],
- // qi->unit_info[i].xqd[1]);
}
}
#endif // CONFIG_CNN_GUIDED_QUADTREE
@@ -1922,18 +1983,10 @@
#if CONFIG_CNN_GUIDED_QUADTREE
if (cm->use_cnn[0] && !cm->cnn_quad_info.signaled) {
QUADInfo *qi = (QUADInfo *)&cm->cnn_quad_info;
- for (int s = 0; s < qi->split_info_length; s += 2) {
- const int split_index = aom_read_symbol(
- reader, xd->tile_ctx->cnn_guided_quad_cdf, 4, ACCT_STR);
- qi->split_info[s].split = split_index >> 1;
- qi->split_info[s + 1].split = split_index & 1;
- }
- qi->unit_info_length = quad_tree_get_unit_info_length(
+ read_filter_quadtree(
+ xd->tile_ctx, cm->quant_params.base_qindex, cm->cnn_indices[0],
cm->superres_upscaled_width, cm->superres_upscaled_height,
- qi->unit_size, qi->split_info, qi->split_info_length);
- read_filter_quadtree(xd->tile_ctx, cm->quant_params.base_qindex,
- cm->cnn_indices[0], cm->superres_scale_denominator,
- frame_is_intra_only(cm), qi, reader);
+ cm->superres_scale_denominator, frame_is_intra_only(cm), qi, reader);
cm->cnn_quad_info.signaled = 1;
}
#endif // CONFIG_CNN_GUIDED_QUADTREE
@@ -2154,7 +2207,7 @@
#if CONFIG_CNN_GUIDED_QUADTREE
// Read quad tree unit index.
static INLINE int quad_tree_read_unit_index(struct aom_read_bit_buffer *rb) {
- const int unit_index = aom_rb_read_bit(rb);
+ const int unit_index = aom_rb_read_literal(rb, GUIDED_QT_UNIT_SIZES_LOG2);
return unit_index;
}
#endif // CONFIG_CNN_GUIDED_QUADTREE
@@ -2191,9 +2244,9 @@
qi->unit_size);
// We allocate unit info assuming maximum number of possible units for now.
// Actual length will be set later after actually reading split info.
- qi->unit_info_length = quad_tree_get_unit_info_length(
+ qi->unit_info_length = quad_tree_get_max_unit_info_length(
cm->superres_upscaled_width, cm->superres_upscaled_height,
- qi->unit_size, NULL, qi->split_info_length);
+ qi->unit_size);
av1_alloc_quadtree_struct(cm, qi);
}
#endif // CONFIG_CNN_GUIDED_QUADTREE
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 0c0c827..88d4b7e 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -85,8 +85,9 @@
#if CONFIG_CNN_GUIDED_QUADTREE
// Write crlc coeffs for one frame
static void write_filter_quadtree(FRAME_CONTEXT *ctx, int qp, int cnn_index,
- int superres_denom, int is_intra_only,
- const QUADInfo *ci, aom_writer *wb);
+ int width, int height, int superres_denom,
+ int is_intra_only, const QUADInfo *qi,
+ aom_writer *wb);
#endif // CONFIG_CNN_GUIDED_QUADTREE
#if CONFIG_IBC_SR_EXT
@@ -2334,22 +2335,10 @@
#if CONFIG_CNN_GUIDED_QUADTREE
if (cm->use_cnn[0] && !cm->cnn_quad_info.signaled) {
QUADInfo *qi = (QUADInfo *)&cm->cnn_quad_info;
- assert(qi->split_info_length ==
- quad_tree_get_split_info_length(cm->superres_upscaled_width,
- cm->superres_upscaled_height,
- qi->unit_size));
- for (int s = 0; s < qi->split_info_length; s += 2) {
- const int split_index =
- qi->split_info[s].split * 2 + qi->split_info[s + 1].split;
- aom_write_symbol(w, split_index, xd->tile_ctx->cnn_guided_quad_cdf, 4);
- }
- assert(qi->unit_info_length ==
- quad_tree_get_unit_info_length(
- cm->superres_upscaled_width, cm->superres_upscaled_height,
- qi->unit_size, qi->split_info, qi->split_info_length));
- write_filter_quadtree(xd->tile_ctx, cm->quant_params.base_qindex,
- cm->cnn_indices[0], cm->superres_scale_denominator,
- frame_is_intra_only(cm), qi, w);
+ write_filter_quadtree(
+ xd->tile_ctx, cm->quant_params.base_qindex, cm->cnn_indices[0],
+ cm->superres_upscaled_width, cm->superres_upscaled_height,
+ cm->superres_scale_denominator, frame_is_intra_only(cm), qi, w);
qi->signaled = 1;
}
#endif // CONFIG_CNN_GUIDED_QUADTREE
@@ -3062,9 +3051,75 @@
#endif // !CONFIG_NEW_DF
#if CONFIG_CNN_GUIDED_QUADTREE
+static void write_quadtree_split_info(FRAME_CONTEXT *ctx, int width, int height,
+ const QUADInfo *qi, aom_writer *wb) {
+#ifndef NDEBUG
+ int unit_info_length = 0;
+#endif // NDEBUG
+ int split_info_index = 0;
+ const int ext_size = qi->unit_size * 3 / 2;
+ for (int row = 0; row < height;) {
+ const int remaining_height = height - row;
+ const int this_unit_height =
+ (remaining_height < ext_size) ? remaining_height : qi->unit_size;
+ for (int col = 0; col < width;) {
+ const int remaining_width = width - col;
+ const int this_unit_width =
+ (remaining_width < ext_size) ? remaining_width : qi->unit_size;
+ // Check for special cases near boundary.
+ const bool is_horz_partitioning_allowed =
+ (this_unit_height >= qi->unit_size);
+ const bool is_vert_partitioning_allowed =
+ (this_unit_width >= qi->unit_size);
+ if (!is_horz_partitioning_allowed && !is_vert_partitioning_allowed) {
+ // Implicitly no split, and single unit info will be signaled.
+#ifndef NDEBUG
+ ++unit_info_length;
+#endif // NDEBUG
+ } else {
+ assert(split_info_index < qi->split_info_length);
+ // Write split info, depending on how many partition types are allowed.
+ const bool all_partition_types_allowed =
+ (is_horz_partitioning_allowed && is_vert_partitioning_allowed);
+ const GuidedQuadTreePartitionType partition_type =
+ qi->split_info[split_info_index].split;
+ if (all_partition_types_allowed) {
+ aom_write_symbol(wb, partition_type, ctx->cnn_guided_quad_cdf,
+ GUIDED_QT_TYPES);
+ } else {
+ aom_write_symbol(wb, partition_type != GUIDED_QT_NONE,
+ ctx->cnn_guided_binary_cdf, 2);
+ }
+#ifndef NDEBUG
+ // Look at the split info to determine number of (sub)units.
+ switch (partition_type) {
+ case GUIDED_QT_NONE: unit_info_length += 1; break;
+ case GUIDED_QT_HORZ:
+ case GUIDED_QT_VERT: unit_info_length += 2; break;
+ case GUIDED_QT_SPLIT: unit_info_length += 4; break;
+ default: assert(0 && "Wrong guided quadtree split type."); break;
+ }
+#endif // NDEBUG
+ ++split_info_index;
+ }
+ col += this_unit_width;
+ }
+ row += this_unit_height;
+ }
+ assert(qi->split_info_length == split_info_index);
+#ifndef NDEBUG
+ assert(qi->unit_info_length == unit_info_length);
+#endif // NDEBUG
+}
+
static void write_filter_quadtree(FRAME_CONTEXT *ctx, int QP, int cnn_index,
- int superres_denom, int is_intra_only,
- const QUADInfo *ci, aom_writer *wb) {
+ int width, int height, int superres_denom,
+ int is_intra_only, const QUADInfo *qi,
+ aom_writer *wb) {
+ // Write partitioning info.
+ write_quadtree_split_info(ctx, width, height, qi, wb);
+
+ // Write weight parameters 'a'.
const int *const quadtset =
get_quadparm_from_qindex(QP, superres_denom, is_intra_only, 1, cnn_index);
const int norestore_ctx =
@@ -3073,9 +3128,9 @@
const int A1_min = quadtset[3];
int ref_0 = GUIDED_A_MID;
int ref_1 = GUIDED_A_MID;
- for (int i = 0; i < ci->unit_info_length; i++) {
- const int a0 = ci->unit_info[i].xqd[0];
- const int a1 = ci->unit_info[i].xqd[1];
+ for (int i = 0; i < qi->unit_info_length; i++) {
+ const int a0 = qi->unit_info[i].xqd[0];
+ const int a1 = qi->unit_info[i].xqd[1];
int norestore;
if (norestore_ctx != -1) {
norestore = (a0 == 0 && a1 == 0);
@@ -3107,8 +3162,8 @@
// Write quad tree unit index.
static INLINE void quad_tree_write_unit_index(struct aom_write_bit_buffer *wb,
const QUADInfo *const qi) {
- assert(qi->unit_index >= 0 && qi->unit_index <= 1);
- aom_wb_write_bit(wb, qi->unit_index);
+ assert(qi->unit_index >= 0 && qi->unit_index < GUIDED_QT_UNIT_SIZES);
+ aom_wb_write_literal(wb, qi->unit_index, GUIDED_QT_UNIT_SIZES_LOG2);
}
#endif // CONFIG_CNN_GUIDED_QUADTREE
diff --git a/av1/encoder/block.h b/av1/encoder/block.h
index 02cc6a3..d740110 100644
--- a/av1/encoder/block.h
+++ b/av1/encoder/block.h
@@ -884,6 +884,10 @@
*/
int cnn_guided_quad_cost[4];
/*!
+ * cnn_guided binary split cost
+ */
+ int cnn_guided_binary_cost[2];
+ /*!
* cnn_guided quad norestore cost
*/
int cnn_guided_norestore_cost[2][2];
diff --git a/av1/encoder/encodeframe_utils.c b/av1/encoder/encodeframe_utils.c
index 8f5780b..20e1191 100644
--- a/av1/encoder/encodeframe_utils.c
+++ b/av1/encoder/encodeframe_utils.c
@@ -1271,6 +1271,8 @@
#endif // CONFIG_RST_MERGECEOFFS
#if CONFIG_CNN_GUIDED_QUADTREE
AVERAGE_CDF(ctx_left->cnn_guided_quad_cdf, ctx_tr->cnn_guided_quad_cdf, 4);
+ AVERAGE_CDF(ctx_left->cnn_guided_binary_cdf, ctx_tr->cnn_guided_binary_cdf,
+ 2);
AVERAGE_CDF(ctx_left->cnn_guided_norestore_cdf,
ctx_tr->cnn_guided_norestore_cdf, 2);
#endif // CONFIG_CNN_GUIDED_QUADTREE
diff --git a/av1/encoder/encoder.c b/av1/encoder/encoder.c
index 0b7e69a..e47994d 100644
--- a/av1/encoder/encoder.c
+++ b/av1/encoder/encoder.c
@@ -2339,6 +2339,7 @@
if (!av1_restore_cnn_quadtree_encode_tflite(
cm, cpi->source, cpi->rd.RDMULT,
x->mode_costs.cnn_guided_quad_cost,
+ x->mode_costs.cnn_guided_binary_cost,
x->mode_costs.cnn_guided_norestore_cost,
cpi->mt_info.num_workers, quadtree_cnn, curr_cnn_indices,
&quad_info, &curr_cnn_rdcosts)) {
diff --git a/av1/encoder/rd.c b/av1/encoder/rd.c
index ebb64df..d8c2073 100644
--- a/av1/encoder/rd.c
+++ b/av1/encoder/rd.c
@@ -455,6 +455,8 @@
#if CONFIG_CNN_GUIDED_QUADTREE
av1_cost_tokens_from_cdf(mode_costs->cnn_guided_quad_cost,
fc->cnn_guided_quad_cdf, NULL);
+ av1_cost_tokens_from_cdf(mode_costs->cnn_guided_binary_cost,
+ fc->cnn_guided_binary_cdf, NULL);
for (int i = 0; i < GUIDED_NORESTORE_CONTEXTS; ++i)
av1_cost_tokens_from_cdf(mode_costs->cnn_guided_norestore_cost[i],
fc->cnn_guided_norestore_cdf[i], NULL);
@@ -502,6 +504,8 @@
#if CONFIG_CNN_GUIDED_QUADTREE
av1_cost_tokens_from_cdf(mode_costs->cnn_guided_quad_cost,
fc->cnn_guided_quad_cdf, NULL);
+ av1_cost_tokens_from_cdf(mode_costs->cnn_guided_binary_cost,
+ fc->cnn_guided_binary_cdf, NULL);
for (int c = 0; c < GUIDED_NORESTORE_CONTEXTS; ++c)
av1_cost_tokens_from_cdf(mode_costs->cnn_guided_norestore_cost[c],
fc->cnn_guided_norestore_cdf[c], NULL);
