soft_palettegen.go

Documentation: github.com/lucasb-eyer/go-colorful

     1  // Largely inspired by the descriptions in http://lab.medialab.sciences-po.fr/iwanthue/
     2  // but written from scratch.
     3  
     4  package colorful
     5  
     6  import (
     7  	"fmt"
     8  	"math"
     9  	"math/rand"
    10  )
    11  
    12  // The algorithm works in L*a*b* color space and converts to RGB in the end.
    13  // L* in [0..1], a* and b* in [-1..1]
    14  type lab_t struct {
    15  	L, A, B float64
    16  }
    17  
    18  type SoftPaletteSettings struct {
    19  	// A function which can be used to restrict the allowed color-space.
    20  	CheckColor func(l, a, b float64) bool
    21  
    22  	// The higher, the better quality but the slower. Usually two figures.
    23  	Iterations int
    24  
    25  	// Use up to 160000 or 8000 samples of the L*a*b* space (and thus calls to CheckColor).
    26  	// Set this to true only if your CheckColor shapes the Lab space weirdly.
    27  	ManySamples bool
    28  }
    29  
    30  // Yeah, windows-stype Foo, FooEx, screw you golang...
    31  // Uses K-means to cluster the color-space and return the means of the clusters
    32  // as a new palette of distinctive colors. Falls back to K-medoid if the mean
    33  // happens to fall outside of the color-space, which can only happen if you
    34  // specify a CheckColor function.
    35  func SoftPaletteEx(colorsCount int, settings SoftPaletteSettings) ([]Color, error) {
    36  
    37  	// Checks whether it's a valid RGB and also fulfills the potentially provided constraint.
    38  	check := func(col lab_t) bool {
    39  		c := Lab(col.L, col.A, col.B)
    40  		return c.IsValid() && (settings.CheckColor == nil || settings.CheckColor(col.L, col.A, col.B))
    41  	}
    42  
    43  	// Sample the color space. These will be the points k-means is run on.
    44  	dl := 0.05
    45  	dab := 0.1
    46  	if settings.ManySamples {
    47  		dl = 0.01
    48  		dab = 0.05
    49  	}
    50  
    51  	samples := make([]lab_t, 0, int(1.0/dl*2.0/dab*2.0/dab))
    52  	for l := 0.0; l <= 1.0; l += dl {
    53  		for a := -1.0; a <= 1.0; a += dab {
    54  			for b := -1.0; b <= 1.0; b += dab {
    55  				if check(lab_t{l, a, b}) {
    56  					samples = append(samples, lab_t{l, a, b})
    57  				}
    58  			}
    59  		}
    60  	}
    61  
    62  	// That would cause some infinite loops down there...
    63  	if len(samples) < colorsCount {
    64  		return nil, fmt.Errorf("palettegen: more colors requested (%v) than samples available (%v). Your requested color count may be wrong, you might want to use many samples or your constraint function makes the valid color space too small", colorsCount, len(samples))
    65  	} else if len(samples) == colorsCount {
    66  		return labs2cols(samples), nil // Oops?
    67  	}
    68  
    69  	// We take the initial means out of the samples, so they are in fact medoids.
    70  	// This helps us avoid infinite loops or arbitrary cutoffs with too restrictive constraints.
    71  	means := make([]lab_t, colorsCount)
    72  	for i := 0; i < colorsCount; i++ {
    73  		for means[i] = samples[rand.Intn(len(samples))]; in(means, i, means[i]); means[i] = samples[rand.Intn(len(samples))] {
    74  		}
    75  	}
    76  
    77  	clusters := make([]int, len(samples))
    78  	samples_used := make([]bool, len(samples))
    79  
    80  	// The actual k-means/medoid iterations
    81  	for i := 0; i < settings.Iterations; i++ {
    82  		// Reassing the samples to clusters, i.e. to their closest mean.
    83  		// By the way, also check if any sample is used as a medoid and if so, mark that.
    84  		for isample, sample := range samples {
    85  			samples_used[isample] = false
    86  			mindist := math.Inf(+1)
    87  			for imean, mean := range means {
    88  				dist := lab_dist(sample, mean)
    89  				if dist < mindist {
    90  					mindist = dist
    91  					clusters[isample] = imean
    92  				}
    93  
    94  				// Mark samples which are used as a medoid.
    95  				if lab_eq(sample, mean) {
    96  					samples_used[isample] = true
    97  				}
    98  			}
    99  		}
   100  
   101  		// Compute new means according to the samples.
   102  		for imean := range means {
   103  			// The new mean is the average of all samples belonging to it..
   104  			nsamples := 0
   105  			newmean := lab_t{0.0, 0.0, 0.0}
   106  			for isample, sample := range samples {
   107  				if clusters[isample] == imean {
   108  					nsamples++
   109  					newmean.L += sample.L
   110  					newmean.A += sample.A
   111  					newmean.B += sample.B
   112  				}
   113  			}
   114  			if nsamples > 0 {
   115  				newmean.L /= float64(nsamples)
   116  				newmean.A /= float64(nsamples)
   117  				newmean.B /= float64(nsamples)
   118  			} else {
   119  				// That mean doesn't have any samples? Get a new mean from the sample list!
   120  				var inewmean int
   121  				for inewmean = rand.Intn(len(samples_used)); samples_used[inewmean]; inewmean = rand.Intn(len(samples_used)) {
   122  				}
   123  				newmean = samples[inewmean]
   124  				samples_used[inewmean] = true
   125  			}
   126  
   127  			// But now we still need to check whether the new mean is an allowed color.
   128  			if nsamples > 0 && check(newmean) {
   129  				// It does, life's good (TM)
   130  				means[imean] = newmean
   131  			} else {
   132  				// New mean isn't an allowed color or doesn't have any samples!
   133  				// Switch to medoid mode and pick the closest (unused) sample.
   134  				// This should always find something thanks to len(samples) >= colorsCount
   135  				mindist := math.Inf(+1)
   136  				for isample, sample := range samples {
   137  					if !samples_used[isample] {
   138  						dist := lab_dist(sample, newmean)
   139  						if dist < mindist {
   140  							mindist = dist
   141  							newmean = sample
   142  						}
   143  					}
   144  				}
   145  			}
   146  		}
   147  	}
   148  	return labs2cols(means), nil
   149  }
   150  
   151  // A wrapper which uses common parameters.
   152  func SoftPalette(colorsCount int) ([]Color, error) {
   153  	return SoftPaletteEx(colorsCount, SoftPaletteSettings{nil, 50, false})
   154  }
   155  
   156  func in(haystack []lab_t, upto int, needle lab_t) bool {
   157  	for i := 0; i < upto && i < len(haystack); i++ {
   158  		if haystack[i] == needle {
   159  			return true
   160  		}
   161  	}
   162  	return false
   163  }
   164  
   165  const LAB_DELTA = 1e-6
   166  
   167  func lab_eq(lab1, lab2 lab_t) bool {
   168  	return math.Abs(lab1.L-lab2.L) < LAB_DELTA &&
   169  		math.Abs(lab1.A-lab2.A) < LAB_DELTA &&
   170  		math.Abs(lab1.B-lab2.B) < LAB_DELTA
   171  }
   172  
   173  // That's faster than using colorful's DistanceLab since we would have to
   174  // convert back and forth for that. Here is no conversion.
   175  func lab_dist(lab1, lab2 lab_t) float64 {
   176  	return math.Sqrt(sq(lab1.L-lab2.L) + sq(lab1.A-lab2.A) + sq(lab1.B-lab2.B))
   177  }
   178  
   179  func labs2cols(labs []lab_t) (cols []Color) {
   180  	cols = make([]Color, len(labs))
   181  	for k, v := range labs {
   182  		cols[k] = Lab(v.L, v.A, v.B)
   183  	}
   184  	return cols
   185  }
   186
View as plain text