Calculate the most dominant color in an icon from Chromium

// RGBA KMean Constants
const uint kNumberOfClusters = 4;
const int kNumberOfIterations = 50;
const uint kMaxBrightness = 665;
const uint kMinDarkness = 100;

private static uint calls_ = 0;

// Support class to hold information about each cluster of pixel data in
// the KMean algorithm. While this class does not contain all of the points
// that exist in the cluster, it keeps track of the aggregate sum so it can
// compute the new center appropriately.
internal class KMeanCluster {

  private uint[] centroid = new uint[3]; //[3];

  // Holds the sum of all the points that make up this cluster. Used to
  // generate the next centroid as well as to check for convergence.
  private uint[] aggregate = new uint[3]; //[3];
  private uint counter;

  // The weight of the cluster, determined by how many points were used
  // to generate the previous centroid.
  private uint weight;


  public KMeanCluster() {
    Reset();
  }

  public void Reset() {
    centroid[0] = centroid[1] = centroid[2] = 0;
    aggregate[0] = aggregate[1] = aggregate[2] = 0;
    counter = 0;
    weight = 0;
  }

  public void SetCentroid(uint r, uint g, uint b) {
    centroid[0] = r;
    centroid[1] = g;
    centroid[2] = b;
  }

  public void GetCentroid(ref uint r, ref uint g, ref uint b) {
    r = centroid[0];
    g = centroid[1];
    b = centroid[2];
  }

  public bool IsAtCentroid(uint r, uint g, uint b) {
    return r == centroid[0] && g == centroid[1] && b == centroid[2];
  }

  // Recomputes the centroid of the cluster based on the aggregate data. The
  // number of points used to calculate this center is stored for weighting
  // purposes. The aggregate and counter are then cleared to be ready for the
  // next iteration.
  public void RecomputeCentroid() {
    if (counter > 0) {
      centroid[0] = aggregate[0] / counter;
      centroid[1] = aggregate[1] / counter;
      centroid[2] = aggregate[2] / counter;

      aggregate[0] = aggregate[1] = aggregate[2] = 0;
      weight = counter;
      counter = 0;
    }
  }

  public void AddPoint(uint r, uint g, uint b) {
    aggregate[0] += r;
    aggregate[1] += g;
    aggregate[2] += b;
    ++counter;
  }

  // Just returns the distance^2. Since we are comparing relative distances
  // there is no need to perform the expensive sqrt() operation.
  public uint GetDistanceSqr(uint r, uint g, uint b) {
    return (r - centroid[0]) * (r - centroid[0]) +
           (g - centroid[1]) * (g - centroid[1]) +
           (b - centroid[2]) * (b - centroid[2]);
  }

  // In order to determine if we have hit convergence or not we need to see
  // if the centroid of the cluster has moved. This determines whether or
  // not the centroid is the same as the aggregate sum of points that will be
  // used to generate the next centroid.
  public bool CompareCentroidWithAggregate() {
    if (counter == 0)
      return false;

    return aggregate[0] / counter == centroid[0] &&
           aggregate[1] / counter == centroid[1] &&
           aggregate[2] / counter == centroid[2];
  }

  // Returns the previous counter, which is used to determine the weight
  // of the cluster for sorting.
  public uint GetWeight()  {
    return weight;
  }

  static bool SortKMeanClusterByWeight(KMeanCluster a, KMeanCluster b) {
    return a.GetWeight() > b.GetWeight();
  }

};


private static int GridSampler_GetSample(int width, int height) {
  // Hand-drawn bitmaps often have special outlines or feathering at the edges.
  // Start our sampling inset from the top and left edges. For example, a 10x10
  // image with 4 clusters would be sampled like this:
  // ..........
  // .0.4.8....
  // ..........
  // .1.5.9....
  // ..........
  // .2.6......
  // ..........
  // .3.7......
  // ..........
  const int kPadX = 1;
  const int kPadY = 1;
  int x = kPadX + (calls_ / kNumberOfClusters) * ((width - 2 * kPadX) / kNumberOfClusters);
  int y = kPadY + (calls_ % kNumberOfClusters) * ((height - 2 * kPadY) / kNumberOfClusters);
  int index = x + (y * width);
  ++calls_;
  return index % (width * height);
}

// Returns the color in an ARGB |image| that is closest in RGB-space to the
// provided |color|. Exported for testing.
public static Color FindClosestColor(LockedBitmap image,int width, int height, Color color) {
  uint in_r = color.R;
  uint in_g = color.G;
  uint in_b = color.B;

  // Search using distance-squared to avoid expensive sqrt() operations.
  int best_distance_squared = kint32max;
  Color best_color = color;
  const uint* byte = image;
  for (int i = 0; i < width * height; ++i) {
    uint b = *(byte++);
    uint g = *(byte++);
    uint r = *(byte++);
    uint a = *(byte++);

    // Ignore fully transparent pixels.
    if (a == 0)
      continue;

    int distance_squared = (in_b - b) * (in_b - b) + (in_g - g) * (in_g - g) +(in_r - r) * (in_r - r);
    if (distance_squared < best_distance_squared) {
      best_distance_squared = distance_squared;
      best_color = new Color(r, g, b);
    }
  }
  return best_color;
}

// Returns an Color that represents the calculated dominant color in the png.
// This uses a KMean clustering algorithm to find clusters of pixel colors in
// RGB space.
// |png| represents the data of a png encoded image.
// |darkness_limit| represents the minimum sum of the RGB components that is
// acceptable as a color choice. This can be from 0 to 765.
// |brightness_limit| represents the maximum sum of the RGB components that is
// acceptable as a color choice. This can be from 0 to 765.
//
// RGB KMean Algorithm (N clusters, M iterations):
// 1.Pick N starting colors by randomly sampling the pixels. If you see a
//   color you already saw keep sampling. After a certain number of tries
//   just remove the cluster and continue with N = N-1 clusters (for an image
//   with just one color this should devolve to N=1). These colors are the
//   centers of your N clusters.
// 2.For each pixel in the image find the cluster that it is closest to in RGB
//   space. Add that pixel's color to that cluster (we keep a sum and a count
//   of all of the pixels added to the space, so just add it to the sum and
//   increment count).
// 3.Calculate the new cluster centroids by getting the average color of all of
//   the pixels in each cluster (dividing the sum by the count).
// 4.See if the new centroids are the same as the old centroids.
//     a) If this is the case for all N clusters than we have converged and
//        can move on.
//     b) If any centroid moved, repeat step 2 with the new centroids for up
//        to M iterations.
// 5.Once the clusters have converged or M iterations have been tried, sort
//   the clusters by weight (where weight is the number of pixels that make up
//   this cluster).
// 6.Going through the sorted list of clusters, pick the first cluster with the
//   largest weight that's centroid fulfills the equation
//   |darkness_limit| < SUM(R, G, B) < |brightness_limit|. Return that color.
//   If no color fulfills that requirement return the color with the largest
//   weight regardless of whether or not it fulfills the equation above.
//
// Note: Switching to HSV space did not improve the results of this algorithm
// for typical favicon images.

// For a 16x16 icon on an Intel Core i5 this function takes approximately
// 0.5 ms to run.
// TODO(port): This code assumes the CPU architecture is little-endian.
public static Color CalculateKMeanColor(LockedBitmap decoded_data,int img_width,int img_height,uint darkness_limit,uint brightness_limit) {
  Color color = Color.White;
  calls_ = 0;
  if (img_width > 0 && img_height > 0) {
    std::vector<KMeanCluster> clusters;
    clusters.resize(kNumberOfClusters, KMeanCluster());

    // Pick a starting point for each cluster
    std::vector<KMeanCluster>::iterator cluster = clusters.begin();
    while (cluster != clusters.end()) {
      // Try up to 10 times to find a unique color. If no unique color can be
      // found, destroy this cluster.
      bool color_unique = false;
      for (int i = 0; i < 10; ++i) {
        int pixel_pos = GridSampler_GetSample(img_width, img_height) % (img_width * img_height);

        uint b = decoded_data[pixel_pos * 4];
        uint g = decoded_data[pixel_pos * 4 + 1];
        uint r = decoded_data[pixel_pos * 4 + 2];
        uint a = decoded_data[pixel_pos * 4 + 3];
        // Skip fully transparent pixels as they usually contain black in their
        // RGB channels but do not contribute to the visual image.
        if (a == 0)
          continue;

        // Loop through the previous clusters and check to see if we have seen
        // this color before.
        color_unique = true;
        for (std::vector<KMeanCluster>::iterator; cluster_check = clusters.begin(); cluster_check != cluster; ++cluster_check) {
          if (cluster_check->IsAtCentroid(r, g, b)) {
            color_unique = false;
            break;
          }
        }

        // If we have a unique color set the center of the cluster to
        // that color.
        if (color_unique) {
          cluster->SetCentroid(r, g, b);
          break;
        }
      }

      // If we don't have a unique color erase this cluster.
      if (!color_unique) {
        cluster = clusters.erase(cluster);
      } else {
        // Have to increment the iterator here, otherwise the increment in the
        // for loop will skip a cluster due to the erase if the color wasn't
        // unique.
        ++cluster;
      }
    }

    // If all pixels in the image are transparent we will have no clusters.
    if (clusters.empty())
      return color;

    bool convergence = false;
    for (int iteration = 0;iteration < kNumberOfIterations && !convergence;++iteration) {

      // Loop through each pixel so we can place it in the appropriate cluster.
      uint* pixel = decoded_data;
      uint* decoded_data_end = decoded_data + (img_width * img_height * 4);
      while (pixel < decoded_data_end) {
        uint b = *(pixel++);
        uint g = *(pixel++);
        uint r = *(pixel++);
        uint a = *(pixel++);
        // Skip transparent pixels, see above.
        if (a == 0)
          continue;

        uint distance_sqr_to_closest_cluster = UINT_MAX;
        std::vector<KMeanCluster>::iterator closest_cluster = clusters.begin();

        // Figure out which cluster this color is closest to in RGB space.
        for (std::vector<KMeanCluster>::iterator cluster = clusters.begin();
            cluster != clusters.end(); ++cluster) {
          uint distance_sqr = cluster->GetDistanceSqr(r, g, b);

          if (distance_sqr < distance_sqr_to_closest_cluster) {
            distance_sqr_to_closest_cluster = distance_sqr;
            closest_cluster = cluster;
          }
        }

        closest_cluster->AddPoint(r, g, b);
      }

      // Calculate the new cluster centers and see if we've converged or not.
      convergence = true;
      for (std::vector<KMeanCluster>::iterator cluster = clusters.begin();
          cluster != clusters.end(); ++cluster) {
        convergence &= cluster->CompareCentroidWithAggregate();

        cluster->RecomputeCentroid();
      }
    }

    // Sort the clusters by population so we can tell what the most popular
    // color is.
    std::sort(clusters.begin(), clusters.end(), KMeanCluster::SortKMeanClusterByWeight);

    // Loop through the clusters to figure out which cluster has an appropriate
    // color. Skip any that are too bright/dark and go in order of weight.
    for (std::vector<KMeanCluster>::iterator cluster = clusters.begin(); cluster != clusters.end(); ++cluster) {
      uint r, g, b;
      cluster->GetCentroid(&r, &g, &b);
      // Sum the RGB components to determine if the color is too bright or too
      // dark.
      // TODO (dtrainor): Look into using HSV here instead. This approximation
      // might be fine though.
      uint summed_color = r + g + b;

      if (summed_color < brightness_limit && summed_color > darkness_limit) {
        // If we found a valid color just set it and break. We don't want to
        // check the other ones.
        color = SkColorSetARGB(0xFF, r, g, b);
        break;
      } else if (cluster == clusters.begin()) {
        // We haven't found a valid color, but we are at the first color so
        // set the color anyway to make sure we at least have a value here.
        color = SkColorSetARGB(0xFF, r, g, b);
      }
    }
  }

  // Find a color that actually appears in the image (the K-mean cluster center
  // will not usually be a color that appears in the image).
  return FindClosestColor(decoded_data, img_width, img_height, color);
}

// Compute color covariance matrix for the input bitmap.
public static gfx::Matrix3F ComputeColorCovariance(LockedBitmap bitmap) {
  // First need basic stats to normalize each channel separately.
  SkAutoLockPixels bitmap_lock(bitmap);
  gfx::Matrix3F covariance = gfx::Matrix3F::Zeros();
  if (!bitmap.getPixels())
    return covariance;

  // Assume ARGB_8888 format.
  //DCHECK(bitmap.config() == SkBitmap::kARGB_8888_Config);

  long r_sum = 0;
  long g_sum = 0;
  long b_sum = 0;
  long rr_sum = 0;
  long gg_sum = 0;
  long bb_sum = 0;
  long rg_sum = 0;
  long rb_sum = 0;
  long gb_sum = 0;

  for (int y = 0; y < bitmap.height(); ++y) {
    SkPMColor* current_color = static_cast<uint*>(bitmap.getAddr32(0, y));
    for (int x = 0; x < bitmap.width(); ++x, ++current_color) {
      Color c = SkUnPreMultiply::PMColorToColor(*current_color);
      Color r = SkColorGetR(c);
      Color g = SkColorGetG(c);
      Color b = SkColorGetB(c);

      r_sum += r;
      g_sum += g;
      b_sum += b;
      rr_sum += r * r;
      gg_sum += g * g;
      bb_sum += b * b;
      rg_sum += r * g;
      rb_sum += r * b;
      gb_sum += g * b;
    }
  }

  // Covariance (not normalized) is E(X*X.t) - m * m.t and this is how it
  // is calculated below.
  // Each row below represents a row of the matrix describing (co)variances
  // of R, G and B channels with (R, G, B)
  int pixel_n = bitmap.width() * bitmap.height();
  covariance.set(
      (static_cast<double>(rr_sum) / pixel_n -
       static_cast<double>(r_sum * r_sum) / pixel_n / pixel_n),
      (static_cast<double>(rg_sum) / pixel_n -
       static_cast<double>(r_sum * g_sum) / pixel_n / pixel_n),
      (static_cast<double>(rb_sum) / pixel_n -
       static_cast<double>(r_sum * b_sum) / pixel_n / pixel_n),
      (static_cast<double>(rg_sum) / pixel_n -
       static_cast<double>(r_sum * g_sum) / pixel_n / pixel_n),
      (static_cast<double>(gg_sum) / pixel_n -
       static_cast<double>(g_sum * g_sum) / pixel_n / pixel_n),
      (static_cast<double>(gb_sum) / pixel_n -
       static_cast<double>(g_sum * b_sum) / pixel_n / pixel_n),
      (static_cast<double>(rb_sum) / pixel_n -
       static_cast<double>(r_sum * b_sum) / pixel_n / pixel_n),
      (static_cast<double>(gb_sum) / pixel_n -
       static_cast<double>(g_sum * b_sum) / pixel_n / pixel_n),
      (static_cast<double>(bb_sum) / pixel_n -
       static_cast<double>(b_sum * b_sum) / pixel_n / pixel_n));
  return covariance;
}

// Apply a color reduction transform defined by |color_transform| vector to
// |source_bitmap|. The result is put into |target_bitmap|, which is expected
// to be initialized to the required size and type (SkBitmap::kA8_Config).
// If |fit_to_range|, result is transfored linearly to fit 0-0xFF range.
// Otherwise, data is clipped.
// Returns true if the target has been computed.
public static bool ApplyColorReduction(LockedBitmap source_bitmap, const gfx::Vector3dF& color_transform,bool fit_to_range, LockedBitmap target_bitmap) {
  //DCHECK(target_bitmap);
  SkAutoLockPixels source_lock(source_bitmap);
  SkAutoLockPixels target_lock(*target_bitmap);

  //DCHECK(source_bitmap.getPixels());
  //DCHECK(target_bitmap->getPixels());
  //DCHECK_EQ(SkBitmap::kARGB_8888_Config, source_bitmap.config());
  //DCHECK_EQ(SkBitmap::kA8_Config, target_bitmap->config());
  //DCHECK_EQ(source_bitmap.height(), target_bitmap->height());
  //DCHECK_EQ(source_bitmap.width(), target_bitmap->width());
  //DCHECK(!source_bitmap.empty());

  // Elements of color_transform are explicitly off-loaded to local values for
  // efficiency reasons. Note that in practice images may correspond to entire
  // tab captures.
  float t0 = 0.0;
  float tr = color_transform.x();
  float tg = color_transform.y();
  float tb = color_transform.z();

  if (fit_to_range) {
    // We will figure out min/max in a preprocessing step and adjust
    // actual_transform as required.
    float max_val = std::numeric_limits<float>::min();
    float min_val = std::numeric_limits<float>::max();
    for (int y = 0; y < source_bitmap.height(); ++y) {
      const SkPMColor* source_color_row = static_cast<SkPMColor*>(
          source_bitmap.getAddr32(0, y));
      for (int x = 0; x < source_bitmap.width(); ++x) {
        Color c = SkUnPreMultiply::PMColorToColor(source_color_row[x]);
        float r = SkColorGetR(c);
        float g = SkColorGetG(c);
        float b = SkColorGetB(c);
        float gray_level = tr * r + tg * g + tb * b;
        max_val = std::max(max_val, gray_level);
        min_val = std::min(min_val, gray_level);
      }
    }

    // Adjust the transform so that the result is scaling.
    float scale = 0.0;
    t0 = -min_val;
    if (max_val > min_val)
      scale = 255.0 / (max_val - min_val);
    t0 *= scale;
    tr *= scale;
    tg *= scale;
    tb *= scale;
  }

  for (int y = 0; y < source_bitmap.height(); ++y) {
    const SkPMColor* source_color_row = static_cast<SkPMColor*>(
        source_bitmap.getAddr32(0, y));
    uint* target_color_row = target_bitmap->getAddr8(0, y);
    for (int x = 0; x < source_bitmap.width(); ++x) {
      Color c = SkUnPreMultiply::PMColorToColor(source_color_row[x]);
      float r = SkColorGetR(c);
      float g = SkColorGetG(c);
      float b = SkColorGetB(c);

      float gl = t0 + tr * r + tg * g + tb * b;
      if (gl < 0)
        gl = 0;
      if (gl > 0xFF)
        gl = 0xFF;
      target_color_row[x] = static_cast<uint>(gl);
    }
  }

  return true;
}

// Compute a monochrome image representing the principal color component of
// the |source_bitmap|. The result is stored in |target_bitmap|, which must be
// initialized to the required size and type (SkBitmap::kA8_Config).
// Returns true if the conversion succeeded. Note that there might be legitimate
// reasons for the process to fail even if all input was correct. This is a
// condition the caller must be able to handle.
public static bool ComputePrincipalComponentImage(LockedBitmap source_bitmap, LockedBitmap target_bitmap) {

  gfx::Matrix3F covariance = ComputeColorCovariance(source_bitmap);
  gfx::Matrix3F eigenvectors = gfx::Matrix3F::Zeros();
  gfx::Vector3dF eigenvals = covariance.SolveEigenproblem(&eigenvectors);
  gfx::Vector3dF principal = eigenvectors.get_column(0);
  if (eigenvals == gfx::Vector3dF() || principal == gfx::Vector3dF())
    return false;  // This may happen for some edge cases.
  return ApplyColorReduction(source_bitmap, principal, true, target_bitmap);
}



// Un-premultiplies each pixel in |bitmap| into an output |buffer|. Requires
// approximately 10 microseconds for a 16x16 icon on an Intel Core i5.
/*public static void UnPreMultiply(SkBitmap& bitmap, uint* buffer, int buffer_size) {
  SkAutoLockPixels auto_lock(bitmap);
  uint* in = static_cast<uint*>(bitmap.getPixels());
  uint* out = buffer;
  int pixel_count = std::min(bitmap.width() * bitmap.height(), buffer_size);
  for (int i = 0; i < pixel_count; ++i){
    *out++ = SkUnPreMultiply::PMColorToColor(*in++);
    }
}*/