 
     // Allocate memory on the GPU and copy data
     float *dU, *dV, *dP, *dD, *dDiaD;
-    cudaMalloc(&dU, (size_t)IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
-    cudaMalloc(&dV, (size_t)IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
-    cudaMalloc(&dD, (size_t)IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
-    cudaMalloc(&dP, (size_t)3*IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
-    cudaMalloc(&dDiaD, (size_t)3*IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
-    cudaMemcpy(dU, fInDepth, (size_t)IMG_WIDTH*IMG_HEIGHT*sizeof(float), cudaMemcpyHostToDevice); CUDA_CHECK;
-    cudaMemcpy(dV, dU, (size_t)IMG_WIDTH*IMG_HEIGHT*sizeof(float), cudaMemcpyDeviceToDevice); CUDA_CHECK;
-    cudaMemset(dD, 0, (size_t)IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
-    //cudaMemset(dP, 0, (size_t)3*IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
-    cudaMemset(dDiaD, 0, (size_t)3*IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
+    cudaMalloc(&dU, (size_t)width*height*sizeof(float)); CUDA_CHECK;
+    cudaMalloc(&dV, (size_t)width*height*sizeof(float)); CUDA_CHECK;
+    cudaMalloc(&dD, (size_t)width*height*sizeof(float)); CUDA_CHECK;
+    cudaMalloc(&dP, (size_t)3*width*height*sizeof(float)); CUDA_CHECK;
+    cudaMalloc(&dDiaD, (size_t)3*width*height*sizeof(float)); CUDA_CHECK;
+    cudaMemcpy(dU, fInDepth, (size_t)width*height*sizeof(float), cudaMemcpyHostToDevice); CUDA_CHECK;
+    cudaMemcpy(dV, dU, (size_t)width*height*sizeof(float), cudaMemcpyDeviceToDevice); CUDA_CHECK;
+    cudaMemset(dD, 0, (size_t)width*height*sizeof(float)); CUDA_CHECK;
+    cudaMemset(dP, 0, (size_t)3*width*height*sizeof(float)); CUDA_CHECK;
+    cudaMemset(dDiaD, 0, (size_t)3*width*height*sizeof(float)); CUDA_CHECK;
 
     // Init block and grid sizes
     dim3 block = dim3(blockX, blockY, blockZ);
-    dim3 grid = dim3((IMG_WIDTH+block.x-1)/block.x, (IMG_HEIGHT+block.y-1)/block.y, 1);
-
-    for(int n1=0; n1<N1; n1++)
-    {
+    dim3 grid = dim3((width+block.x-1)/block.x, (height+block.y-1)/block.y, 1);
+
+    for(int n=0; n<N; n++)
+    {
+        DiamondDotProduct<<<grid, block>>>(dP, dD, width, height);
+        cudaDeviceSynchronize();
+
         theta *= decay;
-        //cudaMemcpy(dV, dU, (size_t)IMG_WIDTH*IMG_HEIGHT*sizeof(float), cudaMemcpyDeviceToDevice); CUDA_CHECK;
-        cudaMemset(dP, 0, (size_t)3*IMG_WIDTH*IMG_HEIGHT*sizeof(float)); CUDA_CHECK;
-
-        for(int n2=0; n2<N2; n2++)
-        {
-            DiamondDotProduct<<<grid, block>>>(dP, dD, IMG_WIDTH, IMG_HEIGHT);
-            cudaDeviceSynchronize();
-
-            ComputeU<<<grid, block>>>(dV, dD, dU, IMG_WIDTH, IMG_HEIGHT, theta);
-            cudaDeviceSynchronize();
-            
-            DiamondOperator<<<grid, block>>>(dU, dDiaD, IMG_WIDTH, IMG_HEIGHT);
-            cudaDeviceSynchronize();
-
-            ComputeDualVariable<<<grid, block>>>(dDiaD, dP, IMG_WIDTH, IMG_HEIGHT, theta, tau);
-            cudaDeviceSynchronize();
-        }
-    }
-
-    DiamondDotProduct<<<grid, block>>>(dP, dD, IMG_WIDTH, IMG_HEIGHT);
+        ComputeU<<<grid, block>>>(dV, dD, dU, width, height, theta);
+        cudaDeviceSynchronize();
+        
+        DiamondOperator<<<grid, block>>>(dU, dDiaD, width, height);
+        cudaDeviceSynchronize();
+
+        ComputeDualVariable<<<grid, block>>>(dDiaD, dP, width, height, theta, tau);
+        cudaDeviceSynchronize();
+    }
+
+    DiamondDotProduct<<<grid, block>>>(dP, dD, width, height);
     cudaDeviceSynchronize();
 
-    ComputeU<<<grid, block>>>(dU, dV, dD, IMG_WIDTH, IMG_HEIGHT, theta);
+    //theta *= decay;
+    ComputeU<<<grid, block>>>(dV, dD, dU, width, height, theta);
     
     // Copy data back to CPU
-    cudaMemcpy(fOutDepth, dU, (size_t)IMG_WIDTH*IMG_HEIGHT*sizeof(float), cudaMemcpyDeviceToHost); CUDA_CHECK;
+    cudaMemcpy(fOutDepth, dU, (size_t)width*height*sizeof(float), cudaMemcpyDeviceToHost); CUDA_CHECK;
     
     // Deallocate memory on the GPU
     cudaFree(dU); CUDA_CHECK;


        float pp = p[idx];
        float a1 = 0.0f;    float a2 = 0.0f;    float a3 = 0.0f;
        float b1 = 0.0f;    float b2 = 0.0f;
        float c1 = 0.0f;    float c2 = 0.0f;    float c3 = 0.0f;
        float d1 = 0.0f;    float d2 = 0.0f;
                                                float e3 = 0.0f;

        if(x!=0)            { a1 = p[idx-1]; a2 = p[idx+(size_t)h*w-1]; a3 = p[idx+(size_t)2*h*w-1]; }
        if((x+1)!=w)        { b1 = p[idx+1]; b2 = p[idx+(size_t)h*w+1]; }
        if(y!=0)            { c1 = p[idx-w]; c2 = p[idx+(size_t)h*w-w]; c3 = p[idx+(size_t)2*h*w-w]; }
        if((y+1)!=h)        { d1 = p[idx+w]; d2 = p[idx+(size_t)h*w+w]; }
        if(y!=0 && x!=0)    {                                           e3  = p[idx+(size_t)2*h*w-w-1]; }
        //if((y+1)!=h && (x+1)!=h)    {                                   e3  = p[idx+(size_t)2*h*w+w+1]; }



// Uncomment to use the live Kinect Camera
//#define KINECT

__global__ void DiamondDotProduct(float *p, float *d, int w, int h)
{
