[x265] [PATCH] calculate SSIM for each Row after deblock, sao

[Aarthi Thirumalai]

# HG changeset patch
# User Aarthi Thirumalai
# Date 1381232542 -19800
#      Tue Oct 08 17:12:22 2013 +0530
# Node ID 3ba555e98be2c0641d92d14155c6fb303c0003d9
# Parent  47286888d5a293234340153112810ce2a4f81546
calculate SSIM for each Row after deblock, sao

diff -r 47286888d5a2 -r 3ba555e98be2 source/Lib/TLibEncoder/TEncTop.cpp
--- a/source/Lib/TLibEncoder/TEncTop.cpp	Tue Oct 08 17:06:57 2013 +0530
+++ b/source/Lib/TLibEncoder/TEncTop.cpp	Tue Oct 08 17:12:22 2013 +0530
@@ -513,7 +513,7 @@
     int width  = recon->getWidth() - getPad(0);
     int height = recon->getHeight() - getPad(1);
     int size = width * height;
-
+    double ssim = 0;
     UInt64 ssdY = computeSSD(orig->getLumaAddr(), recon->getLumaAddr(), stride, width, height);
 
     height >>= 1;
@@ -604,7 +604,14 @@
     {
         m_analyzeB.addResult(psnrY, psnrU, psnrV, (double)bits);
     }
-
+    if (param.bEnableSsim)
+    {
+        if(pic->getSlice()->m_ssimCnt > 0)
+        {
+            ssim += pic->getSlice()->m_ssim / pic->getSlice()->m_ssimCnt;
+            m_globalSsim += ssim;
+        }
+    }
     if (param.logLevel >= X265_LOG_DEBUG)
     {
         char c = (slice->isIntra() ? 'I' : slice->isInterP() ? 'P' : 'B');
diff -r 47286888d5a2 -r 3ba555e98be2 source/encoder/frameencoder.cpp
--- a/source/encoder/frameencoder.cpp	Tue Oct 08 17:06:57 2013 +0530
+++ b/source/encoder/frameencoder.cpp	Tue Oct 08 17:12:22 2013 +0530
@@ -92,7 +92,7 @@
     }
 
     m_frameFilter.destroy();
-
+    X265_FREE(ssimBuf);
     // wait for worker thread to exit
     stop();
 }
@@ -111,6 +111,9 @@
         m_rows[i].create(top);
     }
 
+    if (m_cfg->param.bEnableSsim)
+        CHECKED_MALLOC(ssimBuf, ssim_t, 8 * (m_cfg->param.sourceWidth / 4 + 3));
+
     // NOTE: 2 times of numRows because both Encoder and Filter in same queue
     if (!WaveFront::init(m_numRows * 2))
     {
@@ -168,6 +171,8 @@
         assert(0);
     }
     start();
+fail:
+    return;
 }
 
 int FrameEncoder::getStreamHeaders(NALUnitEBSP **nalunits)
@@ -535,6 +540,31 @@
         slice->setSaoEnabledFlag((saoParam->bSaoFlag[0] == 1) ? true : false);
     }
 
+    /*Compute SSIM if enabled*/
+    if (m_cfg->param.bEnableSsim && ssimBuf)
+    {
+        pixel *rec = (pixel*)m_pic->getPicYuvRec()->getLumaAddr();
+        pixel *org = (pixel*)m_pic->getPicYuvOrg()->getLumaAddr();
+        int stride1 = m_pic->getPicYuvOrg()->getStride();
+        int stride2 = m_pic->getPicYuvRec()->getStride();
+        for (int row = 0; row < m_numRows; row++)
+        {
+            int bEnd = ((row + 1) == (this->m_numRows - 1));
+            int bStart = (row == 0);
+            int minPixY = row * 64 - 4 * !bStart;
+            int maxPixY = (row + 1) * 64 - 4 * !bEnd;
+            int ssim_cnt;
+            x265_emms();
+
+            /* SSIM is done for each row in blocks of 4x4 . The First blocks are offset by 2 pixels to the right
+            * to avoid alignment of ssim blocks with DCT blocks. */
+            minPixY += bStart ? 2 : -6;
+            slice->m_ssim += calculateSSIM(rec + 2 + minPixY * stride1, stride1, org + 2 + minPixY * stride2, stride2, 
+                m_cfg->param.sourceWidth - 2, maxPixY - minPixY, ssimBuf, &ssim_cnt);
+            slice->m_ssimCnt += ssim_cnt;
+        }
+    }
+
     entropyCoder->setBitstream(NULL);
 
     // Reconstruction slice
@@ -681,6 +711,39 @@
     delete bitstreamRedirect;
 }
 
+/* Function to calculate SSIM for each row */
+float FrameEncoder::calculateSSIM(pixel *pix1, intptr_t stride1, pixel *pix2, intptr_t stride2, int width, int height, void *buf, int *cnt)
+{
+    int z = 0;
+    float ssim = 0.0;
+    ssim_t(*sum0)[4] = (ssim_t(*)[4])buf;
+    ssim_t(*sum1)[4] = sum0 + (width >> 2) + 3;
+    width >>= 2;
+    height >>= 2;
+
+    for (int y = 1; y < height; y++)
+    {
+        for (; z <= y; z++)
+        {
+            void* swap = sum0;
+            sum0 = sum1;
+            sum1 = (ssim_t(*)[4])swap;
+            for (int x = 0; x < width; x += 2)
+            {
+                primitives.ssim_4x4x2_core(&pix1[(4 * x + (z * stride1))], stride1, &pix2[(4 * x + (z * stride2))], stride2, &sum0[x]);
+            }
+        }
+
+        for (int x = 0; x < width - 1; x += 4)
+        {
+            ssim += primitives.ssim_end_4(sum0 + x, sum1 + x, X265_MIN(4, width - x - 1));
+        }
+    }
+
+    *cnt = (height - 1) * (width - 1);
+    return ssim; 
+} 
+
 void FrameEncoder::encodeSlice(TComOutputBitstream* substreams)
 {
     // choose entropy coder
diff -r 47286888d5a2 -r 3ba555e98be2 source/encoder/frameencoder.h
--- a/source/encoder/frameencoder.h	Tue Oct 08 17:06:57 2013 +0530
+++ b/source/encoder/frameencoder.h	Tue Oct 08 17:12:22 2013 +0530
@@ -145,6 +145,9 @@
     /* called by compressFrame to perform wave-front compression analysis */
     void compressCTURows();
 
+    /* called by compressFrame to calculate SSIM for each row . */
+    float calculateSSIM(pixel *pix1, intptr_t stride1, pixel *pix2, intptr_t stride2, int width, int height, void *buf, int *cnt);
+
     void encodeSlice(TComOutputBitstream* substreams);
 
     /* blocks until worker thread is done, returns encoded picture and bitstream */
@@ -185,7 +188,10 @@
     int                      m_filterRowDelay;
     CTURow*                  m_rows;
     Event                    m_completionEvent;
+
+    /* Temp Storage for ssim computation that doesnt need repeated malloc */
+    void *                   ssimBuf;
 };
 }
 
-#endif // ifndef X265_FRAMEENCODER_H
+#endif // ifndef X265_FRAMEENCODER_H
\ No newline at end of file