Extend color restoration to support guided filters

Frame level guided filters are now used as an option for
chroma components in addition to the Wiener filter.

Change-Id: Ie33299d22c15b69741ede55686177b7b8ce8e2b3

av1/decoder/decodeframe.c

@@ -2460,8 +2460,13 @@ static void decode_restoration_mode(AV1_COMMON *cm,
         aom_rb_read_bit(rb) ? RESTORE_SWITCHABLE : RESTORE_NONE;
   }
   for (p = 1; p < MAX_MB_PLANE; ++p) {
-    cm->rst_info[p].frame_restoration_type =
-        aom_rb_read_bit(rb) ? RESTORE_WIENER : RESTORE_NONE;
+    rsi = &cm->rst_info[p];
+    if (aom_rb_read_bit(rb)) {
+      rsi->frame_restoration_type =
+          aom_rb_read_bit(rb) ? RESTORE_SGRPROJ : RESTORE_WIENER;
+    } else {
+      rsi->frame_restoration_type = RESTORE_NONE;
+    }
   }
 
   cm->rst_info[0].restoration_tilesize = RESTORATION_TILESIZE_MAX;
@@ -2594,6 +2599,7 @@ static void decode_restoration(AV1_COMMON *cm, aom_reader *rb) {
   }
   for (p = 1; p < MAX_MB_PLANE; ++p) {
     set_default_wiener(&ref_wiener_info);
+    set_default_sgrproj(&ref_sgrproj_info);
     rsi = &cm->rst_info[p];
     if (rsi->frame_restoration_type == RESTORE_WIENER) {
       for (i = 0; i < ntiles_uv; ++i) {
@@ -2607,6 +2613,21 @@ static void decode_restoration(AV1_COMMON *cm, aom_reader *rb) {
           read_wiener_filter(&rsi->wiener_info[i], &ref_wiener_info, rb);
         }
       }
+    } else if (rsi->frame_restoration_type == RESTORE_SGRPROJ) {
+      for (i = 0; i < ntiles_uv; ++i) {
+        if (ntiles_uv > 1)
+          rsi->restoration_type[i] =
+              aom_read(rb, RESTORE_NONE_SGRPROJ_PROB, ACCT_STR)
+                  ? RESTORE_SGRPROJ
+                  : RESTORE_NONE;
+        else
+          rsi->restoration_type[i] = RESTORE_SGRPROJ;
+        if (rsi->restoration_type[i] == RESTORE_SGRPROJ) {
+          read_sgrproj_filter(&rsi->sgrproj_info[i], &ref_sgrproj_info, rb);
+        }
+      }
+    } else if (rsi->frame_restoration_type != RESTORE_NONE) {
+      assert(0);
     }
   }
 }

av1/encoder/bitstream.c

@@ -3370,7 +3370,14 @@ static void encode_restoration_mode(AV1_COMMON *cm,
     rsi = &cm->rst_info[p];
     switch (rsi->frame_restoration_type) {
       case RESTORE_NONE: aom_wb_write_bit(wb, 0); break;
-      case RESTORE_WIENER: aom_wb_write_bit(wb, 1); break;
+      case RESTORE_WIENER:
+        aom_wb_write_bit(wb, 1);
+        aom_wb_write_bit(wb, 0);
+        break;
+      case RESTORE_SGRPROJ:
+        aom_wb_write_bit(wb, 1);
+        aom_wb_write_bit(wb, 1);
+        break;
       default: assert(0);
     }
   }
@@ -3483,6 +3490,7 @@ static void encode_restoration(AV1_COMMON *cm, aom_writer *wb) {
   }
   for (p = 1; p < MAX_MB_PLANE; ++p) {
     set_default_wiener(&ref_wiener_info);
+    set_default_sgrproj(&ref_sgrproj_info);
     rsi = &cm->rst_info[p];
     if (rsi->frame_restoration_type == RESTORE_WIENER) {
       for (i = 0; i < ntiles_uv; ++i) {
@@ -3493,6 +3501,15 @@ static void encode_restoration(AV1_COMMON *cm, aom_writer *wb) {
           write_wiener_filter(&rsi->wiener_info[i], &ref_wiener_info, wb);
         }
       }
+    } else if (rsi->frame_restoration_type == RESTORE_SGRPROJ) {
+      for (i = 0; i < ntiles_uv; ++i) {
+        if (ntiles_uv > 1)
+          aom_write(wb, rsi->restoration_type[i] != RESTORE_NONE,
+                    RESTORE_NONE_SGRPROJ_PROB);
+        if (rsi->restoration_type[i] != RESTORE_NONE) {
+          write_sgrproj_filter(&rsi->sgrproj_info[i], &ref_sgrproj_info, wb);
+        }
+      }
     } else if (rsi->frame_restoration_type != RESTORE_NONE) {
       assert(0);
     }

av1/encoder/pickrst.c

@@ -364,6 +364,121 @@ static int count_sgrproj_bits(SgrprojInfo *sgrproj_info,
   return bits;
 }
 
+static double search_sgrproj_uv(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
+                                int partial_frame, int plane,
+                                RestorationInfo *info, RestorationType *type,
+                                YV12_BUFFER_CONFIG *dst_frame) {
+  SgrprojInfo *sgrproj_info = info->sgrproj_info;
+  int64_t err;
+  double cost_norestore, cost_sgrproj, cost_sgrproj_frame;
+  int bits;
+  MACROBLOCK *x = &cpi->td.mb;
+  AV1_COMMON *const cm = &cpi->common;
+  const YV12_BUFFER_CONFIG *dgd = cm->frame_to_show;
+
+  const int width = src->uv_crop_width;
+  const int height = src->uv_crop_height;
+  const int src_stride = src->uv_stride;
+  const int dgd_stride = dgd->uv_stride;
+
+  RestorationInfo *rsi = cpi->rst_search;
+  int tile_idx, tile_width, tile_height, nhtiles, nvtiles;
+  int h_start, h_end, v_start, v_end;
+  // Allocate for the src buffer at high precision
+  const int ntiles = av1_get_rest_ntiles(
+      width, height, cm->rst_info[plane].restoration_tilesize, &tile_width,
+      &tile_height, &nhtiles, &nvtiles);
+  SgrprojInfo ref_sgrproj_info;
+  set_default_sgrproj(&ref_sgrproj_info);
+  assert(width == dgd->uv_crop_width);
+  assert(height == dgd->uv_crop_height);
+
+  rsi[plane].frame_restoration_type = RESTORE_SGRPROJ;
+
+  for (tile_idx = 0; tile_idx < ntiles; ++tile_idx) {
+    rsi[plane].restoration_type[tile_idx] = RESTORE_NONE;
+  }
+  // Compute best Sgrproj filters for each tile
+  for (tile_idx = 0; tile_idx < ntiles; ++tile_idx) {
+    av1_get_rest_tile_limits(tile_idx, 0, 0, nhtiles, nvtiles, tile_width,
+                             tile_height, width, height, 0, 0, &h_start, &h_end,
+                             &v_start, &v_end);
+    err = sse_restoration_tile(src, cm->frame_to_show, cm, h_start,
+                               h_end - h_start, v_start, v_end - v_start,
+                               (1 << plane));
+    // #bits when a tile is not restored
+    bits = av1_cost_bit(RESTORE_NONE_SGRPROJ_PROB, 0);
+    cost_norestore = RDCOST_DBL(x->rdmult, x->rddiv, (bits >> 4), err);
+    if (plane == AOM_PLANE_U) {
+      search_selfguided_restoration(
+          dgd->u_buffer + v_start * dgd_stride + h_start, h_end - h_start,
+          v_end - v_start, dgd_stride,
+          src->u_buffer + v_start * src_stride + h_start, src_stride,
+#if CONFIG_HIGHBITDEPTH
+          cm->bit_depth,
+#else
+          8,
+#endif  // CONFIG_HIGHBITDEPTH
+          &rsi[plane].sgrproj_info[tile_idx].ep,
+          rsi[plane].sgrproj_info[tile_idx].xqd, cm->rst_internal.tmpbuf);
+    } else if (plane == AOM_PLANE_V) {
+      search_selfguided_restoration(
+          dgd->v_buffer + v_start * dgd_stride + h_start, h_end - h_start,
+          v_end - v_start, dgd_stride,
+          src->v_buffer + v_start * src_stride + h_start, src_stride,
+#if CONFIG_HIGHBITDEPTH
+          cm->bit_depth,
+#else
+          8,
+#endif  // CONFIG_HIGHBITDEPTH
+          &rsi[plane].sgrproj_info[tile_idx].ep,
+          rsi[plane].sgrproj_info[tile_idx].xqd, cm->rst_internal.tmpbuf);
+    } else {
+      assert(0);
+    }
+    rsi[plane].restoration_type[tile_idx] = RESTORE_SGRPROJ;
+    err = try_restoration_tile(src, cpi, rsi, (1 << plane), partial_frame,
+                               tile_idx, 0, 0, dst_frame);
+    bits = count_sgrproj_bits(&rsi[plane].sgrproj_info[tile_idx],
+                              &ref_sgrproj_info)
+           << AV1_PROB_COST_SHIFT;
+    bits += av1_cost_bit(RESTORE_NONE_SGRPROJ_PROB, 1);
+    cost_sgrproj = RDCOST_DBL(x->rdmult, x->rddiv, (bits >> 4), err);
+    if (cost_sgrproj >= cost_norestore) {
+      type[tile_idx] = RESTORE_NONE;
+    } else {
+      type[tile_idx] = RESTORE_SGRPROJ;
+      memcpy(&sgrproj_info[tile_idx], &rsi[plane].sgrproj_info[tile_idx],
+             sizeof(sgrproj_info[tile_idx]));
+      memcpy(&ref_sgrproj_info, &sgrproj_info[tile_idx],
+             sizeof(ref_sgrproj_info));
+    }
+    rsi[plane].restoration_type[tile_idx] = RESTORE_NONE;
+  }
+  // Cost for Sgrproj filtering
+  set_default_sgrproj(&ref_sgrproj_info);
+  bits = frame_level_restore_bits[rsi[plane].frame_restoration_type]
+         << AV1_PROB_COST_SHIFT;
+  for (tile_idx = 0; tile_idx < ntiles; ++tile_idx) {
+    bits +=
+        av1_cost_bit(RESTORE_NONE_SGRPROJ_PROB, type[tile_idx] != RESTORE_NONE);
+    memcpy(&rsi[plane].sgrproj_info[tile_idx], &sgrproj_info[tile_idx],
+           sizeof(sgrproj_info[tile_idx]));
+    if (type[tile_idx] == RESTORE_SGRPROJ) {
+      bits += count_sgrproj_bits(&rsi[plane].sgrproj_info[tile_idx],
+                                 &ref_sgrproj_info)
+              << AV1_PROB_COST_SHIFT;
+      memcpy(&ref_sgrproj_info, &rsi[plane].sgrproj_info[tile_idx],
+             sizeof(ref_sgrproj_info));
+    }
+    rsi[plane].restoration_type[tile_idx] = type[tile_idx];
+  }
+  err = try_restoration_frame(src, cpi, rsi, (1 << plane), partial_frame,
+                              dst_frame);
+  cost_sgrproj_frame = RDCOST_DBL(x->rdmult, x->rddiv, (bits >> 4), err);
+  return cost_sgrproj_frame;
+}
+
 static double search_sgrproj(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
                              int partial_frame, RestorationInfo *info,
                              RestorationType *type, double *best_tile_cost,
@@ -895,7 +1010,7 @@ static double search_wiener_uv(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   RestorationInfo *rsi = cpi->rst_search;
   int64_t err;
   int bits;
-  double cost_wiener, cost_norestore, cost_wiener_frame, cost_norestore_frame;
+  double cost_wiener, cost_norestore, cost_wiener_frame;
   MACROBLOCK *x = &cpi->td.mb;
   double M[WIENER_WIN2];
   double H[WIENER_WIN2 * WIENER_WIN2];
@@ -908,19 +1023,14 @@ static double search_wiener_uv(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   double score;
   int tile_idx, tile_width, tile_height, nhtiles, nvtiles;
   int h_start, h_end, v_start, v_end;
-  const int ntiles =
-      av1_get_rest_ntiles(width, height, cm->rst_info[1].restoration_tilesize,
-                          &tile_width, &tile_height, &nhtiles, &nvtiles);
+  const int ntiles = av1_get_rest_ntiles(
+      width, height, cm->rst_info[plane].restoration_tilesize, &tile_width,
+      &tile_height, &nhtiles, &nvtiles);
   WienerInfo ref_wiener_info;
   set_default_wiener(&ref_wiener_info);
   assert(width == dgd->uv_crop_width);
   assert(height == dgd->uv_crop_height);
 
-  rsi[plane].frame_restoration_type = RESTORE_NONE;
-  err = sse_restoration_frame(cm, src, cm->frame_to_show, (1 << plane));
-  bits = 0;
-  cost_norestore_frame = RDCOST_DBL(x->rdmult, x->rddiv, (bits >> 4), err);
-
   rsi[plane].frame_restoration_type = RESTORE_WIENER;
 
   for (tile_idx = 0; tile_idx < ntiles; ++tile_idx) {
@@ -1036,15 +1146,7 @@ static double search_wiener_uv(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   err = try_restoration_frame(src, cpi, rsi, 1 << plane, partial_frame,
                               dst_frame);
   cost_wiener_frame = RDCOST_DBL(x->rdmult, x->rddiv, (bits >> 4), err);
-
-  if (cost_wiener_frame < cost_norestore_frame) {
-    info->frame_restoration_type = RESTORE_WIENER;
-  } else {
-    info->frame_restoration_type = RESTORE_NONE;
-  }
-
-  return info->frame_restoration_type == RESTORE_WIENER ? cost_wiener_frame
-                                                        : cost_norestore_frame;
+  return cost_wiener_frame;
 }
 
 static double search_wiener(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
@@ -1183,11 +1285,41 @@ static double search_wiener(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   return cost_wiener;
 }
 
+static double search_norestore_uv(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
+                                  int partial_frame, int plane,
+                                  RestorationInfo *info, RestorationType *type,
+                                  YV12_BUFFER_CONFIG *dst_frame) {
+  double cost_norestore;
+  int64_t err;
+  int bits;
+  MACROBLOCK *x = &cpi->td.mb;
+  AV1_COMMON *const cm = &cpi->common;
+  const int width = src->uv_crop_width;
+  const int height = src->uv_crop_height;
+  int tile_idx, tile_width, tile_height, nhtiles, nvtiles;
+  const int ntiles = av1_get_rest_ntiles(
+      width, height, cm->rst_info[plane].restoration_tilesize, &tile_width,
+      &tile_height, &nhtiles, &nvtiles);
+  (void)dst_frame;
+  (void)partial_frame;
+
+  info->frame_restoration_type = RESTORE_NONE;
+  for (tile_idx = 0; tile_idx < ntiles; ++tile_idx) {
+    type[tile_idx] = RESTORE_NONE;
+  }
+  // RD cost associated with no restoration
+  err = sse_restoration_frame(cm, src, cm->frame_to_show, (1 << plane));
+  bits = frame_level_restore_bits[RESTORE_NONE] << AV1_PROB_COST_SHIFT;
+  cost_norestore = RDCOST_DBL(x->rdmult, x->rddiv, (bits >> 4), err);
+  return cost_norestore;
+}
+
 static double search_norestore(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
                                int partial_frame, RestorationInfo *info,
                                RestorationType *type, double *best_tile_cost,
                                YV12_BUFFER_CONFIG *dst_frame) {
-  double err, cost_norestore;
+  int64_t err;
+  double cost_norestore;
   int bits;
   MACROBLOCK *x = &cpi->td.mb;
   AV1_COMMON *const cm = &cpi->common;
@@ -1328,14 +1460,39 @@ void av1_pick_filter_restoration(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   }
 
   // Color components
-  search_wiener_uv(src, cpi, method == LPF_PICK_FROM_SUBIMAGE, AOM_PLANE_U,
-                   &cm->rst_info[AOM_PLANE_U],
-                   cm->rst_info[AOM_PLANE_U].restoration_type,
-                   &cpi->trial_frame_rst);
-  search_wiener_uv(src, cpi, method == LPF_PICK_FROM_SUBIMAGE, AOM_PLANE_V,
-                   &cm->rst_info[AOM_PLANE_V],
-                   cm->rst_info[AOM_PLANE_V].restoration_type,
-                   &cpi->trial_frame_rst);
+  const int ntiles_uv = av1_get_rest_ntiles(
+      ROUND_POWER_OF_TWO(cm->width, cm->subsampling_x),
+      ROUND_POWER_OF_TWO(cm->height, cm->subsampling_y),
+      cm->rst_info[1].restoration_tilesize, NULL, NULL, NULL, NULL);
+  for (int plane = AOM_PLANE_U; plane <= AOM_PLANE_V; ++plane) {
+    double cost_uv[RESTORE_SWITCHABLE_TYPES];
+    cost_uv[RESTORE_NONE] = search_norestore_uv(
+        src, cpi, method == LPF_PICK_FROM_SUBIMAGE, plane, &cm->rst_info[plane],
+        restore_types[0], &cpi->trial_frame_rst);
+    cost_uv[RESTORE_WIENER] = search_wiener_uv(
+        src, cpi, method == LPF_PICK_FROM_SUBIMAGE, plane, &cm->rst_info[plane],
+        restore_types[RESTORE_WIENER], &cpi->trial_frame_rst);
+    cost_uv[RESTORE_SGRPROJ] = search_sgrproj_uv(
+        src, cpi, method == LPF_PICK_FROM_SUBIMAGE, plane, &cm->rst_info[plane],
+        restore_types[RESTORE_SGRPROJ], &cpi->trial_frame_rst);
+    if (cost_uv[RESTORE_SGRPROJ] < cost_uv[RESTORE_WIENER] &&
+        cost_uv[RESTORE_SGRPROJ] < cost_uv[RESTORE_NONE]) {
+      cm->rst_info[plane].frame_restoration_type = RESTORE_SGRPROJ;
+      memcpy(cm->rst_info[plane].restoration_type,
+             restore_types[RESTORE_SGRPROJ],
+             ntiles_uv * sizeof(restore_types[RESTORE_SGRPROJ][0]));
+    } else if (cost_uv[RESTORE_WIENER] < cost_uv[RESTORE_NONE] &&
+               cost_uv[RESTORE_WIENER] < cost_uv[RESTORE_SGRPROJ]) {
+      cm->rst_info[plane].frame_restoration_type = RESTORE_WIENER;
+      memcpy(cm->rst_info[plane].restoration_type,
+             restore_types[RESTORE_WIENER],
+             ntiles_uv * sizeof(restore_types[RESTORE_WIENER][0]));
+    } else {
+      cm->rst_info[plane].frame_restoration_type = RESTORE_NONE;
+      memcpy(cm->rst_info[plane].restoration_type, restore_types[RESTORE_NONE],
+             ntiles_uv * sizeof(restore_types[RESTORE_NONE][0]));
+    }
+  }
   /*
   printf("Frame %d/%d restore types: %d %d %d\n", cm->current_video_frame,
          cm->show_frame, cm->rst_info[0].frame_restoration_type,