git.sesse.net Git - narabu/blob - psnr.cpp

   1 #include <stdio.h>
   2 #include <stdint.h>
   3 #include <stdlib.h>
   4 #include <assert.h>
   5 #include <math.h>
   6 #include <algorithm>
   7
   8 #define WIDTH 1280
   9 #define HEIGHT 720
  10 #define NUM_SYMS 256
  11 #define ESCAPE_LIMIT (NUM_SYMS - 1)
  12
  13 using namespace std;
  14
  15 unsigned char pix1[WIDTH * HEIGHT], pix2[WIDTH * HEIGHT];
  16
  17 void readpix(unsigned char *ptr, const char *filename)
  18 {
  19         FILE *fp = fopen(filename, "rb");
  20         if (fp == nullptr) {
  21                 perror(filename);
  22                 exit(1);
  23         }
  24
  25         fseek(fp, 0, SEEK_END);
  26         long len = ftell(fp);
  27         assert(len >= WIDTH * HEIGHT);
  28         fseek(fp, len - WIDTH * HEIGHT, SEEK_SET);
  29
  30         fread(ptr, 1, WIDTH * HEIGHT, fp);
  31         fclose(fp);
  32 }
  33
  34 /****************************************************************************
  35  * structural similarity metric
  36  ****************************************************************************/
  37 static void ssim_4x4x2_core( const uint8_t *pix1, intptr_t stride1,
  38                              const uint8_t *pix2, intptr_t stride2,
  39                              int sums[2][4] )
  40 {
  41     for( int z = 0; z < 2; z++ )
  42     {
  43         uint32_t s1 = 0, s2 = 0, ss = 0, s12 = 0;
  44         for( int y = 0; y < 4; y++ )
  45             for( int x = 0; x < 4; x++ )
  46             {
  47                 int a = pix1[x+y*stride1];
  48                 int b = pix2[x+y*stride2];
  49                 s1  += a;
  50                 s2  += b;
  51                 ss  += a*a;
  52                 ss  += b*b;
  53                 s12 += a*b;
  54             }
  55         sums[z][0] = s1;
  56         sums[z][1] = s2;
  57         sums[z][2] = ss;
  58         sums[z][3] = s12;
  59         pix1 += 4;
  60         pix2 += 4;
  61     }
  62 }
  63
  64 #define PIXEL_MAX 255
  65
  66 static float ssim_end1( int s1, int s2, int ss, int s12 )
  67 {
  68 /* Maximum value for 10-bit is: ss*64 = (2^10-1)^2*16*4*64 = 4286582784, which will overflow in some cases.
  69  * s1*s1, s2*s2, and s1*s2 also obtain this value for edge cases: ((2^10-1)*16*4)^2 = 4286582784.
  70  * Maximum value for 9-bit is: ss*64 = (2^9-1)^2*16*4*64 = 1069551616, which will not overflow. */
  71 #define type int
  72     static const int ssim_c1 = (int)(.01*.01*PIXEL_MAX*PIXEL_MAX*64 + .5);
  73     static const int ssim_c2 = (int)(.03*.03*PIXEL_MAX*PIXEL_MAX*64*63 + .5);
  74     type fs1 = s1;
  75     type fs2 = s2;
  76     type fss = ss;
  77     type fs12 = s12;
  78     type vars = fss*64 - fs1*fs1 - fs2*fs2;
  79     type covar = fs12*64 - fs1*fs2;
  80     return (float)(2*fs1*fs2 + ssim_c1) * (float)(2*covar + ssim_c2)
  81          / ((float)(fs1*fs1 + fs2*fs2 + ssim_c1) * (float)(vars + ssim_c2));
  82 #undef type
  83 }
  84
  85 static float ssim_end4( int sum0[5][4], int sum1[5][4], int width )
  86 {
  87     float ssim = 0.0;
  88     for( int i = 0; i < width; i++ )
  89         ssim += ssim_end1( sum0[i][0] + sum0[i+1][0] + sum1[i][0] + sum1[i+1][0],
  90                            sum0[i][1] + sum0[i+1][1] + sum1[i][1] + sum1[i+1][1],
  91                            sum0[i][2] + sum0[i+1][2] + sum1[i][2] + sum1[i+1][2],
  92                            sum0[i][3] + sum0[i+1][3] + sum1[i][3] + sum1[i+1][3] );
  93     return ssim;
  94 }
  95
  96 float x264_pixel_ssim_wxh(
  97                            uint8_t *pix1, intptr_t stride1,
  98                            uint8_t *pix2, intptr_t stride2,
  99                            int width, int height, void * buf, int *cnt )
 100 {
 101     int z = 0;
 102     float ssim = 0.0;
 103     int (*sum0)[4] = buf;
 104     int (*sum1)[4] = sum0 + (width >> 2) + 3;
 105     width >>= 2;
 106     height >>= 2;
 107     for( int y = 1; y < height; y++ )
 108     {
 109         for( ; z <= y; z++ )
 110         {
 111             swap(sum0, sum1);
 112             for( int x = 0; x < width; x+=2 )
 113                 ssim_4x4x2_core( &pix1[4*(x+z*stride1)], stride1, &pix2[4*(x+z*stride2)], stride2, &sum0[x] );
 114         }
 115         for( int x = 0; x < width-1; x += 4 )
 116             ssim += ssim_end4( sum0+x, sum1+x, min(4,width-x-1) );
 117     }
 118     *cnt = (height-1) * (width-1);
 119     return ssim;
 120 }
 121
 122 int main(int argc, char **argv)
 123 {
 124         readpix(pix1, argv[1]);
 125         readpix(pix2, argv[2]);
 126
 127         double sum_sq_err = 0.0;
 128
 129         for (unsigned y = 0; y < HEIGHT; ++y) {
 130                 for (unsigned x = 0; x < WIDTH; ++x) {
 131                         int k1 = pix1[y * WIDTH + x];
 132                         int k2 = pix2[y * WIDTH + x];
 133                         sum_sq_err += (k1 - k2) * (k1 - k2);
 134                 }
 135         }
 136         double mse = sum_sq_err / double(WIDTH * HEIGHT);
 137         double psnr_db = 20 * log10(255.0 / sqrt(mse));
 138
 139         void *scratch = new int[8 * WIDTH];
 140         int cnt = 0;
 141         float ssim = x264_pixel_ssim_wxh(
 142                                    pix1 + 2 + WIDTH*2, WIDTH,
 143                                    pix2 + 2 + WIDTH*2, WIDTH,
 144                                    WIDTH - 2, HEIGHT - 2, scratch, &cnt);
 145         ssim /= cnt;
 146
 147         printf("%.2f %.2f\n", psnr_db, -10.0 * log10(1 - ssim));
 148 }