From 3e52b9675379bdfd4d29052e64a6ea39a67fb0e2 Mon Sep 17 00:00:00 2001
From: Angie Chiang <angiebird@google.com>
Date: Fri, 15 Mar 2019 18:30:38 -0700
Subject: [PATCH] Compute count_ls in vp9_kmeans

Change-Id: Id5a83554f2037b03a3a7d86f83e47cac311fbe1d
---
 vp9/encoder/vp9_encodeframe.c | 14 +++++++++-----
 vp9/encoder/vp9_encodeframe.h |  2 +-
 vp9/encoder/vp9_encoder.h     |  1 +
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index d0ebbae8d..9bf2fd189 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -5714,8 +5714,8 @@ int vp9_get_group_idx(double value, double *boundary_ls, int k) {
   return group_idx;
 }
 
-void vp9_kmeans(double *ctr_ls, double *boundary_ls, int k, KMEANS_DATA *arr,
-                int size) {
+void vp9_kmeans(double *ctr_ls, double *boundary_ls, int *count_ls, int k,
+                KMEANS_DATA *arr, int size) {
   double min, max;
   double step;
   int i, j;
@@ -5765,7 +5765,10 @@ void vp9_kmeans(double *ctr_ls, double *boundary_ls, int k, KMEANS_DATA *arr,
     }
   }
 
-  // compute group_idx
+  // compute group_idx, boundary_ls and count_ls
+  for (j = 0; j < k; ++j) {
+    count_ls[j] = 0;
+  }
   compute_boundary_ls(ctr_ls, k, boundary_ls);
   group_idx = 0;
   for (i = 0; i < size; ++i) {
@@ -5776,6 +5779,7 @@ void vp9_kmeans(double *ctr_ls, double *boundary_ls, int k, KMEANS_DATA *arr,
       }
     }
     arr[i].group_idx = group_idx;
+    ++count_ls[group_idx];
   }
 }
 
@@ -5890,8 +5894,8 @@ static void encode_frame_internal(VP9_COMP *cpi) {
           wiener_var_rdmult(cpi, BLOCK_64X64, mi_row, mi_col, cpi->rd.RDMULT);
 
       vp9_kmeans(cpi->kmeans_ctr_ls, cpi->kmeans_boundary_ls,
-                 cpi->kmeans_ctr_num, cpi->kmeans_data_arr,
-                 cpi->kmeans_data_size);
+                 cpi->kmeans_count_ls, cpi->kmeans_ctr_num,
+                 cpi->kmeans_data_arr, cpi->kmeans_data_size);
 
       vp9_perceptual_aq_mode_setup(cpi, &cm->seg);
     }
diff --git a/vp9/encoder/vp9_encodeframe.h b/vp9/encoder/vp9_encodeframe.h
index 29a56b923..fd0a9c517 100644
--- a/vp9/encoder/vp9_encodeframe.h
+++ b/vp9/encoder/vp9_encodeframe.h
@@ -46,7 +46,7 @@ void vp9_set_variance_partition_thresholds(struct VP9_COMP *cpi, int q,
                                            int content_state);
 
 struct KMEANS_DATA;
-void vp9_kmeans(double *ctr_ls, double *boundary_ls, int k,
+void vp9_kmeans(double *ctr_ls, double *boundary_ls, int *count_ls, int k,
                 struct KMEANS_DATA *arr, int size);
 int vp9_get_group_idx(double value, double *boundary_ls, int k);
 
diff --git a/vp9/encoder/vp9_encoder.h b/vp9/encoder/vp9_encoder.h
index 567f05934..67d764f80 100644
--- a/vp9/encoder/vp9_encoder.h
+++ b/vp9/encoder/vp9_encoder.h
@@ -611,6 +611,7 @@ typedef struct VP9_COMP {
   int kmeans_data_stride;
   double kmeans_ctr_ls[MAX_KMEANS_GROUPS];
   double kmeans_boundary_ls[MAX_KMEANS_GROUPS];
+  int kmeans_count_ls[MAX_KMEANS_GROUPS];
   int kmeans_ctr_num;
 #if CONFIG_NON_GREEDY_MV
   int tpl_ready;
-- 
2.40.0