Get dominant colors of image

TTVLImageDominantColor.h

//
//  TTVLImageDominantColor.h
//  KMeans
//
//  Created by PhilCai on 2020/8/6.
//  Copyright © 2020 PhilCai. All rights reserved.
//

#import <UIKit/UIKit.h>

NS_ASSUME_NONNULL_BEGIN

/// Getting the dominant colors of an image using the CIE LAB color space,
/// and the k-means clustering algorithm.
@interface TTVLImageDominantColor : NSObject

/// Get the 1st dominant color of the image
/// It's prefered to be invoked on background threads
/// @param image image
+ (nullable UIColor *)getColorOf:(UIImage *)image;

/// Get the dominant colors of the image
/// It's prefered to be invoked on background threads
/// @param image image
+ (nullable NSArray <UIColor *> *)getColorsOf:(UIImage *)image;

@end

NS_ASSUME_NONNULL_END

TTVLImageDominantColor.mm

//
//  TTVLImageDominantColor.mm
//  KMeans
//
//  Created by PhilCai on 2020/8/6.
//  Copyright © 2020 PhilCai. All rights reserved.
//

#import "TTVLImageDominantColor.h"
#import <UIKit/UIKit.h>
#import <algorithm>
#import <vector>

// clang-format off
namespace ttvl {

class Vector3 {
   public:
    CGFloat x;
    CGFloat y;
    CGFloat z;
    
    Vector3(CGFloat x_, CGFloat y_, CGFloat z_) {
        x = x_;
        y = y_;
        z = z_;
    }
    
    Vector3 operator+(const Vector3& another) {
        Vector3 rgb = Vector3(this->x + another.x,
                              this->y + another.y,
                              this->z + another.z);
        return rgb;
    }
    
    Vector3 operator/(const NSUInteger divider) {
        Vector3 rgb = Vector3(this->x / divider,
                              this->y / divider,
                              this->z / divider);
        return rgb;
    }
    
    Vector3 operator*(const NSUInteger multiplier) {
        Vector3 rgb = Vector3(this->x * multiplier,
                              this->y * multiplier,
                              this->z * multiplier);
        return rgb;
    }
    
};

class Cluster {
   public:
    Vector3 centroid = Vector3(0, 0, 0);
    NSUInteger size;
    Cluster(Vector3 centroid_, NSUInteger size_) {
        centroid = centroid_;
        size = size_;
    }
};

bool operator<(const Cluster& a, const Cluster& b) {
    return a.size < b.size;
}

#pragma mark SRGB

static Vector3 SRGBToLinearSRGB(Vector3 srgbVector) {
    auto f = [](CGFloat c) -> CGFloat {
        if (c <= 0.04045) {
            return c / 12.92;
        } else {
            return powf((c + 0.055) / 1.055, 2.4);
        }
    };
    return Vector3(f(srgbVector.x), f(srgbVector.y), f(srgbVector.z));
}

static Vector3 LinearSRGBToSRGB(Vector3 lSrgbVector) {
    auto f = [](CGFloat c) -> CGFloat {
        if (c <= 0.0031308) {
            return c * 12.92;
        } else {
            return (1.055 * powf(c, 1.0 / 2.4)) - 0.055;
        }
    };
    return Vector3(f(lSrgbVector.x), f(lSrgbVector.y), f(lSrgbVector.z));
}

#pragma mark - XYZ

static Vector3 XYZToLAB(Vector3 xyzVector, Vector3 tristimulus) {
    auto f = [](CGFloat t) -> CGFloat {
        if (t > powf(6.0 / 29.0, 3.0)) {
            return powf(t, 1.0 / 3.0);
        } else {
            return ((1.0 / 3.0) * powf(29.0 / 6.0, 2.0) * t) + (4.0 / 29.0);
        }
    };
    auto fx = f(xyzVector.x / tristimulus.x);
    auto fy = f(xyzVector.y / tristimulus.y);
    auto fz = f(xyzVector.z / tristimulus.z);

    auto l = (116.0 * fy) - 16.0;
    auto a = 500 * (fx - fy);
    auto b = 200 * (fy - fz);

    return Vector3(l, a, b);
}

static Vector3 LABToXYZ(Vector3 labVector, Vector3 tristimulus) {
    auto f = [](CGFloat t) -> CGFloat {
        if (t > (6.0 / 29.0)) {
            return powf(t, 3.0);
        } else {
            return 3.0 * powf(6.0 / 29.0, 2.0) * (t - (4.0 / 29.0));
        }
    };
    auto c = (1.0 / 116.0) * (labVector.x + 16.0);

    auto y = tristimulus.y * f(c);
    auto x = tristimulus.x * f(c + ((1.0 / 500.0) * labVector.y));
    auto z = tristimulus.z * f(c - ((1.0 / 200.0) * labVector.z));

    return Vector3(x, y, z);
}

static Vector3 LinearSRGBToXYZ(Vector3 linearSrgbVector) {
    CGFloat LinearSRGBToXYZMatrix[3][3] = {{0.4124, 0.2126, 0.0193},
                                           {0.3576, 0.7152, 0.1192},
                                           {0.1805, 0.0722, 0.9505}};
    
    auto unscaledXYZVector =
        Vector3(LinearSRGBToXYZMatrix[0][0] * linearSrgbVector.x +
                LinearSRGBToXYZMatrix[1][0] * linearSrgbVector.y +
                LinearSRGBToXYZMatrix[2][0] * linearSrgbVector.z,
                LinearSRGBToXYZMatrix[0][1] * linearSrgbVector.x +
                LinearSRGBToXYZMatrix[1][1] * linearSrgbVector.y +
                LinearSRGBToXYZMatrix[2][1] * linearSrgbVector.z,
                LinearSRGBToXYZMatrix[0][2] * linearSrgbVector.x +
                LinearSRGBToXYZMatrix[1][2] * linearSrgbVector.y +
                LinearSRGBToXYZMatrix[2][2] * linearSrgbVector.z);
    
    return unscaledXYZVector * 100.0;
}

static Vector3 XYZToLinearSRGB(Vector3 xyzVector) {
    auto scaledXYZVector = xyzVector / 100.0;
    CGFloat XYZToLinearSRGBMatrix[3][3] = {{3.2406, -0.9689, 0.0557},
                                           {-1.5372, 1.8758, -0.2040},
                                           {-0.4986, 0.0415, 1.0570}};
    auto x = scaledXYZVector.x;
    auto y = scaledXYZVector.y;
    auto z = scaledXYZVector.z;
    
    return Vector3(XYZToLinearSRGBMatrix[0][0] * x +
                   XYZToLinearSRGBMatrix[1][0] * y +
                   XYZToLinearSRGBMatrix[2][0] * z,
                   XYZToLinearSRGBMatrix[0][1] * x +
                   XYZToLinearSRGBMatrix[1][1] * y +
                   XYZToLinearSRGBMatrix[2][1] * z,
                   XYZToLinearSRGBMatrix[0][2] * x +
                   XYZToLinearSRGBMatrix[1][2] * y +
                   XYZToLinearSRGBMatrix[2][2] * z);
    
}

static Vector3 RGBToLAB(Vector3 gVector) {
    auto srgbVector = gVector;
    auto lSrgbVector = SRGBToLinearSRGB(srgbVector);
    auto xyzVector = LinearSRGBToXYZ(lSrgbVector);
    auto D65Tristimulus = Vector3(5.047, 100.0, 108.883);
    auto labVector = XYZToLAB(xyzVector, D65Tristimulus);
    return labVector;
}

static Vector3 LABToRGB(Vector3 gVector) {
    auto D65Tristimulus = Vector3(5.047, 100.0, 108.883);
    auto xyzVector = LABToXYZ(gVector, D65Tristimulus);
    auto lSrgbVector = XYZToLinearSRGB(xyzVector);
    auto srgbVector = LinearSRGBToSRGB(lSrgbVector);
    auto rgbVector = srgbVector;
    return rgbVector;
}

#pragma mark - kmeans

static CGFloat distanceFunc(Vector3 a, Vector3 b) {
    auto r1 = a.x;
    auto g1 = a.y;
    auto b1 = a.z;
    auto r2 = b.x;
    auto g2 = b.y;
    auto b2 = b.z;

    return pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2);
}

static NSInteger findNearestCluster(Vector3 point,
                                    std::vector<Vector3> centroids,
                                    NSInteger k) {
    auto minDistance = CGFLOAT_MAX;
    NSInteger clusterIndex = 0;
    for (NSInteger i = 0; i < k; i++) {
        auto distance = distanceFunc(point, centroids[i]);
        if (distance < minDistance) {
            minDistance = distance;
            clusterIndex = i;
        }
    }
    return clusterIndex;
}

static std::vector<Cluster> kmeans(std::vector<Vector3> points,
                                   NSInteger k,
                                   CGFloat threadhold) {
    auto n = points.size();
    auto centroids = points;
    auto memberships = std::vector<NSInteger>(n, -1);
    auto clusterSizes = std::vector<NSInteger>(k, 0);

    auto error = 0;
    auto previousError = 0;

    do {
        error = 0;
        auto newCentroids = std::vector<Vector3>(k, Vector3(0, 0, 0));
        auto newClusterSizes = std::vector<NSInteger>(k, 0);

        for (NSUInteger i = 0; i < n; i++) {
            auto point = points[i];
            auto clusterIndex = findNearestCluster(point, centroids, k);
            if (memberships[i] != clusterIndex) {
                error += 1;
                memberships[i] = clusterIndex;
            }
            newClusterSizes[clusterIndex] += 1;
            newCentroids[clusterIndex] = newCentroids[clusterIndex] + point;
        }
        for (NSInteger i = 0; i < k; i++) {
            auto size = newClusterSizes[i];
            if (size > 0) {
                centroids[i] = newCentroids[i] / size;
            }
        }

        clusterSizes = newClusterSizes;
        previousError = error;

    } while (abs(error - previousError) > threadhold);

    auto ret = std::vector<Cluster>();
    for (NSInteger i = 0; i < MIN(centroids.size(), clusterSizes.size()); i++) {
        ret.push_back(Cluster(centroids[i], clusterSizes[i]));
    }

    return ret;
}

#pragma mark - UIImage

static std::vector<Vector3> getLABsFromImage(UIImage* uiimage) {
    auto scaledDimensionsForPixelLimit =
        [](size_t limit, size_t width,
           size_t height) -> std::tuple<size_t, size_t> {
        if (width * height > limit) {
            auto ratio = (CGFloat)(width) / (CGFloat)(height);
            auto maxWidth = sqrtf(ratio * (CGFloat)(limit));

            return std::tuple<size_t, size_t>(
                (size_t)(maxWidth), (size_t)((CGFloat)(limit) / maxWidth));
        }
        return std::tuple<size_t, size_t>(width, height);
    };

    auto maxSampledPixels = 100;
    auto image = uiimage.CGImage;
    auto width = 0;
    auto height = 0;

    std::tie(width, height) = scaledDimensionsForPixelLimit(maxSampledPixels,
                                                            CGImageGetWidth(image),
                                                            CGImageGetHeight(image));
    auto colorSpace = CGColorSpaceCreateDeviceRGB();
    auto rgbaCtx = CGBitmapContextCreateWithData(NULL,
                                                 width,
                                                 height,
                                                 8,
                                                 width * 4,
                                                 colorSpace,
                                                 kCGImageAlphaPremultipliedLast,
                                                 NULL,
                                                 NULL);
    CGContextDrawImage(rgbaCtx, CGRectMake(0, 0, width, height), image);

    auto labValues = std::vector<Vector3>();
    struct RGBA {
        uint8_t red;
        uint8_t green;
        uint8_t blue;
        uint8_t alpha;
    };
    RGBA* data = (RGBA*)CGBitmapContextGetData(rgbaCtx);

    for (NSInteger i = 0; i < height; i++) {
        for (NSInteger j = 0; j < width; j++) {
            auto pixel = data[j + i * width];
            if (pixel.alpha == UINT8_MAX) {
                /// alpha == 1
                auto rgb = Vector3(pixel.red, pixel.green, pixel.blue);
                auto lab = RGBToLAB(rgb / 255.0);
                labValues.push_back(lab);
            } else {
                /// no-op:
                /// ignore pixels with alpha channel
            }
        }
    }

    if (rgbaCtx) {
        CGContextRelease(rgbaCtx);
    }
    if (colorSpace) {
        CGColorSpaceRelease(colorSpace);
    }
    return labValues;
}
    
}

@implementation TTVLImageDominantColor

+ (UIColor*)getColorOf:(UIImage*)image {
    if (!image) {
        return nil;
    }
    auto labs = ttvl::getLABsFromImage(image);
    auto k = 16;
    auto clusters = ttvl::kmeans(labs, k, 0.001);
    if (clusters.empty() == true) {
        return nil;
    }
    auto dominant = std::max_element(clusters.begin(), clusters.end());
    auto lab = dominant->centroid;
    auto rgb = LABToRGB(lab);
    UIColor* color = [UIColor colorWithRed:rgb.x
                                     green:rgb.y
                                      blue:rgb.z
                                     alpha:1];
    return color;
}

+ (NSArray<UIColor*>*)getColorsOf:(UIImage*)image {
    if (!image) {
        return nil;
    }
    auto labs = ttvl::getLABsFromImage(image);
    auto k = 16;
    auto clusters = ttvl::kmeans(labs, k, 0.001);
    if (clusters.empty() == true) {
        return nil;
    }
    NSMutableArray<UIColor*>* colors = NSMutableArray.array;
    std::sort(clusters.begin(), clusters.end());
    auto currentIdx = 0;
    for (auto cluster : clusters) {
        if (currentIdx == k) {
            break;
        }
        auto lab = cluster.centroid;
        auto rgb = LABToRGB(lab);
        UIColor* color = [UIColor colorWithRed:rgb.x
                                         green:rgb.y
                                          blue:rgb.z
                                         alpha:1];

        [colors addObject:color];
        currentIdx ++;
    }
    return colors;
}
@end
// clang-format on