Commit d23ceead authored by Debargha Mukherjee's avatar Debargha Mukherjee
Browse files

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
parent 6171ae77
......@@ -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);
}
}
}
......
......@@ -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);
}
......
......@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment