From e559967790c3f4e13b91d19050eeff58967eb218 Mon Sep 17 00:00:00 2001 From: Manisha Jajoo Date: Thu, 2 Jun 2022 23:25:53 +0530 Subject: [PATCH] Encoder: add support for PSNR export This CL adds support to report PSNR between incoming and resulting encoded frames. Eventually be reported upstream as part of a feedback mechanism. Test: avcenc -c enc.cfg Bug: b/235899898 Change-Id: Id520913ca5964965b5a86bc520240ccd1b28176c --- Android.bp | 1 + encoder/ih264e_api.c | 3 ++ encoder/ih264e_defs.h | 3 ++ encoder/ih264e_encode.c | 18 ++++++++++ encoder/ih264e_master.h | 17 +++++++++ encoder/ih264e_process.c | 5 ++- encoder/ih264e_structs.h | 29 +++++++++++++++- encoder/ih264e_utils.c | 75 ++++++++++++++++++++++++++++++++++++++++ encoder/libavcenc.cmake | 3 +- encoder/psnr.c | 65 ++++++++++++++++++++++++++++++++++ encoder/psnr.h | 30 ++++++++++++++++ 11 files changed, 246 insertions(+), 3 deletions(-) create mode 100644 encoder/psnr.c create mode 100644 encoder/psnr.h diff --git a/Android.bp b/Android.bp index fdd3ec2..992e8ed 100644 --- a/Android.bp +++ b/Android.bp @@ -415,6 +415,7 @@ cc_library_static { "encoder/ime.c", "encoder/ime_distortion_metrics.c", "encoder/ih264e_sei.c", + "encoder/psnr.c", ], arch: { diff --git a/encoder/ih264e_api.c b/encoder/ih264e_api.c index 1f4c532..0660ba6 100644 --- a/encoder/ih264e_api.c +++ b/encoder/ih264e_api.c @@ -2873,6 +2873,7 @@ static WORD32 ih264e_set_default_params(cfg_params_t *ps_cfg) ps_cfg->u4_max_level = DEFAULT_MAX_LEVEL; ps_cfg->e_inp_color_fmt = IV_YUV_420SP_UV; ps_cfg->u4_enable_recon = DEFAULT_RECON_ENABLE; + ps_cfg->u4_enable_quality_metrics = DEFAULT_QUALITY_METRICS_ENABLE; ps_cfg->e_recon_color_fmt = IV_YUV_420P; ps_cfg->u4_enc_speed_preset = IVE_FASTEST; ps_cfg->e_rc_mode = DEFAULT_RC; @@ -3045,6 +3046,8 @@ static WORD32 ih264e_init(codec_t *ps_codec) /* Process thread created status */ memset(ps_codec->ai4_process_thread_created, 0, MAX_PROCESS_THREADS); + memset(&ps_codec->s_global_quality_stats, 0, sizeof(ps_codec->s_global_quality_stats)); + /* Number of MBs processed together */ ps_codec->i4_proc_nmb = 8; diff --git a/encoder/ih264e_defs.h b/encoder/ih264e_defs.h index 33fafd1..7f69554 100644 --- a/encoder/ih264e_defs.h +++ b/encoder/ih264e_defs.h @@ -205,6 +205,7 @@ /* Generic declarations */ #define DEFAULT_MAX_LEVEL 40 #define DEFAULT_RECON_ENABLE 0 +#define DEFAULT_QUALITY_METRICS_ENABLE 0 #define DEFAULT_RC IVE_RC_STORAGE #define DEFAULT_MAX_FRAMERATE 120000 #define DEFAULT_MAX_BITRATE 240000000 @@ -267,6 +268,8 @@ /** Number of buffers Needed for SUBPEL and BIPRED computation */ #define SUBPEL_BUFF_CNT 4 +/** Mask value for PSNR. Needed when quality metrics is enabled */ +#define QUALITY_MASK_PSNR 0x1 /** ***************************************************************************** * Macro to compute total size required to hold on set of scaling matrices diff --git a/encoder/ih264e_encode.c b/encoder/ih264e_encode.c index 6a4e3a2..57cc246 100644 --- a/encoder/ih264e_encode.c +++ b/encoder/ih264e_encode.c @@ -531,6 +531,12 @@ WORD32 ih264e_encode(iv_obj_t *ps_codec_obj, void *pv_api_ip, void *pv_api_op) ih264_list_reset(ps_codec->pv_proc_jobq); ih264_list_reset(ps_codec->pv_entropy_jobq); + + if (ps_codec->s_cfg.u4_enable_quality_metrics & QUALITY_MASK_PSNR) + { + ih264e_compute_quality_stats(ps_proc); + } + } @@ -566,6 +572,7 @@ WORD32 ih264e_encode(iv_obj_t *ps_codec_obj, void *pv_api_ip, void *pv_api_op) IH264_ERROR_T ret = IH264_SUCCESS; pic_buf_t *ps_pic_buf = NULL; WORD32 i4_buf_status, i4_curr_poc = 32768; + WORD8 buf_idx = -1; /* In case of skips we return recon, but indicate that buffer is zero size */ if (ps_codec->s_rate_control.post_encode_skip[ctxt_sel] @@ -593,6 +600,17 @@ WORD32 ih264e_encode(iv_obj_t *ps_codec_obj, void *pv_api_ip, void *pv_api_op) { ps_pic_buf = ps_codec->as_ref_set[i].ps_pic_buf; i4_curr_poc = ps_codec->as_ref_set[i].i4_poc; + buf_idx = i; + } + } + if ((ps_codec->s_cfg.u4_enable_quality_metrics & QUALITY_MASK_PSNR) + && buf_idx >= 0) + { + UWORD8 comp; + for(comp = 0; comp < 3; comp++) + { + DEBUG("PSNR[%d]: %f\n", comp, + ps_codec->as_ref_set[buf_idx].s_pic_quality_stats.total_psnr[comp]); } } diff --git a/encoder/ih264e_master.h b/encoder/ih264e_master.h index 6c7505a..67354fd 100644 --- a/encoder/ih264e_master.h +++ b/encoder/ih264e_master.h @@ -60,6 +60,23 @@ */ void ih264e_join_threads(codec_t *ps_codec); +/** +****************************************************************************** +* +* @brief +* This function calculates various quality metrics; the initial one is PSNR. +* +* @par Description +* +* @param[in] ps_codec +* pointer to process context +* +* @returns none +* +****************************************************************************** +*/ +void ih264e_compute_quality_stats(process_ctxt_t *ps_proc); + /** ****************************************************************************** * diff --git a/encoder/ih264e_process.c b/encoder/ih264e_process.c index bf4c1e7..0381dec 100644 --- a/encoder/ih264e_process.c +++ b/encoder/ih264e_process.c @@ -1260,6 +1260,7 @@ IH264E_ERROR_T ih264e_init_proc_ctxt(process_ctxt_t *ps_proc) convert_uv_only = 1; if (u4_pad_bottom_sz || u4_pad_right_sz || ps_codec->s_cfg.e_inp_color_fmt == IV_YUV_422ILE || + ps_codec->s_cfg.u4_enable_quality_metrics & QUALITY_MASK_PSNR || ps_proc->i4_mb_y == (ps_proc->i4_ht_mbs - 1)) { if (ps_proc->i4_mb_y == ps_proc->i4_ht_mbs - 1) @@ -1279,6 +1280,7 @@ IH264E_ERROR_T ih264e_init_proc_ctxt(process_ctxt_t *ps_proc) if (ps_codec->s_cfg.e_inp_color_fmt == IV_YUV_422ILE || ps_codec->s_cfg.e_inp_color_fmt == IV_YUV_420P || ps_proc->i4_mb_y == (ps_proc->i4_ht_mbs - 1) || + ps_codec->s_cfg.u4_enable_quality_metrics & QUALITY_MASK_PSNR || u4_pad_bottom_sz || u4_pad_right_sz) { if ((ps_codec->s_cfg.e_inp_color_fmt == IV_YUV_420SP_UV) || @@ -2070,7 +2072,8 @@ WORD32 ih264e_process(process_ctxt_t *ps_proc) * 2. dump recon for bit stream sanity check */ ps_proc->u4_compute_recon = ps_codec->u4_is_curr_frm_ref || - ps_codec->s_cfg.u4_enable_recon; + ps_codec->s_cfg.u4_enable_recon || + ps_codec->s_cfg.u4_enable_quality_metrics & QUALITY_MASK_PSNR; /* Encode 'n' macroblocks, * 'n' being the number of mbs dictated by current proc ctxt */ diff --git a/encoder/ih264e_structs.h b/encoder/ih264e_structs.h index 8674470..9b22642 100644 --- a/encoder/ih264e_structs.h +++ b/encoder/ih264e_structs.h @@ -369,9 +369,12 @@ typedef struct /** Input color format */ IV_COLOR_FORMAT_T e_inp_color_fmt; - /** Flag to enable/disable - To be used only for debugging/testing */ + /** Flag to enable/disable recon */ UWORD32 u4_enable_recon; + /** Flag to enable/disable quality metrics */ + UWORD32 u4_enable_quality_metrics; + /** Recon color format */ IV_COLOR_FORMAT_T e_recon_color_fmt; @@ -626,6 +629,22 @@ typedef struct } job_t; +/** + ***************************************************************************** + * @brief Structure to store psnr of the sequence + ***************************************************************************** + */ +typedef struct { + UWORD64 total_sse[3]; + double global_psnr[3]; // total_sse / total_samples + double total_psnr[3]; // sum (per_frame_sse / per_frame_samples) + double avg_psnr[3]; // total_psnr / total_frames + UWORD32 total_samples[3]; + WORD32 total_frames; +} quality_stats_t; + + + /** * Structure to represent a MV Bank buffer */ @@ -669,6 +688,9 @@ typedef struct */ typedef struct { + /** Per pic PSNR */ + quality_stats_t s_pic_quality_stats; + /** Picture count */ WORD32 i4_pic_cnt; @@ -2473,6 +2495,11 @@ struct _codec_t */ void *apv_proc_thread_handle[MAX_PROCESS_THREADS]; + /** + * Structure for global PSNR + */ + quality_stats_t s_global_quality_stats; + /** * Thread created flag for each of the processing threads */ diff --git a/encoder/ih264e_utils.c b/encoder/ih264e_utils.c index 239271e..94050cd 100644 --- a/encoder/ih264e_utils.c +++ b/encoder/ih264e_utils.c @@ -70,6 +70,7 @@ #include "ime_distortion_metrics.h" #include "ime_defs.h" #include "ime_structs.h" +#include "psnr.h" #include "ih264_error.h" #include "ih264_structs.h" #include "ih264_trans_quant_itrans_iquant.h" @@ -2204,3 +2205,77 @@ IH264E_ERROR_T ih264e_pic_init(codec_t *ps_codec, inp_buf_t *ps_inp_buf) return error_status; } + +/** +******************************************************************************* +* +* @brief +* Calculate the per-pic and global PSNR +* +* @par Description: +* This function takes the source and recon luma/chroma buffer pointers from the +* codec context and calculates the per-pic and global PSNR for the current encoding +* frame. +* +* @param[in] ps_codec +* Pointer to process context +* +* @returns none +* +* @remarks +* +* +******************************************************************************* +*/ +void ih264e_compute_quality_stats(process_ctxt_t *ps_proc) +{ + codec_t *ps_codec = ps_proc->ps_codec; + WORD32 wd = ps_codec->s_cfg.u4_wd; + WORD32 ht = ps_codec->s_cfg.u4_ht; + WORD32 disp_wd = ps_codec->s_cfg.u4_disp_wd; + WORD32 disp_ht = ps_codec->s_cfg.u4_disp_ht; + WORD32 src_strds = ps_proc->i4_src_strd; + WORD32 rec_strds = ps_proc->i4_rec_strd; + quality_stats_t *ps_pic_quality_stats = NULL; + double sum_squared_error[3] = {0.0, 0.0, 0.0}; + double total_samples[3]; + WORD32 i; + for (i = 0; i < ps_codec->i4_ref_buf_cnt; i++) + { + if (ps_codec->as_ref_set[i].i4_pic_cnt != -1 && + ps_codec->as_ref_set[i].i4_poc == ps_codec->i4_poc) + { + ps_pic_quality_stats = &ps_codec->as_ref_set[i].s_pic_quality_stats; + break; + } + } + + if(ps_pic_quality_stats == NULL) return; + + get_sse( + ps_proc->pu1_src_buf_luma_base, ps_proc->pu1_rec_buf_luma_base, + ps_proc->pu1_src_buf_chroma_base, ps_proc->pu1_rec_buf_chroma_base, + src_strds, rec_strds, wd, ht, sum_squared_error); + + total_samples[0] = disp_wd * disp_ht; + total_samples[1] = total_samples[2] = total_samples[0] / 4; + + ps_pic_quality_stats->total_frames = 1; + ps_codec->s_global_quality_stats.total_frames += 1; + for (i = 0; i < 3; i++) + { + double psnr = sse_to_psnr(total_samples[i], sum_squared_error[i]); + ps_pic_quality_stats->total_samples[i] = total_samples[i]; + ps_pic_quality_stats->total_sse[i] = sum_squared_error[i]; + ps_pic_quality_stats->global_psnr[i] = ps_pic_quality_stats->avg_psnr[i] = + ps_pic_quality_stats->total_psnr[i] = psnr; + ps_codec->s_global_quality_stats.total_sse[i] += sum_squared_error[i]; + ps_codec->s_global_quality_stats.global_psnr[i] = + sse_to_psnr(ps_codec->s_global_quality_stats.total_samples[i], + ps_codec->s_global_quality_stats.total_sse[i]); + ps_codec->s_global_quality_stats.total_psnr[i] += psnr; + ps_codec->s_global_quality_stats.avg_psnr[i] = + ps_codec->s_global_quality_stats.total_psnr[i] / + ps_codec->s_global_quality_stats.total_frames; + } +} diff --git a/encoder/libavcenc.cmake b/encoder/libavcenc.cmake index b948a6e..cf7be9c 100644 --- a/encoder/libavcenc.cmake +++ b/encoder/libavcenc.cmake @@ -39,7 +39,8 @@ list( "${AVC_ROOT}/encoder/irc_rate_control_api.c" "${AVC_ROOT}/encoder/irc_rd_model.c" "${AVC_ROOT}/encoder/irc_vbr_storage_vbv.c" - "${AVC_ROOT}/encoder/irc_vbr_str_prms.c") + "${AVC_ROOT}/encoder/irc_vbr_str_prms.c" + "${AVC_ROOT}/encoder/psnr.c") include_directories(${AVC_ROOT}/encoder) diff --git a/encoder/psnr.c b/encoder/psnr.c new file mode 100644 index 0000000..5f53a1f --- /dev/null +++ b/encoder/psnr.c @@ -0,0 +1,65 @@ +/****************************************************************************** + * + * Copyright (C) 2022 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at: + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + *****************************************************************************/ + +/* System Include Files */ +#include + +#include "ih264_macros.h" +#include "ih264_typedefs.h" +#include "psnr.h" + +void get_sse(UWORD8 *pu1_src_luma, UWORD8 *pu1_rec_luma, UWORD8 *pu1_src_chroma, + UWORD8 *pu1_rec_chroma, WORD32 src_strd, WORD32 rec_strd, WORD32 width, WORD32 height, + DOUBLE pd_sse[3]) +{ + WORD32 i, j; + + for(j = 0; j < height; j++) + { + for(i = 0; i < width; i++) + { + WORD32 diff = pu1_src_luma[i] - pu1_rec_luma[i]; + pd_sse[0] += diff * diff; + } + pu1_src_luma += src_strd; + pu1_rec_luma += rec_strd; + } + + for(j = 0; j < height / 2; j++) + { + for(i = 0; i < width / 2; i++) + { + WORD32 diff = pu1_src_chroma[i * 2] - pu1_rec_chroma[i * 2]; + pd_sse[1] += diff * diff; + diff = pu1_src_chroma[i * 2 + 1] - pu1_rec_chroma[i * 2 + 1]; + pd_sse[2] += diff * diff; + } + pu1_src_chroma += src_strd; + pu1_rec_chroma += rec_strd; + } +} + +DOUBLE sse_to_psnr(DOUBLE samples, DOUBLE sse) +{ + DOUBLE psnr; + if(samples <= 0) return -1; + if (sse<=0) return MAX_PSNR; + psnr = 10.0 * (log10(samples) + 2*log10(255) - log10(sse)); + psnr = MIN(MAX_PSNR, psnr); + return psnr; +} diff --git a/encoder/psnr.h b/encoder/psnr.h new file mode 100644 index 0000000..546dd0d --- /dev/null +++ b/encoder/psnr.h @@ -0,0 +1,30 @@ +/****************************************************************************** + * + * Copyright (C) 2022 The Android Open Source Project + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at: + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + *****************************************************************************/ + +#ifndef __PSNR_H__ +#define __PSNR_H__ + +#define MAX_PSNR 100.0 + +void get_sse(UWORD8 *pu1_src_luma, UWORD8 *pu1_rec_luma, UWORD8 *pu1_src_chroma, + UWORD8 *pu1_rec_chroma, WORD32 src_strd, WORD32 rec_strd, WORD32 width, WORD32 height, + DOUBLE *pd_sse); + +DOUBLE sse_to_psnr(DOUBLE samples, DOUBLE sse); + +#endif