1// Copyright (c) 2012 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include "ui/gfx/color_analysis.h"
6
7#include <algorithm>
8#include <limits>
9#include <vector>
10
11#include "base/logging.h"
12#include "base/memory/scoped_ptr.h"
13#include "third_party/skia/include/core/SkBitmap.h"
14#include "third_party/skia/include/core/SkColorPriv.h"
15#include "third_party/skia/include/core/SkUnPreMultiply.h"
16#include "ui/gfx/codec/png_codec.h"
17#include "ui/gfx/color_utils.h"
18
19namespace color_utils {
20namespace {
21
22// RGBA KMean Constants
23const uint32_t kNumberOfClusters = 4;
24const int kNumberOfIterations = 50;
25
26const HSL kDefaultLowerHSLBound = {-1, -1, 0.15};
27const HSL kDefaultUpperHSLBound = {-1, -1, 0.85};
28
29// Background Color Modification Constants
30const SkColor kDefaultBgColor = SK_ColorWHITE;
31
32// Support class to hold information about each cluster of pixel data in
33// the KMean algorithm. While this class does not contain all of the points
34// that exist in the cluster, it keeps track of the aggregate sum so it can
35// compute the new center appropriately.
36class KMeanCluster {
37 public:
38  KMeanCluster() {
39    Reset();
40  }
41
42  void Reset() {
43    centroid[0] = centroid[1] = centroid[2] = 0;
44    aggregate[0] = aggregate[1] = aggregate[2] = 0;
45    counter = 0;
46    weight = 0;
47  }
48
49  inline void SetCentroid(uint8_t r, uint8_t g, uint8_t b) {
50    centroid[0] = r;
51    centroid[1] = g;
52    centroid[2] = b;
53  }
54
55  inline void GetCentroid(uint8_t* r, uint8_t* g, uint8_t* b) {
56    *r = centroid[0];
57    *g = centroid[1];
58    *b = centroid[2];
59  }
60
61  inline bool IsAtCentroid(uint8_t r, uint8_t g, uint8_t b) {
62    return r == centroid[0] && g == centroid[1] && b == centroid[2];
63  }
64
65  // Recomputes the centroid of the cluster based on the aggregate data. The
66  // number of points used to calculate this center is stored for weighting
67  // purposes. The aggregate and counter are then cleared to be ready for the
68  // next iteration.
69  inline void RecomputeCentroid() {
70    if (counter > 0) {
71      centroid[0] = aggregate[0] / counter;
72      centroid[1] = aggregate[1] / counter;
73      centroid[2] = aggregate[2] / counter;
74
75      aggregate[0] = aggregate[1] = aggregate[2] = 0;
76      weight = counter;
77      counter = 0;
78    }
79  }
80
81  inline void AddPoint(uint8_t r, uint8_t g, uint8_t b) {
82    aggregate[0] += r;
83    aggregate[1] += g;
84    aggregate[2] += b;
85    ++counter;
86  }
87
88  // Just returns the distance^2. Since we are comparing relative distances
89  // there is no need to perform the expensive sqrt() operation.
90  inline uint32_t GetDistanceSqr(uint8_t r, uint8_t g, uint8_t b) {
91    return (r - centroid[0]) * (r - centroid[0]) +
92           (g - centroid[1]) * (g - centroid[1]) +
93           (b - centroid[2]) * (b - centroid[2]);
94  }
95
96  // In order to determine if we have hit convergence or not we need to see
97  // if the centroid of the cluster has moved. This determines whether or
98  // not the centroid is the same as the aggregate sum of points that will be
99  // used to generate the next centroid.
100  inline bool CompareCentroidWithAggregate() {
101    if (counter == 0)
102      return false;
103
104    return aggregate[0] / counter == centroid[0] &&
105           aggregate[1] / counter == centroid[1] &&
106           aggregate[2] / counter == centroid[2];
107  }
108
109  // Returns the previous counter, which is used to determine the weight
110  // of the cluster for sorting.
111  inline uint32_t GetWeight() const {
112    return weight;
113  }
114
115  static bool SortKMeanClusterByWeight(const KMeanCluster& a,
116                                       const KMeanCluster& b) {
117    return a.GetWeight() > b.GetWeight();
118  }
119
120 private:
121  uint8_t centroid[3];
122
123  // Holds the sum of all the points that make up this cluster. Used to
124  // generate the next centroid as well as to check for convergence.
125  uint32_t aggregate[3];
126  uint32_t counter;
127
128  // The weight of the cluster, determined by how many points were used
129  // to generate the previous centroid.
130  uint32_t weight;
131};
132
133// Un-premultiplies each pixel in |bitmap| into an output |buffer|.
134void UnPreMultiply(const SkBitmap& bitmap, uint32_t* buffer, int buffer_size) {
135  SkAutoLockPixels auto_lock(bitmap);
136  uint32_t* in = static_cast<uint32_t*>(bitmap.getPixels());
137  uint32_t* out = buffer;
138  int pixel_count = std::min(bitmap.width() * bitmap.height(), buffer_size);
139  for (int i = 0; i < pixel_count; ++i) {
140    int alpha = SkGetPackedA32(*in);
141    if (alpha != 0 && alpha != 255)
142      *out++ = SkUnPreMultiply::PMColorToColor(*in++);
143    else
144      *out++ = *in++;
145  }
146}
147
148} // namespace
149
150KMeanImageSampler::KMeanImageSampler() {
151}
152
153KMeanImageSampler::~KMeanImageSampler() {
154}
155
156GridSampler::GridSampler() : calls_(0) {
157}
158
159GridSampler::~GridSampler() {
160}
161
162int GridSampler::GetSample(int width, int height) {
163  // Hand-drawn bitmaps often have special outlines or feathering at the edges.
164  // Start our sampling inset from the top and left edges. For example, a 10x10
165  // image with 4 clusters would be sampled like this:
166  // ..........
167  // .0.4.8....
168  // ..........
169  // .1.5.9....
170  // ..........
171  // .2.6......
172  // ..........
173  // .3.7......
174  // ..........
175  const int kPadX = 1;
176  const int kPadY = 1;
177  int x = kPadX +
178      (calls_ / kNumberOfClusters) * ((width - 2 * kPadX) / kNumberOfClusters);
179  int y = kPadY +
180      (calls_ % kNumberOfClusters) * ((height - 2 * kPadY) / kNumberOfClusters);
181  int index = x + (y * width);
182  ++calls_;
183  return index % (width * height);
184}
185
186SkColor FindClosestColor(const uint8_t* image,
187                         int width,
188                         int height,
189                         SkColor color) {
190  uint8_t in_r = SkColorGetR(color);
191  uint8_t in_g = SkColorGetG(color);
192  uint8_t in_b = SkColorGetB(color);
193  // Search using distance-squared to avoid expensive sqrt() operations.
194  int best_distance_squared = kint32max;
195  SkColor best_color = color;
196  const uint8_t* byte = image;
197  for (int i = 0; i < width * height; ++i) {
198    uint8_t b = *(byte++);
199    uint8_t g = *(byte++);
200    uint8_t r = *(byte++);
201    uint8_t a = *(byte++);
202    // Ignore fully transparent pixels.
203    if (a == 0)
204      continue;
205    int distance_squared =
206        (in_b - b) * (in_b - b) +
207        (in_g - g) * (in_g - g) +
208        (in_r - r) * (in_r - r);
209    if (distance_squared < best_distance_squared) {
210      best_distance_squared = distance_squared;
211      best_color = SkColorSetRGB(r, g, b);
212    }
213  }
214  return best_color;
215}
216
217// For a 16x16 icon on an Intel Core i5 this function takes approximately
218// 0.5 ms to run.
219// TODO(port): This code assumes the CPU architecture is little-endian.
220SkColor CalculateKMeanColorOfBuffer(uint8_t* decoded_data,
221                                    int img_width,
222                                    int img_height,
223                                    const HSL& lower_bound,
224                                    const HSL& upper_bound,
225                                    KMeanImageSampler* sampler) {
226  SkColor color = kDefaultBgColor;
227  if (img_width > 0 && img_height > 0) {
228    std::vector<KMeanCluster> clusters;
229    clusters.resize(kNumberOfClusters, KMeanCluster());
230
231    // Pick a starting point for each cluster
232    std::vector<KMeanCluster>::iterator cluster = clusters.begin();
233    while (cluster != clusters.end()) {
234      // Try up to 10 times to find a unique color. If no unique color can be
235      // found, destroy this cluster.
236      bool color_unique = false;
237      for (int i = 0; i < 10; ++i) {
238        int pixel_pos = sampler->GetSample(img_width, img_height) %
239            (img_width * img_height);
240
241        uint8_t b = decoded_data[pixel_pos * 4];
242        uint8_t g = decoded_data[pixel_pos * 4 + 1];
243        uint8_t r = decoded_data[pixel_pos * 4 + 2];
244        uint8_t a = decoded_data[pixel_pos * 4 + 3];
245        // Skip fully transparent pixels as they usually contain black in their
246        // RGB channels but do not contribute to the visual image.
247        if (a == 0)
248          continue;
249
250        // Loop through the previous clusters and check to see if we have seen
251        // this color before.
252        color_unique = true;
253        for (std::vector<KMeanCluster>::iterator
254            cluster_check = clusters.begin();
255            cluster_check != cluster; ++cluster_check) {
256          if (cluster_check->IsAtCentroid(r, g, b)) {
257            color_unique = false;
258            break;
259          }
260        }
261
262        // If we have a unique color set the center of the cluster to
263        // that color.
264        if (color_unique) {
265          cluster->SetCentroid(r, g, b);
266          break;
267        }
268      }
269
270      // If we don't have a unique color erase this cluster.
271      if (!color_unique) {
272        cluster = clusters.erase(cluster);
273      } else {
274        // Have to increment the iterator here, otherwise the increment in the
275        // for loop will skip a cluster due to the erase if the color wasn't
276        // unique.
277        ++cluster;
278      }
279    }
280
281    // If all pixels in the image are transparent we will have no clusters.
282    if (clusters.empty())
283      return color;
284
285    bool convergence = false;
286    for (int iteration = 0;
287        iteration < kNumberOfIterations && !convergence;
288        ++iteration) {
289
290      // Loop through each pixel so we can place it in the appropriate cluster.
291      uint8_t* pixel = decoded_data;
292      uint8_t* decoded_data_end = decoded_data + (img_width * img_height * 4);
293      while (pixel < decoded_data_end) {
294        uint8_t b = *(pixel++);
295        uint8_t g = *(pixel++);
296        uint8_t r = *(pixel++);
297        uint8_t a = *(pixel++);
298        // Skip transparent pixels, see above.
299        if (a == 0)
300          continue;
301
302        uint32_t distance_sqr_to_closest_cluster = UINT_MAX;
303        std::vector<KMeanCluster>::iterator closest_cluster = clusters.begin();
304
305        // Figure out which cluster this color is closest to in RGB space.
306        for (std::vector<KMeanCluster>::iterator cluster = clusters.begin();
307            cluster != clusters.end(); ++cluster) {
308          uint32_t distance_sqr = cluster->GetDistanceSqr(r, g, b);
309
310          if (distance_sqr < distance_sqr_to_closest_cluster) {
311            distance_sqr_to_closest_cluster = distance_sqr;
312            closest_cluster = cluster;
313          }
314        }
315
316        closest_cluster->AddPoint(r, g, b);
317      }
318
319      // Calculate the new cluster centers and see if we've converged or not.
320      convergence = true;
321      for (std::vector<KMeanCluster>::iterator cluster = clusters.begin();
322          cluster != clusters.end(); ++cluster) {
323        convergence &= cluster->CompareCentroidWithAggregate();
324
325        cluster->RecomputeCentroid();
326      }
327    }
328
329    // Sort the clusters by population so we can tell what the most popular
330    // color is.
331    std::sort(clusters.begin(), clusters.end(),
332              KMeanCluster::SortKMeanClusterByWeight);
333
334    // Loop through the clusters to figure out which cluster has an appropriate
335    // color. Skip any that are too bright/dark and go in order of weight.
336    for (std::vector<KMeanCluster>::iterator cluster = clusters.begin();
337        cluster != clusters.end(); ++cluster) {
338      uint8_t r, g, b;
339      cluster->GetCentroid(&r, &g, &b);
340
341      SkColor current_color = SkColorSetARGB(SK_AlphaOPAQUE, r, g, b);
342      HSL hsl;
343      SkColorToHSL(current_color, &hsl);
344      if (IsWithinHSLRange(hsl, lower_bound, upper_bound)) {
345        // If we found a valid color just set it and break. We don't want to
346        // check the other ones.
347        color = current_color;
348        break;
349      } else if (cluster == clusters.begin()) {
350        // We haven't found a valid color, but we are at the first color so
351        // set the color anyway to make sure we at least have a value here.
352        color = current_color;
353      }
354    }
355  }
356
357  // Find a color that actually appears in the image (the K-mean cluster center
358  // will not usually be a color that appears in the image).
359  return FindClosestColor(decoded_data, img_width, img_height, color);
360}
361
362SkColor CalculateKMeanColorOfPNG(scoped_refptr<base::RefCountedMemory> png,
363                                 const HSL& lower_bound,
364                                 const HSL& upper_bound,
365                                 KMeanImageSampler* sampler) {
366  int img_width = 0;
367  int img_height = 0;
368  std::vector<uint8_t> decoded_data;
369  SkColor color = kDefaultBgColor;
370
371  if (png.get() && png->size() &&
372      gfx::PNGCodec::Decode(png->front(),
373                            png->size(),
374                            gfx::PNGCodec::FORMAT_BGRA,
375                            &decoded_data,
376                            &img_width,
377                            &img_height)) {
378    return CalculateKMeanColorOfBuffer(&decoded_data[0],
379                                       img_width,
380                                       img_height,
381                                       lower_bound,
382                                       upper_bound,
383                                       sampler);
384  }
385  return color;
386}
387
388SkColor CalculateKMeanColorOfPNG(scoped_refptr<base::RefCountedMemory> png) {
389  GridSampler sampler;
390  return CalculateKMeanColorOfPNG(
391      png, kDefaultLowerHSLBound, kDefaultUpperHSLBound, &sampler);
392}
393
394SkColor CalculateKMeanColorOfBitmap(const SkBitmap& bitmap,
395                                    const HSL& lower_bound,
396                                    const HSL& upper_bound,
397                                    KMeanImageSampler* sampler) {
398  // SkBitmap uses pre-multiplied alpha but the KMean clustering function
399  // above uses non-pre-multiplied alpha. Transform the bitmap before we
400  // analyze it because the function reads each pixel multiple times.
401  int pixel_count = bitmap.width() * bitmap.height();
402  scoped_ptr<uint32_t[]> image(new uint32_t[pixel_count]);
403  UnPreMultiply(bitmap, image.get(), pixel_count);
404
405  return CalculateKMeanColorOfBuffer(reinterpret_cast<uint8_t*>(image.get()),
406                                     bitmap.width(),
407                                     bitmap.height(),
408                                     lower_bound,
409                                     upper_bound,
410                                     sampler);
411}
412
413SkColor CalculateKMeanColorOfBitmap(const SkBitmap& bitmap) {
414  GridSampler sampler;
415  return CalculateKMeanColorOfBitmap(
416      bitmap, kDefaultLowerHSLBound, kDefaultUpperHSLBound, &sampler);
417}
418
419gfx::Matrix3F ComputeColorCovariance(const SkBitmap& bitmap) {
420  // First need basic stats to normalize each channel separately.
421  SkAutoLockPixels bitmap_lock(bitmap);
422  gfx::Matrix3F covariance = gfx::Matrix3F::Zeros();
423  if (!bitmap.getPixels())
424    return covariance;
425
426  // Assume ARGB_8888 format.
427  DCHECK(bitmap.colorType() == kN32_SkColorType);
428
429  int64_t r_sum = 0;
430  int64_t g_sum = 0;
431  int64_t b_sum = 0;
432  int64_t rr_sum = 0;
433  int64_t gg_sum = 0;
434  int64_t bb_sum = 0;
435  int64_t rg_sum = 0;
436  int64_t rb_sum = 0;
437  int64_t gb_sum = 0;
438
439  for (int y = 0; y < bitmap.height(); ++y) {
440    SkPMColor* current_color = static_cast<uint32_t*>(bitmap.getAddr32(0, y));
441    for (int x = 0; x < bitmap.width(); ++x, ++current_color) {
442      SkColor c;
443      int alpha = SkGetPackedA32(*current_color);
444      if (alpha != 0 && alpha != 255)
445        c = SkUnPreMultiply::PMColorToColor(*current_color);
446      else
447        c = *current_color;
448
449      SkColor r = SkColorGetR(c);
450      SkColor g = SkColorGetG(c);
451      SkColor b = SkColorGetB(c);
452
453      r_sum += r;
454      g_sum += g;
455      b_sum += b;
456      rr_sum += r * r;
457      gg_sum += g * g;
458      bb_sum += b * b;
459      rg_sum += r * g;
460      rb_sum += r * b;
461      gb_sum += g * b;
462    }
463  }
464
465  // Covariance (not normalized) is E(X*X.t) - m * m.t and this is how it
466  // is calculated below.
467  // Each row below represents a row of the matrix describing (co)variances
468  // of R, G and B channels with (R, G, B)
469  int pixel_n = bitmap.width() * bitmap.height();
470  covariance.set(
471      (static_cast<double>(rr_sum) / pixel_n -
472       static_cast<double>(r_sum * r_sum) / pixel_n / pixel_n),
473      (static_cast<double>(rg_sum) / pixel_n -
474       static_cast<double>(r_sum * g_sum) / pixel_n / pixel_n),
475      (static_cast<double>(rb_sum) / pixel_n -
476       static_cast<double>(r_sum * b_sum) / pixel_n / pixel_n),
477      (static_cast<double>(rg_sum) / pixel_n -
478       static_cast<double>(r_sum * g_sum) / pixel_n / pixel_n),
479      (static_cast<double>(gg_sum) / pixel_n -
480       static_cast<double>(g_sum * g_sum) / pixel_n / pixel_n),
481      (static_cast<double>(gb_sum) / pixel_n -
482       static_cast<double>(g_sum * b_sum) / pixel_n / pixel_n),
483      (static_cast<double>(rb_sum) / pixel_n -
484       static_cast<double>(r_sum * b_sum) / pixel_n / pixel_n),
485      (static_cast<double>(gb_sum) / pixel_n -
486       static_cast<double>(g_sum * b_sum) / pixel_n / pixel_n),
487      (static_cast<double>(bb_sum) / pixel_n -
488       static_cast<double>(b_sum * b_sum) / pixel_n / pixel_n));
489  return covariance;
490}
491
492bool ApplyColorReduction(const SkBitmap& source_bitmap,
493                         const gfx::Vector3dF& color_transform,
494                         bool fit_to_range,
495                         SkBitmap* target_bitmap) {
496  DCHECK(target_bitmap);
497  SkAutoLockPixels source_lock(source_bitmap);
498  SkAutoLockPixels target_lock(*target_bitmap);
499
500  DCHECK(source_bitmap.getPixels());
501  DCHECK(target_bitmap->getPixels());
502  DCHECK_EQ(kN32_SkColorType, source_bitmap.colorType());
503  DCHECK_EQ(kAlpha_8_SkColorType, target_bitmap->colorType());
504  DCHECK_EQ(source_bitmap.height(), target_bitmap->height());
505  DCHECK_EQ(source_bitmap.width(), target_bitmap->width());
506  DCHECK(!source_bitmap.empty());
507
508  // Elements of color_transform are explicitly off-loaded to local values for
509  // efficiency reasons. Note that in practice images may correspond to entire
510  // tab captures.
511  float t0 = 0.0;
512  float tr = color_transform.x();
513  float tg = color_transform.y();
514  float tb = color_transform.z();
515
516  if (fit_to_range) {
517    // We will figure out min/max in a preprocessing step and adjust
518    // actual_transform as required.
519    float max_val = std::numeric_limits<float>::min();
520    float min_val = std::numeric_limits<float>::max();
521    for (int y = 0; y < source_bitmap.height(); ++y) {
522      const SkPMColor* source_color_row = static_cast<SkPMColor*>(
523          source_bitmap.getAddr32(0, y));
524      for (int x = 0; x < source_bitmap.width(); ++x) {
525        SkColor c;
526        int alpha = SkGetPackedA32(source_color_row[x]);
527        if (alpha != 0 && alpha != 255)
528          c = SkUnPreMultiply::PMColorToColor(source_color_row[x]);
529        else
530          c = source_color_row[x];
531
532        float r = SkColorGetR(c);
533        float g = SkColorGetG(c);
534        float b = SkColorGetB(c);
535        float gray_level = tr * r + tg * g + tb * b;
536        max_val = std::max(max_val, gray_level);
537        min_val = std::min(min_val, gray_level);
538      }
539    }
540
541    // Adjust the transform so that the result is scaling.
542    float scale = 0.0;
543    t0 = -min_val;
544    if (max_val > min_val)
545      scale = 255.0 / (max_val - min_val);
546    t0 *= scale;
547    tr *= scale;
548    tg *= scale;
549    tb *= scale;
550  }
551
552  for (int y = 0; y < source_bitmap.height(); ++y) {
553    const SkPMColor* source_color_row = static_cast<SkPMColor*>(
554        source_bitmap.getAddr32(0, y));
555    uint8_t* target_color_row = target_bitmap->getAddr8(0, y);
556    for (int x = 0; x < source_bitmap.width(); ++x) {
557      SkColor c;
558      int alpha = SkGetPackedA32(source_color_row[x]);
559      if (alpha != 0 && alpha != 255)
560        c = SkUnPreMultiply::PMColorToColor(source_color_row[x]);
561      else
562        c = source_color_row[x];
563
564      float r = SkColorGetR(c);
565      float g = SkColorGetG(c);
566      float b = SkColorGetB(c);
567
568      float gl = t0 + tr * r + tg * g + tb * b;
569      if (gl < 0)
570        gl = 0;
571      if (gl > 0xFF)
572        gl = 0xFF;
573      target_color_row[x] = static_cast<uint8_t>(gl);
574    }
575  }
576
577  return true;
578}
579
580bool ComputePrincipalComponentImage(const SkBitmap& source_bitmap,
581                                    SkBitmap* target_bitmap) {
582  if (!target_bitmap) {
583    NOTREACHED();
584    return false;
585  }
586
587  gfx::Matrix3F covariance = ComputeColorCovariance(source_bitmap);
588  gfx::Matrix3F eigenvectors = gfx::Matrix3F::Zeros();
589  gfx::Vector3dF eigenvals = covariance.SolveEigenproblem(&eigenvectors);
590  gfx::Vector3dF principal = eigenvectors.get_column(0);
591  if (eigenvals == gfx::Vector3dF() || principal == gfx::Vector3dF())
592    return false;  // This may happen for some edge cases.
593  return ApplyColorReduction(source_bitmap, principal, true, target_bitmap);
594}
595
596}  // color_utils
597