1 files changed, 13 insertions, 24 deletions
diff --git a/java/tests/ImageProcessing_jb/src/com/android/rs/image/wbalance.rs b/java/tests/ImageProcessing_jb/src/com/android/rs/image/wbalance.rs
index b5ab14ff..c8036409 100644
--- a/java/tests/ImageProcessing_jb/src/com/android/rs/image/wbalance.rs
+++ b/java/tests/ImageProcessing_jb/src/com/android/rs/image/wbalance.rs
@@ -23,9 +23,7 @@ rs_allocation histogramSource;
 uint32_t histogramHeight;
 uint32_t histogramWidth;
 
-static float scaleR;
-static float scaleG;
-static float scaleB;
+static float3 scale;
 
 static uchar4 estimateWhite() {
 
@@ -115,28 +113,19 @@ void prepareWhiteBalance() {
     int maximum = max(estimation.r, max(estimation.g, estimation.b));
     float avg = (minimum + maximum) / 2.f;
 
-    scaleR =  avg/estimation.r;
-    scaleG =  avg/estimation.g;
-    scaleB =  avg/estimation.b;
-
-}
-
-static unsigned char contrastClamp(int c)
-{
-    int N = 255;
-    c &= ~(c >> 31);
-    c -= N;
-    c &= (c >> 31);
-    c += N;
-    return  (unsigned char) c;
+    scale.r =  avg/estimation.r;
+    scale.g =  avg/estimation.g;
+    scale.b =  avg/estimation.b;
 }
 
-void whiteBalanceKernel(const uchar4 *in, uchar4 *out) {
-    float Rc =  in->r*scaleR;
-    float Gc =  in->g*scaleG;
-    float Bc =  in->b*scaleB;
+uchar4 __attribute__((kernel)) whiteBalanceKernel(uchar4 in) {
+    float3 t = convert_float3(in.rgb);
+    t *= scale;
+    t = min(t, 255.f);
 
-    out->r = contrastClamp(Rc);
-    out->g = contrastClamp(Gc);
-    out->b = contrastClamp(Bc);
+    uchar4 out;
+    out.rgb = convert_uchar3(t);
+    out.a = 255;
+    return out;
 }
+