Rework transform block partition context model

This commit allows the partition context model to account for the
maximum transform block size of the coding block.

Change-Id: I22b91e85fff70faa974afd362ce327d3f2eda81d

diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index 201bb169d0f3ce14948cbaaa457444e0f16da52d..cea5769a43fd561f3d73b7f4d223e685356fc72c 100644 (file)
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -984,7 +984,7 @@ int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
 
 #if CONFIG_VAR_TX
 static const aom_prob default_txfm_partition_probs[TXFM_PARTITION_CONTEXTS] = {
-  192, 128, 64, 192, 128, 64, 192, 128, 64,
+  250, 231, 212, 241, 166, 66, 241, 230, 135, 243, 154, 64, 248, 161, 63, 128,
 };
 #endif
 

diff --git a/av1/common/enums.h b/av1/common/enums.h
index a684eedff17622e6896351563bd5e9b543550d1b..e274d4b9311573648c689be418f0cb2b2c85d1f3 100644 (file)
--- a/av1/common/enums.h
+++ b/av1/common/enums.h
@@ -408,7 +408,7 @@ typedef enum {
 #define REF_CONTEXTS 5
 
 #if CONFIG_VAR_TX
-#define TXFM_PARTITION_CONTEXTS 9
+#define TXFM_PARTITION_CONTEXTS 16
 typedef TX_SIZE TXFM_CONTEXT;
 #endif
 

diff --git a/av1/common/onyxc_int.h b/av1/common/onyxc_int.h
index 7c12bd358b73294feb64868c862908dac4f9bae5..c0ed53818ef4fdae2d432eadd6f1c370b778eede 100644 (file)
--- a/av1/common/onyxc_int.h
+++ b/av1/common/onyxc_int.h
@@ -756,11 +756,30 @@ static INLINE void txfm_partition_update(TXFM_CONTEXT *above_ctx,
 
 static INLINE int txfm_partition_context(TXFM_CONTEXT *above_ctx,
                                          TXFM_CONTEXT *left_ctx,
-                                         TX_SIZE tx_size) {
-  int above = *above_ctx < tx_size;
-  int left = *left_ctx < tx_size;
+                                         const BLOCK_SIZE bsize,
+                                         const TX_SIZE tx_size) {
+  const int above = *above_ctx < tx_size;
+  const int left = *left_ctx < tx_size;
+  TX_SIZE max_tx_size = max_txsize_lookup[bsize];
+  int category = 15;
+
+  if (max_tx_size == TX_32X32) {
+    if (tx_size == TX_32X32)
+      category = 0;
+    else
+      category = 1;
+  } else if (max_tx_size == TX_16X16) {
+    if (tx_size == TX_16X16)
+      category = 2;
+    else
+      category = 3;
+  } else if (max_tx_size == TX_8X8) {
+    category = 4;
+  }
+
+  if (category == 15) return category;
 
-  return (tx_size - TX_8X8) * 3 + above + left;
+  return category * 3 + above + left;
 }
 #endif
 

diff --git a/av1/decoder/decodemv.c b/av1/decoder/decodemv.c
index 85f81112cdb4ad4a1f178f838caff7a04ede41b1..77cea8a77d77a02a44bf3e18fe18ccf3f3e4b64e 100644 (file)
--- a/av1/decoder/decodemv.c
+++ b/av1/decoder/decodemv.c
@@ -301,7 +301,8 @@ static void read_tx_size_vartx(AV1_COMMON *cm, MACROBLOCKD *xd,
   int max_blocks_high = block_size_high[mbmi->sb_type];
   int max_blocks_wide = block_size_wide[mbmi->sb_type];
   int ctx = txfm_partition_context(xd->above_txfm_context + tx_col,
-                                   xd->left_txfm_context + tx_row, tx_size);
+                                   xd->left_txfm_context + tx_row,
+                                   mbmi->sb_type, tx_size);
   TX_SIZE(*const inter_tx_size)
   [MAX_MIB_SIZE] =
       (TX_SIZE(*)[MAX_MIB_SIZE]) & mbmi->inter_tx_size[tx_row][tx_col];

diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index e232a87964c075b5cfb19dab0f8adc0ce1da8580..427509855b01170ff2d78a0a3a1b0bf45dbfdb4b 100644 (file)
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -360,7 +360,8 @@ static void write_tx_size_vartx(const AV1_COMMON *cm, const MACROBLOCKD *xd,
   const int max_blocks_wide = max_block_wide(xd, mbmi->sb_type, 0);
 
   int ctx = txfm_partition_context(xd->above_txfm_context + tx_col,
-                                   xd->left_txfm_context + tx_row, tx_size);
+                                   xd->left_txfm_context + tx_row,
+                                   mbmi->sb_type, tx_size);
 
   if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
 

diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 5ad334dda911f1f8f431ea48edcdaff8d0f75ee0..d1d6ecc57810e0d624d01be3175ca651493512da 100644 (file)
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -4987,7 +4987,7 @@ static void sum_intra_stats(FRAME_COUNTS *counts, const MODE_INFO *mi,
 
 #if CONFIG_VAR_TX
 static void update_txfm_count(MACROBLOCK *x, MACROBLOCKD *xd,
-                              FRAME_COUNTS *counts, TX_SIZE tx_size,
+                              FRAME_COUNTS *counts, TX_SIZE tx_size, int depth,
                               int blk_row, int blk_col) {
   MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
   const int tx_row = blk_row >> 1;
@@ -4995,7 +4995,8 @@ static void update_txfm_count(MACROBLOCK *x, MACROBLOCKD *xd,
   const int max_blocks_high = max_block_high(xd, mbmi->sb_type, 0);
   const int max_blocks_wide = max_block_wide(xd, mbmi->sb_type, 0);
   int ctx = txfm_partition_context(xd->above_txfm_context + tx_col,
-                                   xd->left_txfm_context + tx_row, tx_size);
+                                   xd->left_txfm_context + tx_row,
+                                   mbmi->sb_type, tx_size);
   const TX_SIZE plane_tx_size = mbmi->inter_tx_size[tx_row][tx_col];
 
   if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
@@ -5023,8 +5024,8 @@ static void update_txfm_count(MACROBLOCK *x, MACROBLOCKD *xd,
     for (i = 0; i < 4; ++i) {
       int offsetr = (i >> 1) * bh / 2;
       int offsetc = (i & 0x01) * bh / 2;
-      update_txfm_count(x, xd, counts, tx_size - 1, blk_row + offsetr,
-                        blk_col + offsetc);
+      update_txfm_count(x, xd, counts, tx_size - 1, depth + 1,
+                        blk_row + offsetr, blk_col + offsetc);
     }
   }
 }
@@ -5046,7 +5047,8 @@ static void tx_partition_count_update(const AV1_COMMON *const cm, MACROBLOCK *x,
 
   for (idy = 0; idy < mi_height; idy += bh)
     for (idx = 0; idx < mi_width; idx += bh)
-      update_txfm_count(x, xd, td_counts, max_tx_size, idy, idx);
+      update_txfm_count(x, xd, td_counts, max_tx_size, mi_width != mi_height,
+                        idy, idx);
 }
 
 static void set_txfm_context(MACROBLOCKD *xd, TX_SIZE tx_size, int blk_row,

diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index 46dbae49a68e670c748a8a70de3203d11ca3007e..b37714cbe78dde0eba0d5ca95b01dfa3776a2271 100644 (file)
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -3064,8 +3064,9 @@ static void select_tx_block(const AV1_COMP *cpi, MACROBLOCK *x, int blk_row,
   ENTROPY_CONTEXT *pta = ta + blk_col;
   ENTROPY_CONTEXT *ptl = tl + blk_row;
   int coeff_ctx, i;
-  int ctx = txfm_partition_context(tx_above + (blk_col >> 1),
-                                   tx_left + (blk_row >> 1), tx_size);
+  int ctx =
+      txfm_partition_context(tx_above + (blk_col >> 1),
+                             tx_left + (blk_row >> 1), mbmi->sb_type, tx_size);
 
   int64_t sum_dist = 0, sum_bsse = 0;
   int64_t sum_rd = INT64_MAX;