img_dist_ms_ssim.c
上传用户:hjq518
上传日期:2021-12-09
资源大小:5084k
文件大小:11k
源码类别:

Audio

开发平台:

Visual C++

  1. /*!
  2.  *************************************************************************************
  3.  * file img_dist_ms_ssim.c
  4.  *
  5.  * brief
  6.  *    Compute structural similarity (SSIM) index using the encoded image and the reference image
  7.  *
  8.  * author
  9.  *    Main contributors (see contributors.h for copyright, address and affiliation details)
  10.  *     - Woo-Shik Kim                    <wooshik.kim@usc.edu>
  11.  *     - Peshala Pahalawatta             <ppaha@dolby.com>
  12.  *     - Zhen Li                         <zli@dolby.com> 
  13.  *     - Alexis Michael Tourapis         <alexismt@ieee.org>
  14.  *************************************************************************************
  15.  */
  16. #include "contributors.h"
  17. #include "global.h"
  18. #include "img_distortion.h"
  19. #include "enc_statistics.h"
  20. #include "memalloc.h"
  21. #include "math.h"
  22. #ifndef UNBIASED_VARIANCE
  23.   #define UNBIASED_VARIANCE // unbiased estimation of the variance
  24. #endif
  25. #define MS_SSIM_PAD  6
  26. #define MS_SSIM_PAD2 3
  27. #define MS_SSIM_BETA0 0.0448
  28. #define MS_SSIM_BETA1 0.2856
  29. #define MS_SSIM_BETA2 0.3001
  30. #define MS_SSIM_BETA3 0.2363
  31. #define MS_SSIM_BETA4 0.1333
  32. #define MAX_SSIM_LEVELS 5
  33. //Computes the product of the contrast and structure componenents of the structural similarity metric.
  34. float compute_structural_components (imgpel **refImg, imgpel **encImg, int height, int width, int win_height, int win_width, int comp)
  35. {
  36.   static const float K2 = 0.03f;
  37.   static float max_pix_value_sqd;
  38.   float C2;
  39.   float win_pixels = (float) (win_width * win_height);
  40. #ifdef UNBIASED_VARIANCE
  41.   float win_pixels_bias = win_pixels - 1;
  42. #else
  43.   float win_pixels_bias = win_pixels;
  44. #endif
  45.   float mb_ssim, meanOrg, meanEnc;
  46.   float varOrg, varEnc, covOrgEnc;
  47.   int imeanOrg, imeanEnc, ivarOrg, ivarEnc, icovOrgEnc;
  48.   float cur_distortion = 0.0;
  49.   int i, j, n, m, win_cnt = 0;
  50.   int overlapSize = params->SSIMOverlapSize;
  51.   max_pix_value_sqd = (float) (img->max_imgpel_value_comp[comp] * img->max_imgpel_value_comp[comp]);
  52.   C2 = K2 * K2 * max_pix_value_sqd;
  53.   for (j = 0; j <= height - win_height; j += overlapSize)
  54.   {
  55.     for (i = 0; i <= width - win_width; i += overlapSize)
  56.     {
  57.       imeanOrg = 0;
  58.       imeanEnc = 0; 
  59.       ivarOrg  = 0;
  60.       ivarEnc  = 0;
  61.       icovOrgEnc = 0;
  62.       for ( n = j;n < j + win_height;n ++)
  63.       {
  64.         for (m = i;m < i + win_width;m ++)
  65.         {
  66.           imeanOrg   += refImg[n][m];
  67.           imeanEnc   += encImg[n][m];
  68.           ivarOrg    += refImg[n][m] * refImg[n][m];
  69.           ivarEnc    += encImg[n][m] * encImg[n][m];
  70.           icovOrgEnc += refImg[n][m] * encImg[n][m];
  71.         }
  72.       }
  73.       meanOrg = (float) imeanOrg / win_pixels;
  74.       meanEnc = (float) imeanEnc / win_pixels;
  75.       varOrg    = ((float) ivarOrg - ((float) imeanOrg) * meanOrg) / win_pixels_bias;
  76.       varEnc    = ((float) ivarEnc - ((float) imeanEnc) * meanEnc) / win_pixels_bias;
  77.       covOrgEnc = ((float) icovOrgEnc - ((float) imeanOrg) * meanEnc) / win_pixels_bias;
  78.       mb_ssim  = (float) (2.0 * covOrgEnc + C2);
  79.       mb_ssim /= (float) (varOrg + varEnc + C2);
  80.       cur_distortion += mb_ssim;
  81.       win_cnt++;
  82.     }
  83.   }
  84.   cur_distortion /= (float) win_cnt;
  85.   if (cur_distortion >= 1.0 && cur_distortion < 1.01) // avoid float accuracy problem at very low QP(e.g.2)
  86.     cur_distortion = 1.0;
  87.   return cur_distortion;
  88. }
  89. float compute_luminance_component (imgpel **refImg, imgpel **encImg, int height, int width, int win_height, int win_width, int comp)
  90. {
  91.   static const float K1 = 0.01f;
  92.   static float max_pix_value_sqd;
  93.   float C1;
  94.   float win_pixels = (float) (win_width * win_height);
  95.   float mb_ssim, meanOrg, meanEnc;
  96.   int imeanOrg, imeanEnc;
  97.   float cur_distortion = 0.0;
  98.   int i, j, n, m, win_cnt = 0;
  99.   int overlapSize = params->SSIMOverlapSize;
  100.   max_pix_value_sqd = (float) (img->max_imgpel_value_comp[comp] * img->max_imgpel_value_comp[comp]);
  101.   C1 = K1 * K1 * max_pix_value_sqd;
  102.   for (j = 0; j <= height - win_height; j += overlapSize)
  103.   {
  104.     for (i = 0; i <= width - win_width; i += overlapSize)
  105.     {
  106.       imeanOrg = 0;
  107.       imeanEnc = 0; 
  108.       
  109.       for ( n = j;n < j + win_height;n ++)
  110.       {
  111.         for (m = i;m < i + win_width;m ++)
  112.         {
  113.           imeanOrg   += refImg[n][m];
  114.           imeanEnc   += encImg[n][m];
  115.         }
  116.       }
  117.       meanOrg = (float) imeanOrg / win_pixels;
  118.       meanEnc = (float) imeanEnc / win_pixels;
  119.       mb_ssim  = (float) (2.0 * meanOrg * meanEnc + C1);
  120.       mb_ssim /= (float) (meanOrg * meanOrg + meanEnc * meanEnc + C1);
  121.       cur_distortion += mb_ssim;
  122.       win_cnt++;
  123.     }
  124.   }
  125.   cur_distortion /= (float) win_cnt;
  126.   if (cur_distortion >= 1.0 && cur_distortion < 1.01) // avoid float accuracy problem at very low QP(e.g.2)
  127.     cur_distortion = 1.0;
  128.   return cur_distortion;
  129. }
  130. void horizontal_symmetric_extension(int **buffer, int width, int height )
  131. {
  132.   int j;
  133.   int* buf;
  134.  
  135.   int height_plus_pad2 = height + MS_SSIM_PAD2;
  136.   int width_plus_pad2_minus_one  = width  + MS_SSIM_PAD2 - 1;
  137.   // horizontal PSE
  138.   for (j = MS_SSIM_PAD2; j < height_plus_pad2; j++ ) {
  139.     // left column
  140.     buf = &buffer[j][MS_SSIM_PAD2];
  141.     buf[-1] = buf[1];
  142.     buf[-2] = buf[2];
  143.     // right column
  144.     buf = &buffer[j][width_plus_pad2_minus_one];
  145.     buf[1] = buf[-1];
  146.     buf[2] = buf[-2];
  147.     buf[3] = buf[-3];
  148.   }
  149. }
  150. void vertical_symmetric_extension(int **buffer, int width, int height)
  151. {
  152.   int i;
  153.   int* bufminus1;
  154.   int* bufminus2;
  155.   int* bufminus3;
  156.   int* bufplus1;
  157.   int* bufplus2;
  158.   int* bufplus3;
  159.   bufminus1 = &buffer[MS_SSIM_PAD2-1][MS_SSIM_PAD2];
  160.   bufminus2 = &buffer[MS_SSIM_PAD2-2][MS_SSIM_PAD2];
  161.   bufplus1  = &buffer[MS_SSIM_PAD2+1][MS_SSIM_PAD2];
  162.   bufplus2  = &buffer[MS_SSIM_PAD2+2][MS_SSIM_PAD2];
  163.   for (i = 0; i < width; i++)
  164.   {
  165.     bufminus1[i] = bufplus1[i];
  166.     bufminus2[i] = bufplus2[i];
  167.   }
  168.   bufminus1 = &buffer[height + MS_SSIM_PAD2-2][MS_SSIM_PAD2];
  169.   bufminus2 = &buffer[height + MS_SSIM_PAD2-3][MS_SSIM_PAD2];
  170.   bufminus3 = &buffer[height + MS_SSIM_PAD2-4][MS_SSIM_PAD2];
  171.   bufplus1  = &buffer[height + MS_SSIM_PAD2][MS_SSIM_PAD2];
  172.   bufplus2  = &buffer[height + MS_SSIM_PAD2+1][MS_SSIM_PAD2];
  173.   bufplus3  = &buffer[height + MS_SSIM_PAD2+2][MS_SSIM_PAD2];
  174.   for (i = 0; i < width; i++)
  175.   {
  176.     bufplus1[i] = bufminus1[i];
  177.     bufplus2[i] = bufminus2[i];
  178.     bufplus3[i] = bufminus3[i];
  179.   }
  180. }
  181. static void imgpel_to_padded_int(imgpel** src, int **buffer, int width, int height)
  182. {
  183.   int i, j;
  184.   int* tmpDst;
  185.   imgpel* tmpSrc;
  186.   tmpDst = buffer[MS_SSIM_PAD2];
  187.   for (j = 0; j < height; j++)
  188.   {
  189.     tmpSrc = src[j];
  190.     tmpDst = &buffer[j + MS_SSIM_PAD2][MS_SSIM_PAD2];
  191.     for (i = 0; i < width; i++)
  192.     {
  193.       tmpDst[i] = (int)tmpSrc[i];
  194.     }
  195.   }
  196. }
  197. void downsample(imgpel** src, imgpel** out, int height, int width)
  198. {
  199.   int height2 = height >> 1;
  200.   int width2  = width  >> 1;
  201.   int i, j;
  202.   int ii, jj;
  203.   int iDst;
  204.   int tmp, tmp1, tmp2;
  205.   int* tmpDst;
  206.   int* tmpSrc;
  207.   
  208.   int** itemp;
  209.   int** dest;
  210.   get_mem2Dint(&itemp, height + MS_SSIM_PAD, width + MS_SSIM_PAD);
  211.   get_mem2Dint(&dest, height + MS_SSIM_PAD, width2 + MS_SSIM_PAD);
  212.   
  213.   imgpel_to_padded_int(src, itemp, width, height);
  214.   horizontal_symmetric_extension(itemp, width, height);
  215.   for (j = 0; j < height; j++)
  216.   {
  217.     tmpDst = dest[j+MS_SSIM_PAD2];
  218.     tmpSrc = itemp[j+MS_SSIM_PAD2];
  219.     iDst   = MS_SSIM_PAD2;
  220.     for (i = 0; i < width2; i++, iDst++)
  221.     {
  222.       ii = (i << 1) + MS_SSIM_PAD2;
  223.       tmp1 = tmpSrc[ii-1] + tmpSrc[ii+2];
  224.       tmp2 = tmpSrc[ii] + tmpSrc[ii+1];
  225.       tmpDst[iDst] = tmpSrc[ii-2] + tmpSrc[ii+3] + (tmp1 << 1) + tmp1 + (tmp2 << 5) - (tmp2 << 2);
  226.       tmpDst[iDst] >>= 6;
  227.     }
  228.   }
  229.  
  230.   //Periodic extension
  231.   vertical_symmetric_extension(dest, width2, height);
  232.   for (i = 0; i < width2; i++) 
  233.   {
  234.     ii = i + MS_SSIM_PAD2;
  235.     for (j = 0; j < height2; j++)
  236.     {
  237.       jj = (j << 1) + MS_SSIM_PAD2;
  238.       tmp1 = dest[jj-1][ii] + dest[jj+2][ii];
  239.       tmp2 = dest[jj][ii]   + dest[jj+1][ii];
  240.       tmp = dest[jj-2][ii] + dest[jj+3][ii] + (tmp1 << 1) + tmp1 + (tmp2 << 5) - (tmp2 << 2);
  241.       out[j][i] = (unsigned char) (tmp >> 6);   //Note: Should change for different bit depths
  242.     }
  243.   }
  244.   free_mem2Dint(itemp);
  245.   free_mem2Dint(dest);
  246. }
  247. float compute_ms_ssim(imgpel **refImg, imgpel **encImg, int height, int width, int win_height, int win_width, int comp)
  248. {
  249.   float structural[MAX_SSIM_LEVELS];
  250.   float cur_distortion;
  251.   float luminance;
  252.   imgpel** dsRef;
  253.   imgpel** dsEnc;
  254.   int m;
  255.   static int max_ssim_levels_minus_one = MAX_SSIM_LEVELS - 1;
  256.   static float exponent[5] = {(float)MS_SSIM_BETA0, (float)MS_SSIM_BETA1, (float)MS_SSIM_BETA2, (float)MS_SSIM_BETA3, (float)MS_SSIM_BETA4};
  257.   dsRef = NULL;
  258.   dsEnc = NULL;
  259.   get_mem2Dpel(&dsRef, height>>1, width>>1);
  260.   get_mem2Dpel(&dsEnc, height>>1, width>>1);
  261.   structural[0] = compute_structural_components(refImg, encImg, height, width, win_height, win_width, comp);
  262.   cur_distortion = (float)pow(structural[0], exponent[0]);
  263.   downsample(refImg, dsRef, height, width);
  264.   downsample(encImg, dsEnc, height, width);
  265.   for (m = 1; m < MAX_SSIM_LEVELS; m++)
  266.   {
  267.     height >>= 1;
  268.     width >>= 1;
  269.     structural[m] = compute_structural_components(dsRef, dsEnc, height, width, imin(win_height,height), imin(win_width,width), comp);
  270.     cur_distortion *= (float)pow(structural[m], exponent[m]);
  271.     if (m < max_ssim_levels_minus_one)
  272.     {
  273.       downsample(dsRef, dsRef, height, width);
  274.       downsample(dsEnc, dsEnc, height, width);
  275.     }
  276.     else
  277.     {
  278.       luminance = compute_luminance_component(dsRef, dsEnc, height, width, imin(win_height,height), imin(win_width,width), comp);
  279.       cur_distortion *= (float)pow(luminance, exponent[m]);
  280.     }
  281.   }
  282.   free_mem2Dpel(dsRef);
  283.   free_mem2Dpel(dsEnc);
  284.   dsRef = NULL;
  285.   dsEnc = NULL;
  286.   return cur_distortion;
  287. }
  288.   
  289. /*!
  290.  ************************************************************************
  291.  * brief
  292.  *    Find MS-SSIM for all three components
  293.  ************************************************************************
  294.  */
  295. void find_ms_ssim (ImageStructure *ref, ImageStructure *src, DistMetric *metricSSIM)
  296. {
  297.   FrameFormat *format = &ref->format;
  298.   metricSSIM->value[0] = compute_ms_ssim (ref->data[0], src->data[0], format->height, format->width, BLOCK_SIZE_8x8, BLOCK_SIZE_8x8, 0);
  299.   // Chroma.
  300.   if (format->yuv_format != YUV400)
  301.   {     
  302.     metricSSIM->value[1]  = compute_ms_ssim (ref->data[1], src->data[1], format->height_cr, format->width_cr, img->mb_cr_size_y, img->mb_cr_size_x, 1);
  303.     metricSSIM->value[2]  = compute_ms_ssim (ref->data[2], src->data[2], format->height_cr, format->width_cr, img->mb_cr_size_y, img->mb_cr_size_x, 2);
  304.   }
  305.   accumulate_average(metricSSIM,  dist->frame_ctr);
  306.   accumulate_avslice(metricSSIM,  img->type, stats->frame_ctr[img->type]);
  307. }