1 /*
2 * Copyright (c) 2016, Alliance for Open Media. All rights reserved
3 *
4 * This source code is subject to the terms of the BSD 2 Clause License and
5 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6 * was not distributed with this source code in the LICENSE file, you can
7 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
8 * Media Patent License 1.0 was not distributed with this source code in the
9 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10 */
11
12 #include <assert.h>
13 #include <stdint.h>
14 #include <string.h>
15
16 #include "av1/encoder/palette.h"
17 #include "av1/encoder/random.h"
18
19 #ifndef AV1_K_MEANS_DIM
20 #error "This template requires AV1_K_MEANS_DIM to be defined"
21 #endif
22
23 #define RENAME_(x, y) AV1_K_MEANS_RENAME(x, y)
24 #define RENAME(x) RENAME_(x, AV1_K_MEANS_DIM)
25
RENAME(calc_dist)26 static int RENAME(calc_dist)(const int *p1, const int *p2) {
27 int dist = 0;
28 for (int i = 0; i < AV1_K_MEANS_DIM; ++i) {
29 const int diff = p1[i] - p2[i];
30 dist += diff * diff;
31 }
32 return dist;
33 }
34
RENAME(av1_calc_indices)35 void RENAME(av1_calc_indices)(const int *data, const int *centroids,
36 uint8_t *indices, int n, int k) {
37 for (int i = 0; i < n; ++i) {
38 int min_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM, centroids);
39 indices[i] = 0;
40 for (int j = 1; j < k; ++j) {
41 const int this_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM,
42 centroids + j * AV1_K_MEANS_DIM);
43 if (this_dist < min_dist) {
44 min_dist = this_dist;
45 indices[i] = j;
46 }
47 }
48 }
49 }
50
RENAME(calc_centroids)51 static void RENAME(calc_centroids)(const int *data, int *centroids,
52 const uint8_t *indices, int n, int k) {
53 int i, j;
54 int count[PALETTE_MAX_SIZE] = { 0 };
55 unsigned int rand_state = (unsigned int)data[0];
56 assert(n <= 32768);
57 memset(centroids, 0, sizeof(centroids[0]) * k * AV1_K_MEANS_DIM);
58
59 for (i = 0; i < n; ++i) {
60 const int index = indices[i];
61 assert(index < k);
62 ++count[index];
63 for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
64 centroids[index * AV1_K_MEANS_DIM + j] += data[i * AV1_K_MEANS_DIM + j];
65 }
66 }
67
68 for (i = 0; i < k; ++i) {
69 if (count[i] == 0) {
70 memcpy(centroids + i * AV1_K_MEANS_DIM,
71 data + (lcg_rand16(&rand_state) % n) * AV1_K_MEANS_DIM,
72 sizeof(centroids[0]) * AV1_K_MEANS_DIM);
73 } else {
74 for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
75 centroids[i * AV1_K_MEANS_DIM + j] =
76 DIVIDE_AND_ROUND(centroids[i * AV1_K_MEANS_DIM + j], count[i]);
77 }
78 }
79 }
80 }
81
RENAME(calc_total_dist)82 static int64_t RENAME(calc_total_dist)(const int *data, const int *centroids,
83 const uint8_t *indices, int n, int k) {
84 int64_t dist = 0;
85 (void)k;
86 for (int i = 0; i < n; ++i) {
87 dist += RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM,
88 centroids + indices[i] * AV1_K_MEANS_DIM);
89 }
90 return dist;
91 }
92
RENAME(av1_k_means)93 void RENAME(av1_k_means)(const int *data, int *centroids, uint8_t *indices,
94 int n, int k, int max_itr) {
95 int pre_centroids[2 * PALETTE_MAX_SIZE];
96 uint8_t pre_indices[MAX_SB_SQUARE];
97
98 RENAME(av1_calc_indices)(data, centroids, indices, n, k);
99 int64_t this_dist = RENAME(calc_total_dist)(data, centroids, indices, n, k);
100
101 for (int i = 0; i < max_itr; ++i) {
102 const int64_t pre_dist = this_dist;
103 memcpy(pre_centroids, centroids,
104 sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM);
105 memcpy(pre_indices, indices, sizeof(pre_indices[0]) * n);
106
107 RENAME(calc_centroids)(data, centroids, indices, n, k);
108 RENAME(av1_calc_indices)(data, centroids, indices, n, k);
109 this_dist = RENAME(calc_total_dist)(data, centroids, indices, n, k);
110
111 if (this_dist > pre_dist) {
112 memcpy(centroids, pre_centroids,
113 sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM);
114 memcpy(indices, pre_indices, sizeof(pre_indices[0]) * n);
115 break;
116 }
117 if (!memcmp(centroids, pre_centroids,
118 sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM))
119 break;
120 }
121 }
122 #undef RENAME_
123 #undef RENAME
124