Merge "Add depth to tx_size mapper to bit-stream coding" into nextgenv2

av1/common/blockd.h

@@ -712,6 +712,14 @@ static INLINE TX_TYPE get_tx_type(PLANE_TYPE plane_type, const MACROBLOCKD *xd,
 
 void av1_setup_block_planes(MACROBLOCKD *xd, int ss_x, int ss_y);
 
+static INLINE int tx_size_to_depth(const TX_SIZE tx_size) {
+  return (int)(tx_size - TX_4X4);
+}
+
+static INLINE TX_SIZE depth_to_tx_size(const int depth) {
+  return (TX_SIZE)(depth + TX_4X4);
+}
+
 static INLINE TX_SIZE get_uv_tx_size(const MB_MODE_INFO *mbmi,
                                      const struct macroblockd_plane *pd) {
   TX_SIZE uv_txsize;

av1/decoder/decodemv.c

@@ -350,11 +350,11 @@ static TX_SIZE read_selected_tx_size(AV1_COMMON *cm, MACROBLOCKD *xd,
                                      int tx_size_cat, aom_reader *r) {
   FRAME_COUNTS *counts = xd->counts;
   const int ctx = get_tx_size_context(xd);
-  int tx_size =
-      aom_read_tree(r, av1_tx_size_tree[tx_size_cat],
-                    cm->fc->tx_size_probs[tx_size_cat][ctx], ACCT_STR);
+  int depth = aom_read_tree(r, av1_tx_size_tree[tx_size_cat],
+                            cm->fc->tx_size_probs[tx_size_cat][ctx], ACCT_STR);
+  TX_SIZE tx_size = depth_to_tx_size(depth);
   if (counts) ++counts->tx_size[tx_size_cat][ctx][tx_size];
-  return (TX_SIZE)tx_size;
+  return tx_size;
 }
 
 static TX_SIZE read_tx_size_intra(AV1_COMMON *cm, MACROBLOCKD *xd,

av1/encoder/bitstream.c

@@ -414,6 +414,7 @@ static void write_selected_tx_size(const AV1_COMMON *cm, const MACROBLOCKD *xd,
     const int tx_size_cat = is_inter ? inter_tx_size_cat_lookup[bsize]
                                      : intra_tx_size_cat_lookup[bsize];
     const TX_SIZE coded_tx_size = txsize_sqr_up_map[tx_size];
+    const int depth = tx_size_to_depth(coded_tx_size);
 
 #if CONFIG_EXT_TX && CONFIG_RECT_TX
     assert(IMPLIES(is_rect_tx(tx_size), is_rect_tx_allowed(xd, mbmi)));
@@ -423,7 +424,7 @@ static void write_selected_tx_size(const AV1_COMMON *cm, const MACROBLOCKD *xd,
 
     av1_write_token(w, av1_tx_size_tree[tx_size_cat],
                     cm->fc->tx_size_probs[tx_size_cat][tx_size_ctx],
-                    &tx_size_encodings[tx_size_cat][coded_tx_size]);
+                    &tx_size_encodings[tx_size_cat][depth]);
   }
 }
 

av1/encoder/encodeframe.c

@@ -5206,6 +5206,7 @@ static void encode_superblock(const AV1_COMP *const cpi, ThreadData *td,
       const int tx_size_cat = is_inter ? inter_tx_size_cat_lookup[bsize]
                                        : intra_tx_size_cat_lookup[bsize];
       const TX_SIZE coded_tx_size = txsize_sqr_up_map[mbmi->tx_size];
+      const int depth = tx_size_to_depth(coded_tx_size);
 #if CONFIG_EXT_TX && CONFIG_RECT_TX
       assert(IMPLIES(is_rect_tx(mbmi->tx_size), is_rect_tx_allowed(xd, mbmi)));
 #endif  // CONFIG_EXT_TX && CONFIG_RECT_TX
@@ -5219,7 +5220,7 @@ static void encode_superblock(const AV1_COMP *const cpi, ThreadData *td,
         if (is_inter) {
           tx_partition_count_update(cm, x, bsize, mi_row, mi_col, td->counts);
         } else {
-          ++td->counts->tx_size[tx_size_cat][tx_size_ctx][coded_tx_size];
+          ++td->counts->tx_size[tx_size_cat][tx_size_ctx][depth];
           if (mbmi->tx_size != max_txsize_lookup[bsize]) ++x->txb_split_count;
         }
 #if CONFIG_EXT_TX && CONFIG_RECT_TX
@@ -5227,7 +5228,7 @@ static void encode_superblock(const AV1_COMP *const cpi, ThreadData *td,
 #endif
 #endif
 #if !CONFIG_VAR_TX
-      ++td->counts->tx_size[tx_size_cat][tx_size_ctx][coded_tx_size];
+      ++td->counts->tx_size[tx_size_cat][tx_size_ctx][depth];
 #endif
     } else {
       int i, j;