Let vp9_kmeans provide boundary for each group

boundary_ls[j] is the upper bound of data centered at ctr_ls[j]

Add vp9_get_group_idx() for computing group_idx

Change-Id: I3b1b488edf8acbfb63c469eeeba15f3e42b0a645

3 files changed, 32 insertions, 13 deletions

diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index 22b5f06ac..57d74b031 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -8,6 +8,7 @@
  *  be found in the AUTHORS file in the root of the source tree.
  */
 
+#include <float.h>
 #include <limits.h>
 #include <math.h>
 #include <stdio.h>
@@ -5688,12 +5689,33 @@ static int compare_kmeans_data(const void *a, const void *b) {
   }
 }
 
-void vp9_kmeans(double *ctr_ls, int k, KMEANS_DATA *arr, int size) {
+static void compute_boundary_ls(const double *ctr_ls, int k,
+                                double *boundary_ls) {
+  // boundary_ls[j] is the upper bound of data centered at ctr_ls[j]
+  int j;
+  for (j = 0; j < k - 1; ++j) {
+    boundary_ls[j] = (ctr_ls[j] + ctr_ls[j + 1]) / 2.;
+  }
+  boundary_ls[k - 1] = DBL_MAX;
+}
+
+int vp9_get_group_idx(double value, double *boundary_ls, int k) {
+  int group_idx = 0;
+  while (value >= boundary_ls[group_idx]) {
+    ++group_idx;
+    if (group_idx == k - 1) {
+      break;
+    }
+  }
+  return group_idx;
+}
+
+void vp9_kmeans(double *ctr_ls, double *boundary_ls, int k, KMEANS_DATA *arr,
+                int size) {
   double min, max;
   double step;
   int i, j;
   int itr;
-  double boundary_ls[MAX_KMEANS_GROUPS] = { 0 };
   int group_idx;
   double sum;
   int count;
@@ -5714,11 +5736,7 @@ void vp9_kmeans(double *ctr_ls, int k, KMEANS_DATA *arr, int size) {
   }
 
   for (itr = 0; itr < 10; ++itr) {
-    for (j = 0; j < k - 1; ++j) {
-      boundary_ls[j] = (ctr_ls[j] + ctr_ls[j + 1]) / 2.;
-    }
-    boundary_ls[k - 1] = max + 1;
-
+    compute_boundary_ls(ctr_ls, k, boundary_ls);
     group_idx = 0;
     count = 0;
     sum = 0;
@@ -5744,10 +5762,7 @@ void vp9_kmeans(double *ctr_ls, int k, KMEANS_DATA *arr, int size) {
   }
 
   // compute group_idx
-  for (j = 0; j < k - 1; ++j) {
-    boundary_ls[j] = (ctr_ls[j] + ctr_ls[j + 1]) / 2.;
-  }
-  boundary_ls[k - 1] = max + 1;
+  compute_boundary_ls(ctr_ls, k, boundary_ls);
   group_idx = 0;
   for (i = 0; i < size; ++i) {
     while (arr[i].value >= boundary_ls[group_idx]) {
@@ -5890,7 +5905,8 @@ static void encode_frame_internal(VP9_COMP *cpi) {
     }
 
     if (cpi->sf.enable_wiener_variance && cm->show_frame) {
-      vp9_kmeans(cpi->kmeans_ctr_ls, cpi->kmeans_ctr_num, cpi->kmeans_data_arr,
+      vp9_kmeans(cpi->kmeans_ctr_ls, cpi->kmeans_boundary_ls,
+                 cpi->kmeans_ctr_num, cpi->kmeans_data_arr,
                  cpi->kmeans_data_size);
     }
 
diff --git a/vp9/encoder/vp9_encodeframe.h b/vp9/encoder/vp9_encodeframe.h
index a761ae68b..29a56b923 100644
--- a/vp9/encoder/vp9_encodeframe.h
+++ b/vp9/encoder/vp9_encodeframe.h
@@ -46,7 +46,9 @@ void vp9_set_variance_partition_thresholds(struct VP9_COMP *cpi, int q,
                                            int content_state);
 
 struct KMEANS_DATA;
-void vp9_kmeans(double *ctr_ls, int k, struct KMEANS_DATA *arr, int size);
+void vp9_kmeans(double *ctr_ls, double *boundary_ls, int k,
+                struct KMEANS_DATA *arr, int size);
+int vp9_get_group_idx(double value, double *boundary_ls, int k);
 
 #ifdef __cplusplus
 }  // extern "C"
diff --git a/vp9/encoder/vp9_encoder.h b/vp9/encoder/vp9_encoder.h
index 278408fd5..24ccaf99d 100644
--- a/vp9/encoder/vp9_encoder.h
+++ b/vp9/encoder/vp9_encoder.h
@@ -607,6 +607,7 @@ typedef struct VP9_COMP {
   int kmeans_data_size;
   int kmeans_data_stride;
   double kmeans_ctr_ls[MAX_KMEANS_GROUPS];
+  double kmeans_boundary_ls[MAX_KMEANS_GROUPS];
   int kmeans_ctr_num;
 #if CONFIG_NON_GREEDY_MV
   int tpl_ready;