Merge "Enable tile-adaptive restoration" into nextgenv2

aom_dsp/psnr.c

@@ -177,6 +177,14 @@ static int64_t highbd_get_sse(const uint8_t *a, int a_stride, const uint8_t *b,
 }
 #endif  // CONFIG_AOM_HIGHBITDEPTH
 
+int64_t aom_get_y_sse_part(const YV12_BUFFER_CONFIG *a,
+                           const YV12_BUFFER_CONFIG *b,
+                           int hstart, int width, int vstart, int height) {
+  return get_sse(a->y_buffer + vstart * a->y_stride + hstart, a->y_stride,
+                 b->y_buffer + vstart * b->y_stride + hstart, b->y_stride,
+                 width, height);
+}
+
 int64_t aom_get_y_sse(const YV12_BUFFER_CONFIG *a,
                       const YV12_BUFFER_CONFIG *b) {
   assert(a->y_crop_width == b->y_crop_width);
@@ -205,6 +213,16 @@ int64_t aom_get_v_sse(const YV12_BUFFER_CONFIG *a,
 }
 
 #if CONFIG_AOM_HIGHBITDEPTH
+int64_t aom_highbd_get_y_sse_part(const YV12_BUFFER_CONFIG *a,
+                                  const YV12_BUFFER_CONFIG *b,
+                                  int hstart, int width,
+                                  int vstart, int height) {
+  return highbd_get_sse(
+      a->y_buffer + vstart * a->y_stride + hstart, a->y_stride,
+      b->y_buffer + vstart * b->y_stride + hstart, b->y_stride,
+      width, height);
+}
+
 int64_t aom_highbd_get_y_sse(const YV12_BUFFER_CONFIG *a,
                              const YV12_BUFFER_CONFIG *b) {
   assert(a->y_crop_width == b->y_crop_width);

aom_dsp/psnr.h

@@ -35,10 +35,17 @@ typedef struct {
 * \param[in]    sse           Sum of squared errors
 */
 double aom_sse_to_psnr(double samples, double peak, double sse);
+int64_t aom_get_y_sse_part(const YV12_BUFFER_CONFIG *a,
+                           const YV12_BUFFER_CONFIG *b,
+                           int hstart, int width, int vstart, int height);
 int64_t aom_get_y_sse(const YV12_BUFFER_CONFIG *a, const YV12_BUFFER_CONFIG *b);
 int64_t aom_get_u_sse(const YV12_BUFFER_CONFIG *a, const YV12_BUFFER_CONFIG *b);
 int64_t aom_get_v_sse(const YV12_BUFFER_CONFIG *a, const YV12_BUFFER_CONFIG *b);
 #if CONFIG_AOM_HIGHBITDEPTH
+int64_t aom_highbd_get_y_sse_part(const YV12_BUFFER_CONFIG *a,
+                                  const YV12_BUFFER_CONFIG *b,
+                                  int hstart, int width,
+                                  int vstart, int height);
 int64_t aom_highbd_get_y_sse(const YV12_BUFFER_CONFIG *a,
                              const YV12_BUFFER_CONFIG *b);
 int64_t aom_highbd_get_u_sse(const YV12_BUFFER_CONFIG *a,

av1/common/alloccommon.c

@@ -83,6 +83,8 @@ void av1_free_ref_frame_buffers(BufferPool *pool) {
 
 #if CONFIG_LOOP_RESTORATION
 void av1_free_restoration_buffers(AV1_COMMON *cm) {
+  aom_free(cm->rst_info.restoration_type);
+  cm->rst_info.restoration_type = NULL;
   aom_free(cm->rst_info.bilateral_level);
   cm->rst_info.bilateral_level = NULL;
   aom_free(cm->rst_info.vfilter);

av1/common/entropymode.c

@@ -865,6 +865,17 @@ static const aom_prob default_tx_size_prob[TX_SIZES - 1][TX_SIZE_CONTEXTS]
                                             },
                                           };
 
+#if CONFIG_LOOP_RESTORATION
+const aom_tree_index
+    av1_switchable_restore_tree[TREE_SIZE(RESTORE_SWITCHABLE_TYPES)] = {
+      -RESTORE_NONE, 2,
+      -RESTORE_BILATERAL, -RESTORE_WIENER,
+    };
+
+static const aom_prob
+    default_switchable_restore_prob[RESTORE_SWITCHABLE_TYPES - 1] = {32, 128};
+#endif  // CONFIG_LOOP_RESTORATION
+
 #if CONFIG_EXT_TX && CONFIG_RECT_TX && CONFIG_VAR_TX
 // the probability of (0) using recursive square tx partition vs.
 // (1) biggest rect tx for 4X8-8X4/8X16-16X8/16X32-32X16 blocks
@@ -1361,6 +1372,9 @@ static void init_mode_probs(FRAME_CONTEXT *fc) {
 #endif  // CONFIG_EXT_INTRA
   av1_copy(fc->inter_ext_tx_prob, default_inter_ext_tx_prob);
   av1_copy(fc->intra_ext_tx_prob, default_intra_ext_tx_prob);
+#if CONFIG_LOOP_RESTORATION
+  av1_copy(fc->switchable_restore_prob, default_switchable_restore_prob);
+#endif  // CONFIG_LOOP_RESTORATION
 }
 
 #if CONFIG_EXT_INTERP

av1/common/entropymode.h

@@ -126,6 +126,9 @@ typedef struct frame_contexts {
 #if CONFIG_GLOBAL_MOTION
   aom_prob global_motion_types_prob[GLOBAL_MOTION_TYPES - 1];
 #endif  // CONFIG_GLOBAL_MOTION
+#if CONFIG_LOOP_RESTORATION
+  aom_prob switchable_restore_prob[RESTORE_SWITCHABLE_TYPES - 1];
+#endif  // CONFIG_LOOP_RESTORATION
 } FRAME_CONTEXT;
 
 typedef struct FRAME_COUNTS {
@@ -263,6 +266,13 @@ extern const aom_tree_index av1_ext_tx_tree[TREE_SIZE(TX_TYPES)];
 extern const aom_tree_index av1_motvar_tree[TREE_SIZE(MOTION_VARIATIONS)];
 #endif  // CONFIG_OBMC || CONFIG_WARPED_MOTION
 
+#if CONFIG_LOOP_RESTORATION
+#define RESTORE_NONE_BILATERAL_PROB 16
+#define RESTORE_NONE_WIENER_PROB 64
+extern const aom_tree_index
+    av1_switchable_restore_tree[TREE_SIZE(RESTORE_SWITCHABLE_TYPES)];
+#endif  // CONFIG_LOOP_RESTORATION
+
 void av1_setup_past_independence(struct AV1Common *cm);
 
 void av1_adapt_intra_frame_probs(struct AV1Common *cm);

av1/common/enums.h

@@ -433,6 +433,16 @@ typedef TX_SIZE TXFM_CONTEXT;
 #define MAX_SUPERTX_BLOCK_SIZE BLOCK_32X32
 #endif  // CONFIG_SUPERTX
 
+#if CONFIG_LOOP_RESTORATION
+typedef enum {
+  RESTORE_NONE,
+  RESTORE_BILATERAL,
+  RESTORE_WIENER,
+  RESTORE_SWITCHABLE,
+  RESTORE_SWITCHABLE_TYPES = RESTORE_SWITCHABLE,
+  RESTORE_TYPES,
+} RestorationType;
+#endif  // CONFIG_LOOP_RESTORATION
 #ifdef __cplusplus
 }  // extern "C"
 #endif

av1/common/loopfilter.c

@@ -16,7 +16,6 @@
 #include "av1/common/loopfilter.h"
 #include "av1/common/onyxc_int.h"
 #include "av1/common/reconinter.h"
-#include "av1/common/restoration.h"
 #include "aom_dsp/aom_dsp_common.h"
 #include "aom_mem/aom_mem.h"
 #include "aom_ports/mem.h"

av1/common/loopfilter.h

@@ -16,7 +16,6 @@
 #include "./aom_config.h"
 
 #include "av1/common/blockd.h"
-#include "av1/common/restoration.h"
 #include "av1/common/seg_common.h"
 
 #ifdef __cplusplus

av1/common/onyxc_int.h

@@ -25,7 +25,9 @@
 #include "av1/common/frame_buffers.h"
 #include "av1/common/quant_common.h"
 #include "av1/common/tile_common.h"
+#if CONFIG_LOOP_RESTORATION
 #include "av1/common/restoration.h"
+#endif  // CONFIG_LOOP_RESTORATION
 
 #ifdef __cplusplus
 extern "C" {

av1/common/restoration.h

@@ -26,9 +26,10 @@ extern "C" {
 #define BILATERAL_LEVELS (1 << BILATERAL_LEVEL_BITS)
 // #define DEF_BILATERAL_LEVEL     2
 
-#define RESTORATION_TILESIZES 3
-#define BILATERAL_TILESIZE 1
-#define WIENER_TILESIZE 2
+#define RESTORATION_TILESIZE_SML 128
+#define RESTORATION_TILESIZE_BIG 256
+#define BILATERAL_SUBTILE_BITS 1
+#define BILATERAL_SUBTILES (1 << (2 * BILATERAL_SUBTILE_BITS))
 
 #define RESTORATION_HALFWIN 3
 #define RESTORATION_HALFWIN1 (RESTORATION_HALFWIN + 1)
@@ -56,43 +57,84 @@ extern "C" {
 #define WIENER_FILT_TAP2_MAXV \
   (WIENER_FILT_TAP2_MINV - 1 + (1 << WIENER_FILT_TAP2_BITS))
 
-typedef enum {
-  RESTORE_NONE,
-  RESTORE_BILATERAL,
-  RESTORE_WIENER,
-} RestorationType;
-
 typedef struct {
-  RestorationType restoration_type;
+  RestorationType frame_restoration_type;
+  RestorationType *restoration_type;
   // Bilateral filter
   int *bilateral_level;
   // Wiener filter
   int *wiener_level;
-  int (*vfilter)[RESTORATION_HALFWIN], (*hfilter)[RESTORATION_HALFWIN];
+  int (*vfilter)[RESTORATION_WIN], (*hfilter)[RESTORATION_WIN];
 } RestorationInfo;
 
 typedef struct {
-  RestorationType restoration_type;
+  RestorationInfo *rsi;
+  int keyframe;
   int subsampling_x;
   int subsampling_y;
-  int tilesize_index;
   int ntiles;
   int tile_width, tile_height;
   int nhtiles, nvtiles;
-  // Bilateral filter
-  int *bilateral_level;
-  uint8_t (**wx_lut)[RESTORATION_WIN];
-  uint8_t **wr_lut;
-  // Wiener filter
-  int *wiener_level;
-  int (*vfilter)[RESTORATION_WIN], (*hfilter)[RESTORATION_WIN];
 } RestorationInternal;
 
+static INLINE int get_rest_tilesize(int width, int height) {
+  if (width * height <= 352 * 288)
+    return RESTORATION_TILESIZE_SML;
+  else
+    return RESTORATION_TILESIZE_BIG;
+}
+
+static INLINE int av1_get_rest_ntiles(int width, int height,
+                                      int *tile_width, int *tile_height,
+                                      int *nhtiles, int *nvtiles) {
+  int nhtiles_, nvtiles_;
+  int tile_width_, tile_height_;
+  int tilesize = get_rest_tilesize(width, height);
+  tile_width_ = (tilesize < 0) ? width : AOMMIN(tilesize, width);
+  tile_height_ = (tilesize < 0) ? height : AOMMIN(tilesize, height);
+  nhtiles_ = (width + (tile_width_ >> 1)) / tile_width_;
+  nvtiles_ = (height + (tile_height_ >> 1)) / tile_height_;
+  if (tile_width) *tile_width = tile_width_;
+  if (tile_height) *tile_height = tile_height_;
+  if (nhtiles) *nhtiles = nhtiles_;
+  if (nvtiles) *nvtiles = nvtiles_;
+  return (nhtiles_ * nvtiles_);
+}
+
+static INLINE void av1_get_rest_tile_limits(
+    int tile_idx, int subtile_idx, int subtile_bits, int nhtiles, int nvtiles,
+    int tile_width, int tile_height, int im_width, int im_height, int clamp_h,
+    int clamp_v, int *h_start, int *h_end, int *v_start, int *v_end) {
+  const int htile_idx = tile_idx % nhtiles;
+  const int vtile_idx = tile_idx / nhtiles;
+  *h_start = htile_idx * tile_width;
+  *v_start = vtile_idx * tile_height;
+  *h_end = (htile_idx < nhtiles - 1) ? *h_start + tile_width : im_width;
+  *v_end = (vtile_idx < nvtiles - 1) ? *v_start + tile_height : im_height;
+  if (subtile_bits) {
+    const int num_subtiles_1d = (1 << subtile_bits);
+    const int subtile_width = (*h_end - *h_start) >> subtile_bits;
+    const int subtile_height = (*v_end - *v_start) >> subtile_bits;
+    const int subtile_idx_h = subtile_idx & (num_subtiles_1d - 1);
+    const int subtile_idx_v = subtile_idx >> subtile_bits;
+    *h_start += subtile_idx_h * subtile_width;
+    *v_start += subtile_idx_v * subtile_height;
+    *h_end = subtile_idx_h == num_subtiles_1d - 1 ? *h_end
+                                                  : *h_start + subtile_width;
+    *v_end = subtile_idx_v == num_subtiles_1d - 1 ? *v_end
+                                                  : *v_start + subtile_height;
+  }
+  if (clamp_h) {
+    *h_start = AOMMAX(*h_start, RESTORATION_HALFWIN);
+    *h_end = AOMMIN(*h_end, im_width - RESTORATION_HALFWIN);
+  }
+  if (clamp_v) {
+    *v_start = AOMMAX(*v_start, RESTORATION_HALFWIN);
+    *v_end = AOMMIN(*v_end, im_height - RESTORATION_HALFWIN);
+  }
+}
+
 int av1_bilateral_level_bits(const struct AV1Common *const cm);
-int av1_get_restoration_ntiles(int tilesize, int width, int height);
-void av1_get_restoration_tile_size(int tilesize, int width, int height,
-                                   int *tile_width, int *tile_height,
-                                   int *nhtiles, int *nvtiles);
 void av1_loop_restoration_init(RestorationInternal *rst, RestorationInfo *rsi,
                                int kf, int width, int height);
 void av1_loop_restoration_frame(YV12_BUFFER_CONFIG *frame, struct AV1Common *cm,

av1/decoder/decodeframe.c

@@ -1899,62 +1899,134 @@ static void setup_segmentation(AV1_COMMON *const cm,
 }
 
 #if CONFIG_LOOP_RESTORATION
-static void setup_restoration(AV1_COMMON *cm, struct aom_read_bit_buffer *rb) {
-  int i;
+static void decode_restoration_mode(AV1_COMMON *cm,
+                                    struct aom_read_bit_buffer *rb) {
   RestorationInfo *rsi = &cm->rst_info;
-  int ntiles;
   if (aom_rb_read_bit(rb)) {
-    if (aom_rb_read_bit(rb)) {
-      rsi->restoration_type = RESTORE_BILATERAL;
-      ntiles =
-          av1_get_restoration_ntiles(BILATERAL_TILESIZE, cm->width, cm->height);
+    rsi->frame_restoration_type =
+        aom_rb_read_bit(rb) ? RESTORE_WIENER : RESTORE_BILATERAL;
+  } else {
+    rsi->frame_restoration_type =
+        aom_rb_read_bit(rb) ? RESTORE_SWITCHABLE : RESTORE_NONE;
+  }
+}
+
+static void decode_restoration(AV1_COMMON *cm, aom_reader *rb) {
+  int i;
+  RestorationInfo *rsi = &cm->rst_info;
+  const int ntiles = av1_get_rest_ntiles(cm->width, cm->height,
+                                         NULL, NULL, NULL, NULL);
+  if (rsi->frame_restoration_type != RESTORE_NONE) {
+    rsi->restoration_type = (RestorationType *)aom_realloc(
+        rsi->restoration_type, sizeof(*rsi->restoration_type) * ntiles);
+    if (rsi->frame_restoration_type == RESTORE_SWITCHABLE) {
       rsi->bilateral_level = (int *)aom_realloc(
-          rsi->bilateral_level, sizeof(*rsi->bilateral_level) * ntiles);
+          rsi->bilateral_level,
+          sizeof(*rsi->bilateral_level) * ntiles * BILATERAL_SUBTILES);
       assert(rsi->bilateral_level != NULL);
+      rsi->wiener_level = (int *)aom_realloc(
+          rsi->wiener_level, sizeof(*rsi->wiener_level) * ntiles);
+      assert(rsi->wiener_level != NULL);
+      rsi->vfilter = (int(*)[RESTORATION_WIN])aom_realloc(
+          rsi->vfilter, sizeof(*rsi->vfilter) * ntiles);
+      assert(rsi->vfilter != NULL);
+      rsi->hfilter = (int(*)[RESTORATION_WIN])aom_realloc(
+          rsi->hfilter, sizeof(*rsi->hfilter) * ntiles);
+      assert(rsi->hfilter != NULL);
       for (i = 0; i < ntiles; ++i) {
-        if (aom_rb_read_bit(rb)) {
-          rsi->bilateral_level[i] =
-              aom_rb_read_literal(rb, av1_bilateral_level_bits(cm));
-        } else {
-          rsi->bilateral_level[i] = -1;
+        rsi->restoration_type[i] = aom_read_tree(
+            rb, av1_switchable_restore_tree, cm->fc->switchable_restore_prob);
+        if (rsi->restoration_type[i] == RESTORE_WIENER) {
+          rsi->wiener_level[i] = 1;
+          rsi->vfilter[i][0] =
+              aom_read_literal(rb, WIENER_FILT_TAP0_BITS) +
+              WIENER_FILT_TAP0_MINV;
+          rsi->vfilter[i][1] =
+              aom_read_literal(rb, WIENER_FILT_TAP1_BITS) +
+              WIENER_FILT_TAP1_MINV;
+          rsi->vfilter[i][2] =
+              aom_read_literal(rb, WIENER_FILT_TAP2_BITS) +
+              WIENER_FILT_TAP2_MINV;
+          rsi->hfilter[i][0] =
+              aom_read_literal(rb, WIENER_FILT_TAP0_BITS) +
+              WIENER_FILT_TAP0_MINV;
+          rsi->hfilter[i][1] =
+              aom_read_literal(rb, WIENER_FILT_TAP1_BITS) +
+              WIENER_FILT_TAP1_MINV;
+          rsi->hfilter[i][2] =
+              aom_read_literal(rb, WIENER_FILT_TAP2_BITS) +
+              WIENER_FILT_TAP2_MINV;
+        } else if (rsi->restoration_type[i] == RESTORE_BILATERAL) {
+          int s;
+          for (s = 0; s < BILATERAL_SUBTILES; ++s) {
+            const int j = i * BILATERAL_SUBTILES + s;
+#if BILATERAL_SUBTILES == 0
+            rsi->bilateral_level[j] =
+                aom_read_literal(rb, av1_bilateral_level_bits(cm));
+#else
+            if (aom_read(rb, RESTORE_NONE_BILATERAL_PROB)) {
+              rsi->bilateral_level[j] =
+                  aom_read_literal(rb, av1_bilateral_level_bits(cm));
+            } else {
+              rsi->bilateral_level[j] = -1;
+            }
+#endif
+          }
         }
       }
-    } else {
-      rsi->restoration_type = RESTORE_WIENER;
-      ntiles =
-          av1_get_restoration_ntiles(WIENER_TILESIZE, cm->width, cm->height);
+    } else if (rsi->frame_restoration_type == RESTORE_WIENER) {
       rsi->wiener_level = (int *)aom_realloc(
           rsi->wiener_level, sizeof(*rsi->wiener_level) * ntiles);
       assert(rsi->wiener_level != NULL);
-      rsi->vfilter = (int(*)[RESTORATION_HALFWIN])aom_realloc(
+      rsi->vfilter = (int(*)[RESTORATION_WIN])aom_realloc(
           rsi->vfilter, sizeof(*rsi->vfilter) * ntiles);
       assert(rsi->vfilter != NULL);
-      rsi->hfilter = (int(*)[RESTORATION_HALFWIN])aom_realloc(
+      rsi->hfilter = (int(*)[RESTORATION_WIN])aom_realloc(
           rsi->hfilter, sizeof(*rsi->hfilter) * ntiles);
       assert(rsi->hfilter != NULL);
       for (i = 0; i < ntiles; ++i) {
-        rsi->wiener_level[i] = aom_rb_read_bit(rb);
-        if (rsi->wiener_level[i]) {
-          rsi->vfilter[i][0] = aom_rb_read_literal(rb, WIENER_FILT_TAP0_BITS) +
+        if (aom_read(rb, RESTORE_NONE_WIENER_PROB)) {
+          rsi->wiener_level[i] = 1;
+          rsi->restoration_type[i] = RESTORE_WIENER;
+          rsi->vfilter[i][0] = aom_read_literal(rb, WIENER_FILT_TAP0_BITS) +
                                WIENER_FILT_TAP0_MINV;
-          rsi->vfilter[i][1] = aom_rb_read_literal(rb, WIENER_FILT_TAP1_BITS) +
+          rsi->vfilter[i][1] = aom_read_literal(rb, WIENER_FILT_TAP1_BITS) +
                                WIENER_FILT_TAP1_MINV;
-          rsi->vfilter[i][2] = aom_rb_read_literal(rb, WIENER_FILT_TAP2_BITS) +
+          rsi->vfilter[i][2] = aom_read_literal(rb, WIENER_FILT_TAP2_BITS) +
                                WIENER_FILT_TAP2_MINV;
-          rsi->hfilter[i][0] = aom_rb_read_literal(rb, WIENER_FILT_TAP0_BITS) +
+          rsi->hfilter[i][0] = aom_read_literal(rb, WIENER_FILT_TAP0_BITS) +
                                WIENER_FILT_TAP0_MINV;
-          rsi->hfilter[i][1] = aom_rb_read_literal(rb, WIENER_FILT_TAP1_BITS) +
+          rsi->hfilter[i][1] = aom_read_literal(rb, WIENER_FILT_TAP1_BITS) +
                                WIENER_FILT_TAP1_MINV;
-          rsi->hfilter[i][2] = aom_rb_read_literal(rb, WIENER_FILT_TAP2_BITS) +
+          rsi->hfilter[i][2] = aom_read_literal(rb, WIENER_FILT_TAP2_BITS) +
                                WIENER_FILT_TAP2_MINV;
         } else {
-          rsi->vfilter[i][0] = rsi->vfilter[i][1] = rsi->vfilter[i][2] = 0;
-          rsi->hfilter[i][0] = rsi->hfilter[i][1] = rsi->hfilter[i][2] = 0;
+          rsi->wiener_level[i] = 0;
+          rsi->restoration_type[i] = RESTORE_NONE;
+        }
+      }
+    } else {
+      rsi->frame_restoration_type = RESTORE_BILATERAL;
+      rsi->bilateral_level = (int *)aom_realloc(
+          rsi->bilateral_level,
+          sizeof(*rsi->bilateral_level) * ntiles * BILATERAL_SUBTILES);
+      assert(rsi->bilateral_level != NULL);
+      for (i = 0; i < ntiles; ++i) {
+        int s;
+        rsi->restoration_type[i] = RESTORE_BILATERAL;
+        for (s = 0; s < BILATERAL_SUBTILES; ++s) {
+          const int j = i * BILATERAL_SUBTILES + s;
+          if (aom_read(rb, RESTORE_NONE_BILATERAL_PROB)) {
+            rsi->bilateral_level[j] =
+                aom_read_literal(rb, av1_bilateral_level_bits(cm));
+          } else {
+            rsi->bilateral_level[j] = -1;
+          }
         }
       }
     }
   } else {
-    rsi->restoration_type = RESTORE_NONE;
+    rsi->frame_restoration_type = RESTORE_NONE;
   }
 }
 #endif  // CONFIG_LOOP_RESTORATION
@@ -3286,7 +3358,7 @@ static size_t read_uncompressed_header(AV1Decoder *pbi,
   setup_dering(cm, rb);
 #endif
 #if CONFIG_LOOP_RESTORATION
-  setup_restoration(cm, rb);
+  decode_restoration_mode(cm, rb);
 #endif  // CONFIG_LOOP_RESTORATION
   setup_quantization(cm, rb);
 #if CONFIG_AOM_HIGHBITDEPTH
@@ -3468,6 +3540,10 @@ static int read_compressed_header(AV1Decoder *pbi, const uint8_t *data,
                        "Failed to allocate compressed header ANS decoder");
 #endif  // !CONFIG_ANS
 
+#if CONFIG_LOOP_RESTORATION
+  decode_restoration(cm, &r);
+#endif
+
   if (cm->tx_mode == TX_MODE_SELECT) {
     for (i = 0; i < TX_SIZES - 1; ++i)
       for (j = 0; j < TX_SIZE_CONTEXTS; ++j)

av1/encoder/bitstream.c

@@ -150,6 +150,9 @@ static struct av1_token interintra_mode_encodings[INTERINTRA_MODES];
 #if CONFIG_OBMC || CONFIG_WARPED_MOTION
 static struct av1_token motvar_encodings[MOTION_VARIATIONS];
 #endif  // CONFIG_OBMC || CONFIG_WARPED_MOTION
+#if CONFIG_LOOP_RESTORATION
+static struct av1_token switchable_restore_encodings[RESTORE_SWITCHABLE_TYPES];
+#endif  // CONFIG_LOOP_RESTORATION
 
 void av1_encode_token_init(void) {
 #if CONFIG_EXT_TX
@@ -176,6 +179,10 @@ void av1_encode_token_init(void) {
   av1_tokens_from_tree(global_motion_types_encodings,
                        av1_global_motion_types_tree);
 #endif  // CONFIG_GLOBAL_MOTION
+#if CONFIG_LOOP_RESTORATION
+  av1_tokens_from_tree(switchable_restore_encodings,
+                       av1_switchable_restore_tree);
+#endif  // CONFIG_LOOP_RESTORATION
 }
 
 static void write_intra_mode(aom_writer *w, PREDICTION_MODE mode,
@@ -2422,42 +2429,102 @@ static void update_coef_probs(AV1_COMP *cpi, aom_writer *w) {
 }
 
 #if CONFIG_LOOP_RESTORATION
-static void encode_restoration(AV1_COMMON *cm,
-                               struct aom_write_bit_buffer *wb) {
-  int i;
+static void encode_restoration_mode(AV1_COMMON *cm,
+                                    struct aom_write_bit_buffer *wb) {
   RestorationInfo *rst = &cm->rst_info;
-  aom_wb_write_bit(wb, rst->restoration_type != RESTORE_NONE);
-  if (rst->restoration_type != RESTORE_NONE) {
-    if (rst->restoration_type == RESTORE_BILATERAL) {
+  switch (rst->frame_restoration_type) {
+    case RESTORE_NONE:
+      aom_wb_write_bit(wb, 0);
+      aom_wb_write_bit(wb, 0);
+      break;
+    case RESTORE_SWITCHABLE:
+      aom_wb_write_bit(wb, 0);
+      aom_wb_write_bit(wb, 1);
+      break;
+    case RESTORE_BILATERAL:
       aom_wb_write_bit(wb, 1);
-      for (i = 0; i < cm->rst_internal.ntiles; ++i) {
-        if (rst->bilateral_level[i] >= 0) {
-          aom_wb_write_bit(wb, 1);
-          aom_wb_write_literal(wb, rst->bilateral_level[i],
-                               av1_bilateral_level_bits(cm));
-        } else {
-          aom_wb_write_bit(wb, 0);
-        }
-      }
-    } else {
       aom_wb_write_bit(wb, 0);
+      break;
+    case RESTORE_WIENER:
+      aom_wb_write_bit(wb, 1);
+      aom_wb_write_bit(wb, 1);
+      break;
+    default: assert(0);
+  }
+}
+
+static void encode_restoration(AV1_COMMON *cm, aom_writer *wb) {
+  int i;
+  RestorationInfo *rst = &cm->rst_info;
+  if (rst->frame_restoration_type != RESTORE_NONE) {
+    if (rst->frame_restoration_type == RESTORE_SWITCHABLE) {
+      // RESTORE_SWITCHABLE
       for (i = 0; i < cm->rst_internal.ntiles; ++i) {
-        if (rst->wiener_level[i]) {
-          aom_wb_write_bit(wb, 1);
-          aom_wb_write_literal(wb, rst->vfilter[i][0] - WIENER_FILT_TAP0_MINV,
-                               WIENER_FILT_TAP0_BITS);
-          aom_wb_write_literal(wb, rst->vfilter[i][1] - WIENER_FILT_TAP1_MINV,
-                               WIENER_FILT_TAP1_BITS);
-          aom_wb_write_literal(wb, rst->vfilter[i][2] - WIENER_FILT_TAP2_MINV,
+        av1_write_token(
+            wb, av1_switchable_restore_tree,
+            cm->fc->switchable_restore_prob,
+            &switchable_restore_encodings[rst->restoration_type[i]]);
+        if (rst->restoration_type[i] == RESTORE_NONE) {
+        } else if (rst->restoration_type[i] == RESTORE_BILATERAL) {
+          int s;
+          for (s = 0; s < BILATERAL_SUBTILES; ++s) {
+            const int j = i * BILATERAL_SUBTILES + s;
+#if BILATERAL_SUBTILES == 0
+            aom_write_literal(wb, rst->bilateral_level[j],
+                              av1_bilateral_level_bits(cm));
+#else
+            aom_write(wb, rst->bilateral_level[j] >= 0,
+                      RESTORE_NONE_BILATERAL_PROB);
+            if (rst->bilateral_level[j] >= 0) {
+              aom_write_literal(wb, rst->bilateral_level[j],
+                                av1_bilateral_level_bits(cm));
+            }
+#endif
+          }
+        } else {
+          aom_write_literal(wb, rst->vfilter[i][0] - WIENER_FILT_TAP0_MINV,
+                            WIENER_FILT_TAP0_BITS);
+          aom_write_literal(wb, rst->vfilter[i][1] - WIENER_FILT_TAP1_MINV,
+                            WIENER_FILT_TAP1_BITS);
+          aom_write_literal(wb, rst->vfilter[i][2] - WIENER_FILT_TAP2_MINV,
                                WIENER_FILT_TAP2_BITS);
-          aom_wb_write_literal(wb, rst->hfilter[i][0] - WIENER_FILT_TAP0_MINV,
+          aom_write_literal(wb, rst->hfilter[i][0] - WIENER_FILT_TAP0_MINV,
                                WIENER_FILT_TAP0_BITS);
-          aom_wb_write_literal(wb, rst->hfilter[i][1] - WIENER_FILT_TAP1_MINV,
+          aom_write_literal(wb, rst->hfilter[i][1] - WIENER_FILT_TAP1_MINV,
                                WIENER_FILT_TAP1_BITS);
-          aom_wb_write_literal(wb, rst->hfilter[i][2] - WIENER_FILT_TAP2_MINV,
+          aom_write_literal(wb, rst->hfilter[i][2] - WIENER_FILT_TAP2_MINV,
                                WIENER_FILT_TAP2_BITS);
-        } else {
-          aom_wb_write_bit(wb, 0);
+        }
+      }
+    } else if (rst->frame_restoration_type == RESTORE_BILATERAL) {
+      for (i = 0; i < cm->rst_internal.ntiles; ++i) {
+        int s;
+        for (s = 0; s < BILATERAL_SUBTILES; ++s) {
+          const int j = i * BILATERAL_SUBTILES + s;
+          aom_write(wb, rst->bilateral_level[j] >= 0,
+                    RESTORE_NONE_BILATERAL_PROB);
+          if (rst->bilateral_level[j] >= 0) {