Line data    Source code

       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             : #ifndef AV1_K_MEANS_DIM
      17             : #error "This template requires AV1_K_MEANS_DIM to be defined"
      18             : #endif
      19             : 
      20             : #define RENAME_(x, y) AV1_K_MEANS_RENAME(x, y)
      21             : #define RENAME(x) RENAME_(x, AV1_K_MEANS_DIM)
      22             : 
      23           0 : static int RENAME(calc_dist)(const int *p1, const int *p2) {
      24           0 :   int dist = 0;
      25           0 :   for (int i = 0; i < AV1_K_MEANS_DIM; ++i) {
      26           0 :     const int diff = p1[i] - p2[i];
      27           0 :     dist += diff * diff;
      28             :   }
      29           0 :   return dist;
      30             : }
      31             : 
      32           0 : void RENAME(av1_calc_indices)(const int *data, const int *centroids,
      33             :                               uint8_t *indices, int n, int k) {
      34           0 :   for (int i = 0; i < n; ++i) {
      35           0 :     int min_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM, centroids);
      36           0 :     indices[i] = 0;
      37           0 :     for (int j = 1; j < k; ++j) {
      38           0 :       const int this_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM,
      39           0 :                                               centroids + j * AV1_K_MEANS_DIM);
      40           0 :       if (this_dist < min_dist) {
      41           0 :         min_dist = this_dist;
      42           0 :         indices[i] = j;
      43             :       }
      44             :     }
      45             :   }
      46           0 : }
      47             : 
      48           0 : static void RENAME(calc_centroids)(const int *data, int *centroids,
      49             :                                    const uint8_t *indices, int n, int k) {
      50             :   int i, j;
      51           0 :   int count[PALETTE_MAX_SIZE] = { 0 };
      52           0 :   unsigned int rand_state = (unsigned int)data[0];
      53           0 :   assert(n <= 32768);
      54           0 :   memset(centroids, 0, sizeof(centroids[0]) * k * AV1_K_MEANS_DIM);
      55             : 
      56           0 :   for (i = 0; i < n; ++i) {
      57           0 :     const int index = indices[i];
      58           0 :     assert(index < k);
      59           0 :     ++count[index];
      60           0 :     for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
      61           0 :       centroids[index * AV1_K_MEANS_DIM + j] += data[i * AV1_K_MEANS_DIM + j];
      62             :     }
      63             :   }
      64             : 
      65           0 :   for (i = 0; i < k; ++i) {
      66           0 :     if (count[i] == 0) {
      67           0 :       memcpy(centroids + i * AV1_K_MEANS_DIM,
      68           0 :              data + (lcg_rand16(&rand_state) % n) * AV1_K_MEANS_DIM,
      69             :              sizeof(centroids[0]) * AV1_K_MEANS_DIM);
      70             :     } else {
      71           0 :       for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
      72           0 :         centroids[i * AV1_K_MEANS_DIM + j] =
      73           0 :             DIVIDE_AND_ROUND(centroids[i * AV1_K_MEANS_DIM + j], count[i]);
      74             :       }
      75             :     }
      76             :   }
      77           0 : }
      78             : 
      79           0 : static int64_t RENAME(calc_total_dist)(const int *data, const int *centroids,
      80             :                                        const uint8_t *indices, int n, int k) {
      81           0 :   int64_t dist = 0;
      82             :   (void)k;
      83           0 :   for (int i = 0; i < n; ++i) {
      84           0 :     dist += RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM,
      85           0 :                               centroids + indices[i] * AV1_K_MEANS_DIM);
      86             :   }
      87           0 :   return dist;
      88             : }
      89             : 
      90           0 : void RENAME(av1_k_means)(const int *data, int *centroids, uint8_t *indices,
      91             :                          int n, int k, int max_itr) {
      92             :   int pre_centroids[2 * PALETTE_MAX_SIZE];
      93             :   uint8_t pre_indices[MAX_SB_SQUARE];
      94             : 
      95           0 :   RENAME(av1_calc_indices)(data, centroids, indices, n, k);
      96           0 :   int64_t this_dist = RENAME(calc_total_dist)(data, centroids, indices, n, k);
      97             : 
      98           0 :   for (int i = 0; i < max_itr; ++i) {
      99           0 :     const int64_t pre_dist = this_dist;
     100           0 :     memcpy(pre_centroids, centroids,
     101             :            sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM);
     102           0 :     memcpy(pre_indices, indices, sizeof(pre_indices[0]) * n);
     103             : 
     104           0 :     RENAME(calc_centroids)(data, centroids, indices, n, k);
     105           0 :     RENAME(av1_calc_indices)(data, centroids, indices, n, k);
     106           0 :     this_dist = RENAME(calc_total_dist)(data, centroids, indices, n, k);
     107             : 
     108           0 :     if (this_dist > pre_dist) {
     109           0 :       memcpy(centroids, pre_centroids,
     110             :              sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM);
     111           0 :       memcpy(indices, pre_indices, sizeof(pre_indices[0]) * n);
     112           0 :       break;
     113             :     }
     114           0 :     if (!memcmp(centroids, pre_centroids,
     115             :                 sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM))
     116           0 :       break;
     117             :   }
     118           0 : }
     119             : #undef RENAME_
     120             : #undef RENAME