--- plugin_psnrhvsm.c	2010/10/10 19:19:46	1.1
+++ plugin_psnrhvsm.c	2010/11/10 21:25:16	1.3
@@ -19,7 +19,7 @@
  *  along with this program ; if not, write to the Free Software
  *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307 USA
  *
- * $Id: plugin_psnrhvsm.c,v 1.1 2010/10/10 19:19:46 Isibaar Exp $
+ * $Id: plugin_psnrhvsm.c,v 1.3 2010/11/10 21:25:16 Isibaar Exp $
  *
  ****************************************************************************/
 
@@ -48,8 +48,10 @@
 
 typedef struct { 
 
-	int pixels; /* width*height */
-	uint64_t mse_sum; /* for avrg psnrhvsm */
+	uint64_t mse_sum_y; /* for avrg psnr-hvs-m */
+	uint64_t mse_sum_u;
+	uint64_t mse_sum_v;
+
 	long frame_cnt;
 
 } psnrhvsm_data_t; /* internal plugin data */
@@ -76,7 +78,7 @@
 	  0.019290f, 0.011815f, 0.011080f, 0.010412f, 0.007972f, 0.010000f, 0.009426f, 0.010203f
 	};
 
-#if 1 /* Floating-point implementation */
+#if 0 /* Floating-point implementation */
 
 static uint32_t Calc_MSE_H(int16_t *DCT_A, int16_t *DCT_B, uint8_t *IMG_A, uint8_t *IMG_B, int stride)
 {
@@ -85,6 +87,7 @@
 	uint32_t Local[8] = {0, 0, 0, 0, 0, 0, 0, 0};
 	uint32_t Local_Square[8] = {0, 0, 0, 0, 0, 0, 0, 0};
 	float MASK_A = 0.f, MASK_B = 0.f;
+	float Mult1 = 1.f, Mult2 = 1.f;
 	uint32_t MSE_H = 0;
 
 	/* Step 1: Calculate CSF weighted energy of DCT coefficients */
@@ -98,14 +101,16 @@
 	/* Step 2: Determine local variances compared to entire block variance */
 	for (y = 0; y < 2; y++) {
 		for (x = 0; x < 2; x++) {
-			for (i = 0; i < 4; i++) {
-				uint8_t A = IMG_A[y*4*stride + 4*x + i];
-				uint8_t B = IMG_B[y*4*stride + 4*x + i];
-
-				Local[y*2 + x] += A;
-				Local[y*2 + x + 4] += B;
-				Local_Square[y*2 + x] += A*A;
-				Local_Square[y*2 + x + 4] += B*B;
+			for (j = 0; j < 4; j++) {
+				for (i = 0; i < 4; i++) {
+					uint8_t A = IMG_A[(y*4+j)*stride + 4*x + i];
+					uint8_t B = IMG_B[(y*4+j)*stride + 4*x + i];
+
+					Local[y*2 + x] += A;
+					Local[y*2 + x + 4] += B;
+					Local_Square[y*2 + x] += A*A;
+					Local_Square[y*2 + x + 4] += B*B;
+				}
 			}
 		}
 	}
@@ -123,8 +128,14 @@
 	Global_B = (Local_Square[4]<<6) - Global_B*Global_B; /* 64*Var(D) */
 
 	/* Step 3: Calculate contrast masking threshold */
-	MASK_A = (float)sqrt(MASK_A*((float)(Local[0]+Local[1]+Local[2]+Local[3])/((float)Global_A/4.f)))/32.f;
-	MASK_B = (float)sqrt(MASK_B*((float)(Local[4]+Local[5]+Local[6]+Local[7])/((float)Global_B/4.f)))/32.f;
+	if (Global_A) 
+		Mult1 = (float)(Local[0]+Local[1]+Local[2]+Local[3])/((float)(Global_A)/4.f);
+
+	if (Global_B)
+		Mult2 = (float)(Local[4]+Local[5]+Local[6]+Local[7])/((float)(Global_B)/4.f);
+
+	MASK_A = (float)sqrt(MASK_A * Mult1) / 32.f;
+	MASK_B = (float)sqrt(MASK_B * Mult2) / 32.f;
 
 	if (MASK_B > MASK_A) MASK_A = MASK_B; /* MAX(MASK_A, MASK_B) */
 
@@ -146,8 +157,10 @@
 	return MSE_H; /* Fixed-point value right-shifted by eight */
 }
 
-#else /* First draft of a fixed-point implementation.
-	     Might serve as a template for MMX/SSE code */
+#else 
+
+/* First draft of a fixed-point implementation.
+   Might serve as a template for MMX/SSE code */
 
 static const uint16_t iMask_Coeff[64] = 
 	{     0, 59577, 65535, 40959, 27306, 16384, 12850, 10743,
@@ -182,6 +195,21 @@
 	   366,  286,  277,  269,  235, 264, 256, 266
 	};
 
+static __inline uint32_t isqrt(unsigned long n)
+{
+    uint32_t c = 0x8000;
+    uint32_t g = 0x8000;
+
+    for(;;) {
+        if(g*g > n)
+            g ^= c;
+        c >>= 1;
+        if(c == 0)
+            return g;
+        g |= c;
+    }
+}
+
 static uint32_t Calc_MSE_H(int16_t *DCT_A, int16_t *DCT_B, uint8_t *IMG_A, uint8_t *IMG_B, int stride)
 {
 	int x, y, i, j;
@@ -205,14 +233,16 @@
 	/* Step 2: Determine local variances compared to entire block variance */
 	for (y = 0; y < 2; y++) {
 		for (x = 0; x < 2; x++) {
-			for (i = 0; i < 4; i++) {
-				uint8_t A = IMG_A[y*4*stride + 4*x + i];
-				uint8_t B = IMG_B[y*4*stride + 4*x + i];
-
-				Local[y*2 + x] += A;
-				Local[y*2 + x + 4] += B;
-				Local_Square[y*2 + x] += A*A;
-				Local_Square[y*2 + x + 4] += B*B;
+			for (j = 0; j < 4; j++) {
+				for (i = 0; i < 4; i++) {
+					uint8_t A = IMG_A[(y*4+j)*stride + 4*x + i];
+					uint8_t B = IMG_B[(y*4+j)*stride + 4*x + i];
+
+					Local[y*2 + x] += A;
+					Local[y*2 + x + 4] += B;
+					Local_Square[y*2 + x] += A*A;
+					Local_Square[y*2 + x + 4] += B*B;
+				}
 			}
 		}
 	}
@@ -231,18 +261,22 @@
 
 	/* Step 3: Calculate contrast masking threshold */
 	{ 
-		float MASK_A, MASK_B; /* TODO */
+		uint32_t MASK_A, MASK_B;
+		uint32_t Mult1 = 64, Mult2 = 64;
+
+		if (Global_A) 
+			Mult1 = ((Local[0]+Local[1]+Local[2]+Local[3])<<8) / Global_A;
 
-		MASK_A = (float)((double)(Sum_A + 4) / 8.);
-		MASK_B = (float)((double)(Sum_B + 4) / 8.);
+		if (Global_B)
+			Mult2 = ((Local[4]+Local[5]+Local[6]+Local[7])<<8) / Global_B;
 
-		MASK_A = sqrt(MASK_A*((float)(Local[0]+Local[1]+Local[2]+Local[3])/((float)Global_A/4.f)))/32.f;
-		MASK_B = sqrt(MASK_B*((float)(Local[4]+Local[5]+Local[6]+Local[7])/((float)Global_B/4.f)))/32.f;
+		MASK_A = isqrt(2*Sum_A*Mult1) + 16;
+		MASK_B = isqrt(2*Sum_B*Mult2) + 16;
 
 		if (MASK_B > MASK_A)  /* MAX(MASK_A, MASK_B) */
-			MASK = (uint32_t) (MASK_B * 1024.f + 0.5f);
+			MASK = (uint32_t) MASK_B;
 		else
-			MASK = (uint32_t) (MASK_A * 1024.f + 0.5f);
+			MASK = (uint32_t) MASK_A;
 	}
 
 	/* Step 4: Calculate MSE of DCT coeffs reduced by masking effect */
@@ -272,17 +306,17 @@
 static void psnrhvsm_after(xvid_plg_data_t *data, psnrhvsm_data_t *psnrhvsm)
 {
 	DECLARE_ALIGNED_MATRIX(DCT, 2, 64, int16_t, CACHE_LINE);
-	int x, y, stride = data->original.stride[0];
+	int32_t x, y, u, v; 
 	int16_t *DCT_A = &DCT[0], *DCT_B = &DCT[64];
-	uint8_t *IMG_A = (uint8_t *) data->original.plane[0];
-	uint8_t *IMG_B = (uint8_t *) data->current.plane[0];
-	uint32_t MSE_H = 0;
+	uint64_t sse_y = 0, sse_u = 0, sse_v = 0;
 
-	psnrhvsm->pixels = data->width * data->height;
+	for (y = 0; y < data->height>>3; y++) {
+		uint8_t *IMG_A = (uint8_t *) data->original.plane[0];
+		uint8_t *IMG_B = (uint8_t *) data->current.plane[0];
+		uint32_t stride = data->original.stride[0];
 
-	for (y = 0; y < data->height; y += 8) {
-		for (x = 0; x < data->width; x += 8) {
-			int offset = y*stride + x;
+		for (x = 0; x < data->width>>3; x++) { /* non multiple of 8 handling ?? */
+			int offset = (y<<3)*stride + (x<<3);
 
 			emms();
 
@@ -297,14 +331,62 @@
 			emms();
 
 			/* Calculate MSE_H reduced by contrast masking effect */
-			MSE_H += Calc_MSE_H(DCT_A, DCT_B, IMG_A + offset, IMG_B + offset, stride);
+			sse_y += Calc_MSE_H(DCT_A, DCT_B, IMG_A + offset, IMG_B + offset, stride);
+		}
+	}
+
+	for (y = 0; y < data->height>>4; y++) {
+		uint8_t *U_A = (uint8_t *) data->original.plane[1];
+		uint8_t *V_A = (uint8_t *) data->original.plane[2];
+		uint8_t *U_B = (uint8_t *) data->current.plane[1];
+		uint8_t *V_B = (uint8_t *) data->current.plane[2];
+		uint32_t stride_uv = data->current.stride[1];
+
+		for (x = 0; x < data->width>>4; x++) { /* non multiple of 8 handling ?? */
+			int offset = (y<<3)*stride_uv + (x<<3);
+
+			emms();
+
+			/* Transfer data */
+			transfer_8to16copy(DCT_A, U_A + offset, stride_uv);
+			transfer_8to16copy(DCT_B, U_B + offset, stride_uv);
+
+			/* Perform DCT */
+			fdct(DCT_A);
+			fdct(DCT_B);
+
+			emms();
+
+			/* Calculate MSE_H reduced by contrast masking effect */
+			sse_u += Calc_MSE_H(DCT_A, DCT_B, U_A + offset, U_B + offset, stride_uv);
+
+			emms();
+
+			/* Transfer data */
+			transfer_8to16copy(DCT_A, V_A + offset, stride_uv);
+			transfer_8to16copy(DCT_B, V_B + offset, stride_uv);
+
+			/* Perform DCT */
+			fdct(DCT_A);
+			fdct(DCT_B);
+
+			emms();
+
+			/* Calculate MSE_H reduced by contrast masking effect */
+			sse_v += Calc_MSE_H(DCT_A, DCT_B, V_A + offset, V_B + offset, stride_uv);
 		}
 	}
 
-	psnrhvsm->mse_sum += MSE_H;
+	y = (int32_t) ( 4*sse_y / (data->width * data->height));
+	u = (int32_t) (16*sse_u / (data->width * data->height));
+	v = (int32_t) (16*sse_v / (data->width * data->height));
+
+	psnrhvsm->mse_sum_y += y;
+	psnrhvsm->mse_sum_u += u;
+	psnrhvsm->mse_sum_v += v;
 	psnrhvsm->frame_cnt++;
 
-	printf("       psnrhvsm: %2.2f\n", sse_to_PSNR(MSE_H, 256*psnrhvsm->pixels));
+	printf("       psnrhvsm y: %2.2f, psnrhvsm u: %2.2f, psnrhvsm v: %2.2f\n", sse_to_PSNR(y, 1024), sse_to_PSNR(u, 1024), sse_to_PSNR(v, 1024));
 }
 
 static int psnrhvsm_create(xvid_plg_create_t *create, void **handle)
@@ -312,10 +394,11 @@
 	psnrhvsm_data_t *psnrhvsm;
 	psnrhvsm = (psnrhvsm_data_t *) malloc(sizeof(psnrhvsm_data_t));
 
-	psnrhvsm->mse_sum = 0;
-	psnrhvsm->frame_cnt = 0;
+	psnrhvsm->mse_sum_y = 0;
+	psnrhvsm->mse_sum_u = 0;
+	psnrhvsm->mse_sum_v = 0;
 
-	psnrhvsm->pixels = 0;
+	psnrhvsm->frame_cnt = 0;
 
 	*(handle) = (void*) psnrhvsm;
 	return 0;
@@ -338,14 +421,17 @@
 			break;
 		case(XVID_PLG_DESTROY):
 			{
-				uint32_t MSE_H;
+				uint32_t y, u, v;
 				psnrhvsm_data_t *psnrhvsm = (psnrhvsm_data_t *)handle;
 				
 				if (psnrhvsm) {
-					MSE_H = (uint32_t) (psnrhvsm->mse_sum / psnrhvsm->frame_cnt);
+					y = (uint32_t) (psnrhvsm->mse_sum_y / psnrhvsm->frame_cnt);
+					u = (uint32_t) (psnrhvsm->mse_sum_u / psnrhvsm->frame_cnt);
+					v = (uint32_t) (psnrhvsm->mse_sum_v / psnrhvsm->frame_cnt);
 
 					emms();
-					printf("Average psnrhvsm: %2.2f\n", sse_to_PSNR(MSE_H, 256*psnrhvsm->pixels));
+					printf("Average psnrhvsm y: %2.2f, psnrhvsm u: %2.2f, psnrhvsm v: %2.2f\n", 
+						sse_to_PSNR(y, 1024), sse_to_PSNR(u, 1024), sse_to_PSNR(v, 1024));
 					free(psnrhvsm);
 				}
 			}