Exploring Core Graphics: Extract Prominent and Unique Colors from UIImage

Ever since I saw Apple Music's lyrics page with beautiful background from the artwork image, I wanted a way to extract the colors in a similar way for my apps. I still have not gotten closer to the magic that Apple Music does, but closer than ever.

Here is a post about extracting prominent and unique colors from a UIImage using Core Graphics in Swift. I will analyze the image's pixels, and use a method called k-means clustering to find the dominant ones.

The Color Extraction Method

I will create an extension on UIImage and add a method extractColors on it:

extension UIImage {

  /// Extracts the most prominent and unique colors from a UIImage.
  /// - Parameters:
  ///   - numberOfColors: The number of prominent colors to extract (default is 4).
  /// - Returns: An array of UIColors representing the prominent colors.
  func extractColors(numberOfColors: Int = 4) throws -> [UIColor] {
    // Function implementation
  }
}

This function takes the number of colors you want to extract. By default, it extracts four colors. Then, I make sure the image has an underlying CGImage:

guard let _ = image.cgImage else {
    throw NSError(domain: "Invalid image", code: 0, userInfo: nil)
}

A CGImage is a Core Graphics image, which I will use for pixel manipulation.

Downsampling the Image

Processing a large image can be slow, so I will downsample it to speed things up:

let size = CGSize(width: 200, height: 200 * image.size.height / image.size.width)
UIGraphicsBeginImageContext(size)
image.draw(in: CGRect(origin: .zero, size: size))
guard let resizedImage = UIGraphicsGetImageFromCurrentImageContext() else {
    UIGraphicsEndImageContext()
    throw NSError(domain: "Failed to resize image", code: 0, userInfo: nil)
}
UIGraphicsEndImageContext()

I resize the image to a width of 200 pixels while maintaining its aspect ratio. This reduces the number of pixels to process.

Creating a Pixel Buffer

Next, I set up a pixel buffer to hold the image data:

guard let data = calloc(height * width, MemoryLayout<UInt32>.size) else {
    throw NSError(domain: "Failed to allocate memory", code: 0, userInfo: nil)
}

defer { free(data) }

I allocate memory for the pixel data and ensure it is freed after the method is done.

Drawing the Image into a CGContext

I create a CGContext to draw the image so I can access its pixel data:

let width = inputCGImage.width
let height = inputCGImage.height
let bytesPerPixel = 4
let bytesPerRow = bytesPerPixel * width
let bitsPerComponent = 8
let colorSpace = CGColorSpaceCreateDeviceRGB()
let bitmapInfo = CGImageAlphaInfo.premultipliedLast.rawValue

let context = CGContext(data: data, width: width, height: height,
                        bitsPerComponent: bitsPerComponent,
                        bytesPerRow: bytesPerRow,
                        space: colorSpace,
                        bitmapInfo: bitmapInfo)!
context.draw(inputCGImage, in: CGRect(x: 0, y: 0, width: width, height: height))

A CGContext is an object that holds the drawing environment that allows to work with image data directly and manipulate the raw pixels of the image.

Once the CGContext is set up, I draw the image into it. This transfers the pixel data from the original image inputCGImage into the CGContext.

Without the CGContext, I cannot easily work with the individual pixel data. This is important because once the image is in the context, I can loop through the pixel buffer, read the red, green, and blue values for each pixel, and analyze them to extract the prominent colors.

Extracting Pixel Data

Now, I read the pixel data into an array:

let pixelBuffer = data.bindMemory(to: UInt8.self, capacity: width * height * bytesPerPixel)
var pixelData = [PixelData]()

for y in 0..<height {
  for x in 0..<width {
    let offset = ((width * y) + x) * bytesPerPixel
    let r = pixelBuffer[offset]
    let g = pixelBuffer[offset + 1]
    let b = pixelBuffer[offset + 2]
    pixelData.append(PixelData(red: Double(r), green: Double(g), blue: Double(b)))
  }
}

I loop through each pixel, extract the red, green, and blue (RGB) values, and store them in an array of PixelData, a simple structure to hold the values:

private struct PixelData {
  let red: Double
  let green: Double
  let blue: Double
}

K-Means Clustering for Color Analysis

To find the most prominent colors, I will use k-means clustering, a method that groups data into clusters. To start with, I define a Cluster structure:

struct Cluster {
  var center: PixelData
  var points: [PixelData]
}

Each cluster has a center (the average color) and a list of points (pixels). Here is the function that performs the clustering:

func kMeansCluster(pixels: [PixelData], k: Int, maxIterations: Int = 10) -> [Cluster] {
  // Function implementation
}

I initialize clusters with random centers and then iterate to refine them:

var clusters = [Cluster]()
for _ in 0..<k {
  if let randomPixel = pixels.randomElement() {
    clusters.append(Cluster(center: randomPixel, points: []))
  }
}

For each pixel, I find the nearest cluster center:

for pixel in pixels {
    var minDistance = Double.greatestFiniteMagnitude
    var closestClusterIndex = 0
    
    for (index, cluster) in clusters.enumerated() {
        let distance = euclideanDistance(pixel1: pixel, pixel2: cluster.center)
        if distance < minDistance {
            minDistance = distance
            closestClusterIndex = index
        }
    }
    clusters[closestClusterIndex].points.append(pixel)
}

I use the Euclidean distance to measure how close a pixel is to a cluster center:

private func euclideanDistance(pixel1: PixelData, pixel2: PixelData) -> Double {
  let dr = pixel1.red - pixel2.red
  let dg = pixel1.green - pixel2.green
  let db = pixel1.blue - pixel2.blue
  return sqrt(dr * dr + dg * dg + db * db)
}

This measures the straight-line distance between two points in RGB space. After assigning pixels, I update each cluster's center:

for clusterIndex in 0..<clusters.count {
  let cluster = clusters[clusterIndex]
  if cluster.points.isEmpty { continue }
  let sum = cluster.points.reduce(PixelData(red: 0, green: 0, blue: 0)) { (result, pixel) -> PixelData in
    return PixelData(red: result.red + pixel.red, green: result.green + pixel.green, blue: result.blue + pixel.blue)
  }
  let count = Double(cluster.points.count)
  clusters[clusterIndex].center = PixelData(red: sum.red / count, green: sum.green / count, blue: sum.blue / count)
}

I calculate the average RGB values of all pixels in the cluster to find the new center.

Converting Clusters to UIColors

Finally, I convert the cluster centers to UIColor objects:

let colors = clusters.map { cluster -> UIColor in
  UIColor(red: CGFloat(cluster.center.red / 255.0),
          green: CGFloat(cluster.center.green / 255.0),
          blue: CGFloat(cluster.center.blue / 255.0),
          alpha: 1.0)
}

These colors represent the most prominent colors in the image. Here is the full code for reference:

extension UIImage {

  /// Extracts the most prominent and unique colors from the image.
  ///
  /// - Parameter numberOfColors: The number of prominent colors to extract (default is 4).
  /// - Returns: An array of UIColors representing the prominent colors.
  func extractColors(numberOfColors: Int = 4) throws -> [UIColor] {
    // Ensure the image has a CGImage
    guard let _ = self.cgImage else {
      throw NSError(domain: "Invalid image", code: 0, userInfo: nil)
    }

    let size = CGSize(width: 200, height: 200 * self.size.height / self.size.width)
    UIGraphicsBeginImageContext(size)
    self.draw(in: CGRect(origin: .zero, size: size))
    guard let resizedImage = UIGraphicsGetImageFromCurrentImageContext() else {
      UIGraphicsEndImageContext()
      throw NSError(domain: "Failed to resize image", code: 0, userInfo: nil)
    }
    UIGraphicsEndImageContext()

    guard let inputCGImage = resizedImage.cgImage else {
      throw NSError(domain: "Invalid resized image", code: 0, userInfo: nil)
    }

    let width = inputCGImage.width
    let height = inputCGImage.height
    let bytesPerPixel = 4
    let bytesPerRow = bytesPerPixel * width
    let bitsPerComponent = 8

    guard let data = calloc(height * width, MemoryLayout<UInt32>.size) else {
      throw NSError(domain: "Failed to allocate memory", code: 0, userInfo: nil)
    }

    defer { free(data) }

    let colorSpace = CGColorSpaceCreateDeviceRGB()
    let bitmapInfo = CGImageAlphaInfo.premultipliedLast.rawValue

    guard let context = CGContext(data: data, width: width, height: height,
                                  bitsPerComponent: bitsPerComponent,
                                  bytesPerRow: bytesPerRow,
                                  space: colorSpace,
                                  bitmapInfo: bitmapInfo) else {
      throw NSError(domain: "Failed to create CGContext", code: 0, userInfo: nil)
    }

    context.draw(inputCGImage, in: CGRect(x: 0, y: 0, width: width, height: height))

    let pixelBuffer = data.bindMemory(to: UInt8.self, capacity: width * height * bytesPerPixel)
    var pixelData = [PixelData]()
    for y in 0..<height {
      for x in 0..<width {
        let offset = ((width * y) + x) * bytesPerPixel
        let r = pixelBuffer[offset]
        let g = pixelBuffer[offset + 1]
        let b = pixelBuffer[offset + 2]
        pixelData.append(PixelData(red: Double(r), green: Double(g), blue: Double(b)))
      }
    }

    let clusters = kMeansCluster(pixels: pixelData, k: numberOfColors)

    let colors = clusters.map { cluster -> UIColor in
      UIColor(red: CGFloat(cluster.center.red / 255.0),
                     green: CGFloat(cluster.center.green / 255.0),
                     blue: CGFloat(cluster.center.blue / 255.0),
                     alpha: 1.0)
    }

    return colors
  }

  private struct PixelData {
    let red: Double
    let green: Double
    let blue: Double
  }

  private struct Cluster {
    var center: PixelData
    var points: [PixelData]
  }

  private func kMeansCluster(pixels: [PixelData], k: Int, maxIterations: Int = 10) -> [Cluster] {
    var clusters = [Cluster]()
    for _ in 0..<k {
      if let randomPixel = pixels.randomElement() {
        clusters.append(Cluster(center: randomPixel, points: []))
      }
    }

    for _ in 0..<maxIterations {
      for clusterIndex in 0..<clusters.count {
        clusters[clusterIndex].points.removeAll()
      }

      for pixel in pixels {
        var minDistance = Double.greatestFiniteMagnitude
        var closestClusterIndex = 0
        for (index, cluster) in clusters.enumerated() {
          let distance = euclideanDistance(pixel1: pixel, pixel2: cluster.center)
          if distance < minDistance {
            minDistance = distance
            closestClusterIndex = index
          }
        }
        clusters[closestClusterIndex].points.append(pixel)
      }

      for clusterIndex in 0..<clusters.count {
        let cluster = clusters[clusterIndex]
        if cluster.points.isEmpty { continue }
        let sum = cluster.points.reduce(PixelData(red: 0, green: 0, blue: 0)) { (result, pixel) -> PixelData in
          return PixelData(red: result.red + pixel.red, green: result.green + pixel.green, blue: result.blue + pixel.blue)
        }
        let count = Double(cluster.points.count)
        clusters[clusterIndex].center = PixelData(red: sum.red / count, green: sum.green / count, blue: sum.blue / count)
      }
    }

    return clusters
  }

  private func euclideanDistance(pixel1: PixelData, pixel2: PixelData) -> Double {
    let dr = pixel1.red - pixel2.red
    let dg = pixel1.green - pixel2.green
    let db = pixel1.blue - pixel2.blue
    return sqrt(dr * dr + dg * dg + db * db)
  }
}

Building the SwiftUI Palette View

Now that I have got the colors, let's display them using SwiftUI. Here is a view that shows the colors in a grid:

struct ColorPaletteView: View {
  let colors: [UIColor]

  let columns = [GridItem(.flexible()), GridItem(.flexible())]

  var body: some View {
    ScrollView {
      LazyVGrid(columns: columns, spacing: 16) {
        ForEach(0..<colors.count, id: \.self) { index in
          let uiColor = colors[index]
          let color = Color(uiColor)
          let hex = uiColor.toHexString()

          VStack {
            Rectangle()
              .fill(color)
              .frame(height: 60)
              .cornerRadius(12)

            Text(hex)
              .font(.caption)
              .foregroundColor(.primary)
          }
        }
      }
      .padding()
    }
  }
}

extension UIColor {
  func toHexString() -> String {
    var rFloat: CGFloat = 0
    var gFloat: CGFloat = 0
    var bFloat: CGFloat = 0
    var aFloat: CGFloat = 0

    self.getRed(&rFloat, green: &gFloat, blue: &bFloat, alpha: &aFloat)

    let rInt = Int(rFloat * 255)
    let gInt = Int(gFloat * 255)
    let bInt = Int(bFloat * 255)

    return String(format: "#%02X%02X%02X", rInt, gInt, bInt)
  }
}

I set up the ProminentColorsView view to show the image and the color pallete from it:

struct ProminentColorsView: View {
  @State private var colors: [UIColor] = []
  @State private var errorMessage: String?
  private let image =  UIImage(named: "example")!

  var body: some View {
    VStack {
      Image(uiImage: image)
        .resizable()
        .aspectRatio(contentMode: .fit)
        .cornerRadius(16)
        .padding(.horizontal)

      if !colors.isEmpty {
        ColorPaletteView(colors: colors)
      }

      if let errorMessage = errorMessage {
        Text("Error: \(errorMessage)")
          .foregroundColor(.red)
      }
    }
    .task {
      do {
        colors = try image.extractColors(numberOfColors: 8)
      } catch {
        errorMessage = error.localizedDescription
      }
    }
  }
}

And here is the beautiful colors generated from my favorite J-Pop song, Summer Haze:

Moving Forward

From the first try a few years ago to now, I am content with this implementation of k-means clustering. Almost as good as Apple Music, but still not there.

Give it a try with your image, and happy coloring!