k_means_template.h (5030B)
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 <stdlib.h> 15 #include <string.h> 16 17 #include "av1/common/blockd.h" 18 #include "av1/encoder/palette.h" 19 #include "av1/encoder/random.h" 20 21 #ifndef AV1_K_MEANS_DIM 22 #error "This template requires AV1_K_MEANS_DIM to be defined" 23 #endif 24 25 #define RENAME_(x, y) AV1_K_MEANS_RENAME(x, y) 26 #define RENAME(x) RENAME_(x, AV1_K_MEANS_DIM) 27 #define K_MEANS_RENAME_C(x, y) x##_dim##y##_c 28 #define RENAME_C_(x, y) K_MEANS_RENAME_C(x, y) 29 #define RENAME_C(x) RENAME_C_(x, AV1_K_MEANS_DIM) 30 31 // Though we want to compute the smallest L2 norm, in 1 dimension, 32 // it is equivalent to find the smallest L1 norm and then square it. 33 // This is preferrable for speed, especially on the SIMD side. 34 static int RENAME(calc_dist)(const int16_t *p1, const int16_t *p2) { 35 #if AV1_K_MEANS_DIM == 1 36 return abs(p1[0] - p2[0]); 37 #else 38 int dist = 0; 39 for (int i = 0; i < AV1_K_MEANS_DIM; ++i) { 40 const int diff = p1[i] - p2[i]; 41 dist += diff * diff; 42 } 43 return dist; 44 #endif 45 } 46 47 void RENAME_C(av1_calc_indices)(const int16_t *data, const int16_t *centroids, 48 uint8_t *indices, int64_t *dist, int n, int k) { 49 if (dist) { 50 *dist = 0; 51 } 52 for (int i = 0; i < n; ++i) { 53 int min_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM, centroids); 54 indices[i] = 0; 55 for (int j = 1; j < k; ++j) { 56 const int this_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM, 57 centroids + j * AV1_K_MEANS_DIM); 58 if (this_dist < min_dist) { 59 min_dist = this_dist; 60 indices[i] = j; 61 } 62 } 63 if (dist) { 64 #if AV1_K_MEANS_DIM == 1 65 *dist += min_dist * min_dist; 66 #else 67 *dist += min_dist; 68 #endif 69 } 70 } 71 } 72 73 static void RENAME(calc_centroids)(const int16_t *data, int16_t *centroids, 74 const uint8_t *indices, int n, int k) { 75 int i, j; 76 int count[PALETTE_MAX_SIZE] = { 0 }; 77 int centroids_sum[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE]; 78 unsigned int rand_state = (unsigned int)data[0]; 79 assert(n <= 32768); 80 memset(centroids_sum, 0, sizeof(centroids_sum[0]) * k * AV1_K_MEANS_DIM); 81 82 for (i = 0; i < n; ++i) { 83 const int index = indices[i]; 84 assert(index < k); 85 ++count[index]; 86 for (j = 0; j < AV1_K_MEANS_DIM; ++j) { 87 centroids_sum[index * AV1_K_MEANS_DIM + j] += 88 data[i * AV1_K_MEANS_DIM + j]; 89 } 90 } 91 92 for (i = 0; i < k; ++i) { 93 if (count[i] == 0) { 94 memcpy(centroids + i * AV1_K_MEANS_DIM, 95 data + (lcg_rand16(&rand_state) % n) * AV1_K_MEANS_DIM, 96 sizeof(centroids[0]) * AV1_K_MEANS_DIM); 97 } else { 98 for (j = 0; j < AV1_K_MEANS_DIM; ++j) { 99 centroids[i * AV1_K_MEANS_DIM + j] = 100 DIVIDE_AND_ROUND(centroids_sum[i * AV1_K_MEANS_DIM + j], count[i]); 101 } 102 } 103 } 104 } 105 106 void RENAME(av1_k_means)(const int16_t *data, int16_t *centroids, 107 uint8_t *indices, int n, int k, int max_itr) { 108 int16_t centroids_tmp[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE]; 109 uint8_t indices_tmp[MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT]; 110 int16_t *meta_centroids[2] = { centroids, centroids_tmp }; 111 uint8_t *meta_indices[2] = { indices, indices_tmp }; 112 int i, l = 0, prev_l, best_l = 0; 113 int64_t this_dist; 114 115 assert(n <= MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT); 116 117 #if AV1_K_MEANS_DIM == 1 118 av1_calc_indices_dim1(data, centroids, indices, &this_dist, n, k); 119 #else 120 av1_calc_indices_dim2(data, centroids, indices, &this_dist, n, k); 121 #endif 122 123 for (i = 0; i < max_itr; ++i) { 124 const int64_t prev_dist = this_dist; 125 prev_l = l; 126 l = (l == 1) ? 0 : 1; 127 128 RENAME(calc_centroids)(data, meta_centroids[l], meta_indices[prev_l], n, k); 129 if (!memcmp(meta_centroids[l], meta_centroids[prev_l], 130 sizeof(centroids[0]) * k * AV1_K_MEANS_DIM)) { 131 break; 132 } 133 #if AV1_K_MEANS_DIM == 1 134 av1_calc_indices_dim1(data, meta_centroids[l], meta_indices[l], &this_dist, 135 n, k); 136 #else 137 av1_calc_indices_dim2(data, meta_centroids[l], meta_indices[l], &this_dist, 138 n, k); 139 #endif 140 141 if (this_dist > prev_dist) { 142 best_l = prev_l; 143 break; 144 } 145 } 146 if (i == max_itr) best_l = l; 147 if (best_l != 0) { 148 memcpy(centroids, meta_centroids[1], 149 sizeof(centroids[0]) * k * AV1_K_MEANS_DIM); 150 memcpy(indices, meta_indices[1], sizeof(indices[0]) * n); 151 } 152 } 153 #undef RENAME_ 154 #undef RENAME 155 #undef K_MEANS_RENAME_C 156 #undef RENAME_C_ 157 #undef RENAME_C