Home | History | Annotate | Download | only in aom_dsp
      1 /*
      2  * Copyright (c) 2016, Alliance for Open Media. All rights reserved
      3  *
      4  * This source code is subject to the terms of the BSD 2 Clause License and
      5  * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
      6  * was not distributed with this source code in the LICENSE file, you can
      7  * obtain it at www.aomedia.org/license/software. If the Alliance for Open
      8  * Media Patent License 1.0 was not distributed with this source code in the
      9  * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
     10  *
     11  *  This code was originally written by: Nathan E. Egge, at the Daala
     12  *  project.
     13  */
     14 #include <assert.h>
     15 #include <math.h>
     16 #include <stdlib.h>
     17 #include <string.h>
     18 
     19 #include "config/aom_config.h"
     20 #include "config/aom_dsp_rtcd.h"
     21 
     22 #include "aom_dsp/ssim.h"
     23 #include "aom_ports/system_state.h"
     24 
     25 typedef struct fs_level fs_level;
     26 typedef struct fs_ctx fs_ctx;
     27 
     28 #define SSIM_C1 (255 * 255 * 0.01 * 0.01)
     29 #define SSIM_C2 (255 * 255 * 0.03 * 0.03)
     30 #define SSIM_C1_10 (1023 * 1023 * 0.01 * 0.01)
     31 #define SSIM_C1_12 (4095 * 4095 * 0.01 * 0.01)
     32 #define SSIM_C2_10 (1023 * 1023 * 0.03 * 0.03)
     33 #define SSIM_C2_12 (4095 * 4095 * 0.03 * 0.03)
     34 
     35 #define FS_MINI(_a, _b) ((_a) < (_b) ? (_a) : (_b))
     36 #define FS_MAXI(_a, _b) ((_a) > (_b) ? (_a) : (_b))
     37 
     38 struct fs_level {
     39   uint32_t *im1;
     40   uint32_t *im2;
     41   double *ssim;
     42   int w;
     43   int h;
     44 };
     45 
     46 struct fs_ctx {
     47   fs_level *level;
     48   int nlevels;
     49   unsigned *col_buf;
     50 };
     51 
     52 static void fs_ctx_init(fs_ctx *_ctx, int _w, int _h, int _nlevels) {
     53   unsigned char *data;
     54   size_t data_size;
     55   int lw;
     56   int lh;
     57   int l;
     58   lw = (_w + 1) >> 1;
     59   lh = (_h + 1) >> 1;
     60   data_size =
     61       _nlevels * sizeof(fs_level) + 2 * (lw + 8) * 8 * sizeof(*_ctx->col_buf);
     62   for (l = 0; l < _nlevels; l++) {
     63     size_t im_size;
     64     size_t level_size;
     65     im_size = lw * (size_t)lh;
     66     level_size = 2 * im_size * sizeof(*_ctx->level[l].im1);
     67     level_size += sizeof(*_ctx->level[l].ssim) - 1;
     68     level_size /= sizeof(*_ctx->level[l].ssim);
     69     level_size += im_size;
     70     level_size *= sizeof(*_ctx->level[l].ssim);
     71     data_size += level_size;
     72     lw = (lw + 1) >> 1;
     73     lh = (lh + 1) >> 1;
     74   }
     75   data = (unsigned char *)malloc(data_size);
     76   _ctx->level = (fs_level *)data;
     77   _ctx->nlevels = _nlevels;
     78   data += _nlevels * sizeof(*_ctx->level);
     79   lw = (_w + 1) >> 1;
     80   lh = (_h + 1) >> 1;
     81   for (l = 0; l < _nlevels; l++) {
     82     size_t im_size;
     83     size_t level_size;
     84     _ctx->level[l].w = lw;
     85     _ctx->level[l].h = lh;
     86     im_size = lw * (size_t)lh;
     87     level_size = 2 * im_size * sizeof(*_ctx->level[l].im1);
     88     level_size += sizeof(*_ctx->level[l].ssim) - 1;
     89     level_size /= sizeof(*_ctx->level[l].ssim);
     90     level_size *= sizeof(*_ctx->level[l].ssim);
     91     _ctx->level[l].im1 = (uint32_t *)data;
     92     _ctx->level[l].im2 = _ctx->level[l].im1 + im_size;
     93     data += level_size;
     94     _ctx->level[l].ssim = (double *)data;
     95     data += im_size * sizeof(*_ctx->level[l].ssim);
     96     lw = (lw + 1) >> 1;
     97     lh = (lh + 1) >> 1;
     98   }
     99   _ctx->col_buf = (unsigned *)data;
    100 }
    101 
    102 static void fs_ctx_clear(fs_ctx *_ctx) { free(_ctx->level); }
    103 
    104 static void fs_downsample_level(fs_ctx *_ctx, int _l) {
    105   const uint32_t *src1;
    106   const uint32_t *src2;
    107   uint32_t *dst1;
    108   uint32_t *dst2;
    109   int w2;
    110   int h2;
    111   int w;
    112   int h;
    113   int i;
    114   int j;
    115   w = _ctx->level[_l].w;
    116   h = _ctx->level[_l].h;
    117   dst1 = _ctx->level[_l].im1;
    118   dst2 = _ctx->level[_l].im2;
    119   w2 = _ctx->level[_l - 1].w;
    120   h2 = _ctx->level[_l - 1].h;
    121   src1 = _ctx->level[_l - 1].im1;
    122   src2 = _ctx->level[_l - 1].im2;
    123   for (j = 0; j < h; j++) {
    124     int j0offs;
    125     int j1offs;
    126     j0offs = 2 * j * w2;
    127     j1offs = FS_MINI(2 * j + 1, h2) * w2;
    128     for (i = 0; i < w; i++) {
    129       int i0;
    130       int i1;
    131       i0 = 2 * i;
    132       i1 = FS_MINI(i0 + 1, w2);
    133       dst1[j * w + i] = src1[j0offs + i0] + src1[j0offs + i1] +
    134                         src1[j1offs + i0] + src1[j1offs + i1];
    135       dst2[j * w + i] = src2[j0offs + i0] + src2[j0offs + i1] +
    136                         src2[j1offs + i0] + src2[j1offs + i1];
    137     }
    138   }
    139 }
    140 
    141 static void fs_downsample_level0(fs_ctx *_ctx, const uint8_t *_src1,
    142                                  int _s1ystride, const uint8_t *_src2,
    143                                  int _s2ystride, int _w, int _h, uint32_t shift,
    144                                  int buf_is_hbd) {
    145   uint32_t *dst1;
    146   uint32_t *dst2;
    147   int w;
    148   int h;
    149   int i;
    150   int j;
    151   w = _ctx->level[0].w;
    152   h = _ctx->level[0].h;
    153   dst1 = _ctx->level[0].im1;
    154   dst2 = _ctx->level[0].im2;
    155   for (j = 0; j < h; j++) {
    156     int j0;
    157     int j1;
    158     j0 = 2 * j;
    159     j1 = FS_MINI(j0 + 1, _h);
    160     for (i = 0; i < w; i++) {
    161       int i0;
    162       int i1;
    163       i0 = 2 * i;
    164       i1 = FS_MINI(i0 + 1, _w);
    165       if (!buf_is_hbd) {
    166         dst1[j * w + i] =
    167             _src1[j0 * _s1ystride + i0] + _src1[j0 * _s1ystride + i1] +
    168             _src1[j1 * _s1ystride + i0] + _src1[j1 * _s1ystride + i1];
    169         dst2[j * w + i] =
    170             _src2[j0 * _s2ystride + i0] + _src2[j0 * _s2ystride + i1] +
    171             _src2[j1 * _s2ystride + i0] + _src2[j1 * _s2ystride + i1];
    172       } else {
    173         uint16_t *src1s = CONVERT_TO_SHORTPTR(_src1);
    174         uint16_t *src2s = CONVERT_TO_SHORTPTR(_src2);
    175         dst1[j * w + i] = (src1s[j0 * _s1ystride + i0] >> shift) +
    176                           (src1s[j0 * _s1ystride + i1] >> shift) +
    177                           (src1s[j1 * _s1ystride + i0] >> shift) +
    178                           (src1s[j1 * _s1ystride + i1] >> shift);
    179         dst2[j * w + i] = (src2s[j0 * _s2ystride + i0] >> shift) +
    180                           (src2s[j0 * _s2ystride + i1] >> shift) +
    181                           (src2s[j1 * _s2ystride + i0] >> shift) +
    182                           (src2s[j1 * _s2ystride + i1] >> shift);
    183       }
    184     }
    185   }
    186 }
    187 
    188 static void fs_apply_luminance(fs_ctx *_ctx, int _l, int bit_depth) {
    189   unsigned *col_sums_x;
    190   unsigned *col_sums_y;
    191   uint32_t *im1;
    192   uint32_t *im2;
    193   double *ssim;
    194   double c1;
    195   int w;
    196   int h;
    197   int j0offs;
    198   int j1offs;
    199   int i;
    200   int j;
    201   double ssim_c1 = SSIM_C1;
    202 
    203   if (bit_depth == 10) ssim_c1 = SSIM_C1_10;
    204   if (bit_depth == 12) ssim_c1 = SSIM_C1_12;
    205 
    206   w = _ctx->level[_l].w;
    207   h = _ctx->level[_l].h;
    208   col_sums_x = _ctx->col_buf;
    209   col_sums_y = col_sums_x + w;
    210   im1 = _ctx->level[_l].im1;
    211   im2 = _ctx->level[_l].im2;
    212   for (i = 0; i < w; i++) col_sums_x[i] = 5 * im1[i];
    213   for (i = 0; i < w; i++) col_sums_y[i] = 5 * im2[i];
    214   for (j = 1; j < 4; j++) {
    215     j1offs = FS_MINI(j, h - 1) * w;
    216     for (i = 0; i < w; i++) col_sums_x[i] += im1[j1offs + i];
    217     for (i = 0; i < w; i++) col_sums_y[i] += im2[j1offs + i];
    218   }
    219   ssim = _ctx->level[_l].ssim;
    220   c1 = (double)(ssim_c1 * 4096 * (1 << 4 * _l));
    221   for (j = 0; j < h; j++) {
    222     unsigned mux;
    223     unsigned muy;
    224     int i0;
    225     int i1;
    226     mux = 5 * col_sums_x[0];
    227     muy = 5 * col_sums_y[0];
    228     for (i = 1; i < 4; i++) {
    229       i1 = FS_MINI(i, w - 1);
    230       mux += col_sums_x[i1];
    231       muy += col_sums_y[i1];
    232     }
    233     for (i = 0; i < w; i++) {
    234       ssim[j * w + i] *= (2 * mux * (double)muy + c1) /
    235                          (mux * (double)mux + muy * (double)muy + c1);
    236       if (i + 1 < w) {
    237         i0 = FS_MAXI(0, i - 4);
    238         i1 = FS_MINI(i + 4, w - 1);
    239         mux += col_sums_x[i1] - col_sums_x[i0];
    240         muy += col_sums_x[i1] - col_sums_x[i0];
    241       }
    242     }
    243     if (j + 1 < h) {
    244       j0offs = FS_MAXI(0, j - 4) * w;
    245       for (i = 0; i < w; i++) col_sums_x[i] -= im1[j0offs + i];
    246       for (i = 0; i < w; i++) col_sums_y[i] -= im2[j0offs + i];
    247       j1offs = FS_MINI(j + 4, h - 1) * w;
    248       for (i = 0; i < w; i++) col_sums_x[i] += im1[j1offs + i];
    249       for (i = 0; i < w; i++) col_sums_y[i] += im2[j1offs + i];
    250     }
    251   }
    252 }
    253 
    254 #define FS_COL_SET(_col, _joffs, _ioffs)                       \
    255   do {                                                         \
    256     unsigned gx;                                               \
    257     unsigned gy;                                               \
    258     gx = gx_buf[((j + (_joffs)) & 7) * stride + i + (_ioffs)]; \
    259     gy = gy_buf[((j + (_joffs)) & 7) * stride + i + (_ioffs)]; \
    260     col_sums_gx2[(_col)] = gx * (double)gx;                    \
    261     col_sums_gy2[(_col)] = gy * (double)gy;                    \
    262     col_sums_gxgy[(_col)] = gx * (double)gy;                   \
    263   } while (0)
    264 
    265 #define FS_COL_ADD(_col, _joffs, _ioffs)                       \
    266   do {                                                         \
    267     unsigned gx;                                               \
    268     unsigned gy;                                               \
    269     gx = gx_buf[((j + (_joffs)) & 7) * stride + i + (_ioffs)]; \
    270     gy = gy_buf[((j + (_joffs)) & 7) * stride + i + (_ioffs)]; \
    271     col_sums_gx2[(_col)] += gx * (double)gx;                   \
    272     col_sums_gy2[(_col)] += gy * (double)gy;                   \
    273     col_sums_gxgy[(_col)] += gx * (double)gy;                  \
    274   } while (0)
    275 
    276 #define FS_COL_SUB(_col, _joffs, _ioffs)                       \
    277   do {                                                         \
    278     unsigned gx;                                               \
    279     unsigned gy;                                               \
    280     gx = gx_buf[((j + (_joffs)) & 7) * stride + i + (_ioffs)]; \
    281     gy = gy_buf[((j + (_joffs)) & 7) * stride + i + (_ioffs)]; \
    282     col_sums_gx2[(_col)] -= gx * (double)gx;                   \
    283     col_sums_gy2[(_col)] -= gy * (double)gy;                   \
    284     col_sums_gxgy[(_col)] -= gx * (double)gy;                  \
    285   } while (0)
    286 
    287 #define FS_COL_COPY(_col1, _col2)                    \
    288   do {                                               \
    289     col_sums_gx2[(_col1)] = col_sums_gx2[(_col2)];   \
    290     col_sums_gy2[(_col1)] = col_sums_gy2[(_col2)];   \
    291     col_sums_gxgy[(_col1)] = col_sums_gxgy[(_col2)]; \
    292   } while (0)
    293 
    294 #define FS_COL_HALVE(_col1, _col2)                         \
    295   do {                                                     \
    296     col_sums_gx2[(_col1)] = col_sums_gx2[(_col2)] * 0.5;   \
    297     col_sums_gy2[(_col1)] = col_sums_gy2[(_col2)] * 0.5;   \
    298     col_sums_gxgy[(_col1)] = col_sums_gxgy[(_col2)] * 0.5; \
    299   } while (0)
    300 
    301 #define FS_COL_DOUBLE(_col1, _col2)                      \
    302   do {                                                   \
    303     col_sums_gx2[(_col1)] = col_sums_gx2[(_col2)] * 2;   \
    304     col_sums_gy2[(_col1)] = col_sums_gy2[(_col2)] * 2;   \
    305     col_sums_gxgy[(_col1)] = col_sums_gxgy[(_col2)] * 2; \
    306   } while (0)
    307 
    308 static void fs_calc_structure(fs_ctx *_ctx, int _l, int bit_depth) {
    309   uint32_t *im1;
    310   uint32_t *im2;
    311   unsigned *gx_buf;
    312   unsigned *gy_buf;
    313   double *ssim;
    314   double col_sums_gx2[8];
    315   double col_sums_gy2[8];
    316   double col_sums_gxgy[8];
    317   double c2;
    318   int stride;
    319   int w;
    320   int h;
    321   int i;
    322   int j;
    323   double ssim_c2 = SSIM_C2;
    324   if (bit_depth == 10) ssim_c2 = SSIM_C2_10;
    325   if (bit_depth == 12) ssim_c2 = SSIM_C2_12;
    326 
    327   w = _ctx->level[_l].w;
    328   h = _ctx->level[_l].h;
    329   im1 = _ctx->level[_l].im1;
    330   im2 = _ctx->level[_l].im2;
    331   ssim = _ctx->level[_l].ssim;
    332   gx_buf = _ctx->col_buf;
    333   stride = w + 8;
    334   gy_buf = gx_buf + 8 * stride;
    335   memset(gx_buf, 0, 2 * 8 * stride * sizeof(*gx_buf));
    336   c2 = ssim_c2 * (1 << 4 * _l) * 16 * 104;
    337   for (j = 0; j < h + 4; j++) {
    338     if (j < h - 1) {
    339       for (i = 0; i < w - 1; i++) {
    340         unsigned g1;
    341         unsigned g2;
    342         unsigned gx;
    343         unsigned gy;
    344         g1 = abs((int)im1[(j + 1) * w + i + 1] - (int)im1[j * w + i]);
    345         g2 = abs((int)im1[(j + 1) * w + i] - (int)im1[j * w + i + 1]);
    346         gx = 4 * FS_MAXI(g1, g2) + FS_MINI(g1, g2);
    347         g1 = abs((int)im2[(j + 1) * w + i + 1] - (int)im2[j * w + i]);
    348         g2 = abs((int)im2[(j + 1) * w + i] - (int)im2[j * w + i + 1]);
    349         gy = 4 * FS_MAXI(g1, g2) + FS_MINI(g1, g2);
    350         gx_buf[(j & 7) * stride + i + 4] = gx;
    351         gy_buf[(j & 7) * stride + i + 4] = gy;
    352       }
    353     } else {
    354       memset(gx_buf + (j & 7) * stride, 0, stride * sizeof(*gx_buf));
    355       memset(gy_buf + (j & 7) * stride, 0, stride * sizeof(*gy_buf));
    356     }
    357     if (j >= 4) {
    358       int k;
    359       col_sums_gx2[3] = col_sums_gx2[2] = col_sums_gx2[1] = col_sums_gx2[0] = 0;
    360       col_sums_gy2[3] = col_sums_gy2[2] = col_sums_gy2[1] = col_sums_gy2[0] = 0;
    361       col_sums_gxgy[3] = col_sums_gxgy[2] = col_sums_gxgy[1] =
    362           col_sums_gxgy[0] = 0;
    363       for (i = 4; i < 8; i++) {
    364         FS_COL_SET(i, -1, 0);
    365         FS_COL_ADD(i, 0, 0);
    366         for (k = 1; k < 8 - i; k++) {
    367           FS_COL_DOUBLE(i, i);
    368           FS_COL_ADD(i, -k - 1, 0);
    369           FS_COL_ADD(i, k, 0);
    370         }
    371       }
    372       for (i = 0; i < w; i++) {
    373         double mugx2;
    374         double mugy2;
    375         double mugxgy;
    376         mugx2 = col_sums_gx2[0];
    377         for (k = 1; k < 8; k++) mugx2 += col_sums_gx2[k];
    378         mugy2 = col_sums_gy2[0];
    379         for (k = 1; k < 8; k++) mugy2 += col_sums_gy2[k];
    380         mugxgy = col_sums_gxgy[0];
    381         for (k = 1; k < 8; k++) mugxgy += col_sums_gxgy[k];
    382         ssim[(j - 4) * w + i] = (2 * mugxgy + c2) / (mugx2 + mugy2 + c2);
    383         if (i + 1 < w) {
    384           FS_COL_SET(0, -1, 1);
    385           FS_COL_ADD(0, 0, 1);
    386           FS_COL_SUB(2, -3, 2);
    387           FS_COL_SUB(2, 2, 2);
    388           FS_COL_HALVE(1, 2);
    389           FS_COL_SUB(3, -4, 3);
    390           FS_COL_SUB(3, 3, 3);
    391           FS_COL_HALVE(2, 3);
    392           FS_COL_COPY(3, 4);
    393           FS_COL_DOUBLE(4, 5);
    394           FS_COL_ADD(4, -4, 5);
    395           FS_COL_ADD(4, 3, 5);
    396           FS_COL_DOUBLE(5, 6);
    397           FS_COL_ADD(5, -3, 6);
    398           FS_COL_ADD(5, 2, 6);
    399           FS_COL_DOUBLE(6, 7);
    400           FS_COL_ADD(6, -2, 7);
    401           FS_COL_ADD(6, 1, 7);
    402           FS_COL_SET(7, -1, 8);
    403           FS_COL_ADD(7, 0, 8);
    404         }
    405       }
    406     }
    407   }
    408 }
    409 
    410 #define FS_NLEVELS (4)
    411 
    412 /*These weights were derived from the default weights found in Wang's original
    413  Matlab implementation: {0.0448, 0.2856, 0.2363, 0.1333}.
    414  We drop the finest scale and renormalize the rest to sum to 1.*/
    415 
    416 static const double FS_WEIGHTS[FS_NLEVELS] = {
    417   0.2989654541015625, 0.3141326904296875, 0.2473602294921875, 0.1395416259765625
    418 };
    419 
    420 static double fs_average(fs_ctx *_ctx, int _l) {
    421   double *ssim;
    422   double ret;
    423   int w;
    424   int h;
    425   int i;
    426   int j;
    427   w = _ctx->level[_l].w;
    428   h = _ctx->level[_l].h;
    429   ssim = _ctx->level[_l].ssim;
    430   ret = 0;
    431   for (j = 0; j < h; j++)
    432     for (i = 0; i < w; i++) ret += ssim[j * w + i];
    433   return pow(ret / (w * h), FS_WEIGHTS[_l]);
    434 }
    435 
    436 static double convert_ssim_db(double _ssim, double _weight) {
    437   assert(_weight >= _ssim);
    438   if ((_weight - _ssim) < 1e-10) return MAX_SSIM_DB;
    439   return 10 * (log10(_weight) - log10(_weight - _ssim));
    440 }
    441 
    442 static double calc_ssim(const uint8_t *_src, int _systride, const uint8_t *_dst,
    443                         int _dystride, int _w, int _h, uint32_t _bd,
    444                         uint32_t _shift, int buf_is_hbd) {
    445   fs_ctx ctx;
    446   double ret;
    447   int l;
    448   ret = 1;
    449   fs_ctx_init(&ctx, _w, _h, FS_NLEVELS);
    450   fs_downsample_level0(&ctx, _src, _systride, _dst, _dystride, _w, _h, _shift,
    451                        buf_is_hbd);
    452   for (l = 0; l < FS_NLEVELS - 1; l++) {
    453     fs_calc_structure(&ctx, l, _bd);
    454     ret *= fs_average(&ctx, l);
    455     fs_downsample_level(&ctx, l + 1);
    456   }
    457   fs_calc_structure(&ctx, l, _bd);
    458   fs_apply_luminance(&ctx, l, _bd);
    459   ret *= fs_average(&ctx, l);
    460   fs_ctx_clear(&ctx);
    461   return ret;
    462 }
    463 
    464 double aom_calc_fastssim(const YV12_BUFFER_CONFIG *source,
    465                          const YV12_BUFFER_CONFIG *dest, double *ssim_y,
    466                          double *ssim_u, double *ssim_v, uint32_t bd,
    467                          uint32_t in_bd) {
    468   double ssimv;
    469   uint32_t bd_shift = 0;
    470   aom_clear_system_state();
    471   assert(bd >= in_bd);
    472   assert(source->flags == dest->flags);
    473   int buf_is_hbd = source->flags & YV12_FLAG_HIGHBITDEPTH;
    474   bd_shift = bd - in_bd;
    475 
    476   *ssim_y = calc_ssim(source->y_buffer, source->y_stride, dest->y_buffer,
    477                       dest->y_stride, source->y_crop_width,
    478                       source->y_crop_height, in_bd, bd_shift, buf_is_hbd);
    479   *ssim_u = calc_ssim(source->u_buffer, source->uv_stride, dest->u_buffer,
    480                       dest->uv_stride, source->uv_crop_width,
    481                       source->uv_crop_height, in_bd, bd_shift, buf_is_hbd);
    482   *ssim_v = calc_ssim(source->v_buffer, source->uv_stride, dest->v_buffer,
    483                       dest->uv_stride, source->uv_crop_width,
    484                       source->uv_crop_height, in_bd, bd_shift, buf_is_hbd);
    485   ssimv = (*ssim_y) * .8 + .1 * ((*ssim_u) + (*ssim_v));
    486   return convert_ssim_db(ssimv, 1.0);
    487 }
    488