Increase the transform block number in context_tree allocation
The pixel number remains the same.
Change-Id: I40e54f4c0893db8f9452624f8167ead994b71f1d
diff --git a/av1/encoder/context_tree.c b/av1/encoder/context_tree.c
index e1db4be..a0ab2e6 100644
--- a/av1/encoder/context_tree.c
+++ b/av1/encoder/context_tree.c
@@ -25,7 +25,7 @@
#endif
PICK_MODE_CONTEXT *ctx) {
const int num_blk = (num_4x4_blk < 4 ? 4 : num_4x4_blk);
- const int num_pix = num_blk << 4;
+ const int num_pix = num_blk * tx_size_2d[0];
int i;
ctx->num_4x4_blk = num_blk;
#if CONFIG_EXT_PARTITION_TYPES
@@ -228,7 +228,11 @@
for (pc_tree_index = 0; pc_tree_index < leaf_nodes; ++pc_tree_index) {
PC_TREE *const tree = &td->pc_tree[pc_tree_index];
tree->block_size = square[0];
+#if CONFIG_CB4X4
+ alloc_tree_contexts(cm, tree, 16);
+#else
alloc_tree_contexts(cm, tree, 4);
+#endif
tree->leaf_split[0] = this_leaf++;
for (j = 1; j < 4; j++) tree->leaf_split[j] = tree->leaf_split[0];
}
@@ -238,7 +242,11 @@
for (nodes = leaf_nodes >> 2; nodes > 0; nodes >>= 2) {
for (i = 0; i < nodes; ++i) {
PC_TREE *const tree = &td->pc_tree[pc_tree_index];
+#if CONFIG_CB4X4
+ alloc_tree_contexts(cm, tree, 16 << (2 * square_index));
+#else
alloc_tree_contexts(cm, tree, 4 << (2 * square_index));
+#endif
tree->block_size = square[square_index];
for (j = 0; j < 4; j++) tree->split[j] = this_pc++;
++pc_tree_index;