kmeansColour.swift

func kMeans(image: CGImage, limit: Int = 10, iterations: Int = 5, perceptual: Bool = true) -> [CGColor] {
    let ciImage = CIImage(cgImage: image)
    let space = image.colorSpace!
    // kMeans filter returns the dominant colours within an image
    let filter = CIFilter.kMeans()
    filter.inputImage = ciImage
    filter.count = limit
    filter.passes = Float(iterations) // why is this a float lol
    filter.perceptual = perceptual
    let ciResult = filter.outputImage!
    // use a CIContext to render these to a CGImage, which can then be turned into bytes
    let context = CIContext()
    guard
        let cgResult = context.createCGImage(
            ciResult, from: ciResult.extent, format: .RGBAf,
            colorSpace: space),
        let data = cgResult.dataProvider?.data,
        let bytes = CFDataGetBytePtr(data)
    else {
        return []
    }

    // walk through pixels and convert them to colours
    let bytesPerPixel = cgResult.bitsPerPixel / 8
    var colors: [CGColor] = []


    // we only have 1 row of pixels so we need not walk through y
    for x in 0..<cgResult.width {
        let offset = (x + 0) * bytesPerPixel
        let raw = UnsafeRawPointer(bytes)
        // Jesus Christ
        let buffer = UnsafeBufferPointer(start: raw.advanced(by: offset).assumingMemoryBound(to: Float.self),
                                         count: 4)
        let cgFloats = buffer.map { CGFloat($0) }

        // alpha in these results denotes the weight of the cluster, so if we want to use these colours, it must be set to 1
        let color = CGColor(colorSpace: space,
                            components: cgFloats)!.copy(alpha: 1.0)!

        // avoid adding dupes
        if !colors.contains(color) {
            colors.append(color)
        }
    }

    return colors
}