K-means colour matching — putting this old file somewhere

K-means colour matching.m

@import Foundation;
@import CoreGraphics;
@import ImageIO;

static NSUInteger maxThumbnailSize = 20;

extern uint64_t dispatch_benchmark(size_t count, void (^block)(void));

/// Returns a CGImage with a +1 retain count.
static CGImageRef createThumbnailFromURL(NSURL *imageURL, NSUInteger maxPixelSize)
{
	CGImageSourceRef source = CGImageSourceCreateWithURL((__bridge CFURLRef)(imageURL), NULL);
	
	NSDictionary *const thumbnailOptions = @{
											 (id)kCGImageSourceThumbnailMaxPixelSize : @(maxPixelSize),
											 (id)kCGImageSourceCreateThumbnailFromImageAlways : @YES,
											 (id)kCGImageSourceShouldCache : @NO,
											 };
	CGImageRef image = CGImageSourceCreateThumbnailAtIndex(source, 0, (__bridge CFDictionaryRef)thumbnailOptions);
	
	CFRelease(source);
	source = NULL;
	
	return image;
}

typedef struct {
	unsigned char red;
	unsigned char green;
	unsigned char blue;
	unsigned char unused;
} ZORGBColour;

static void enumeratePixels(CGImageRef image, void (^block)(ZORGBColour colour, NSUInteger x, NSUInteger y))
{
	if (block == nil) {
		return;
	}
	
	NSUInteger const width = CGImageGetWidth(image);
	NSUInteger const height = CGImageGetHeight(image);
	NSUInteger const pixelCount = width * height;
	
	NSLog(@"%lu x %lu = %lu", width, height, pixelCount);
	
	static NSUInteger const numberOfConcurrentOperations = 4;
	
	// set up a dispatch group
	
	for (NSUInteger concurrencyIndex = 0; concurrencyIndex < numberOfConcurrentOperations; ++concurrencyIndex) {
		
		// dispatch this in the group
		
		NSLog(@"- - -");
		
		ZORGBColour pixelData = {};
		static NSUInteger const contextWidth = 1, contextHeight = 1;
		static NSUInteger const bitsPerComponent = 8;
		static NSUInteger const bytesPerRow = 4 * contextWidth;
		CGColorSpaceRef colourSpace = CGColorSpaceCreateDeviceRGB();
		CGContextRef context = CGBitmapContextCreate(&pixelData, contextWidth, contextHeight, bitsPerComponent, bytesPerRow, colourSpace, (CGBitmapInfo)kCGImageAlphaNoneSkipLast);
		CGColorSpaceRelease(colourSpace);
		colourSpace = NULL;
		
		for (NSUInteger pixelIndex = concurrencyIndex; pixelIndex < pixelCount; pixelIndex += numberOfConcurrentOperations) {
			NSUInteger const x = pixelIndex % width;
			NSUInteger const y = pixelIndex / width;
			
			CGContextDrawTiledImage(context, CGRectMake(width - x, height - y, width, height), image);
			
			block(pixelData, x, y);
		}
		
		CGContextRelease(context);
		context = NULL;
	}
	
	// wait for the dispatch group to finish
}

@implementation NSSet (DHAccumulate)

- (id)dh_objectByReducingWithOptions:(NSEnumerationOptions)options usingBlock:(id (^)(id object1, id object2))reductionBlock
{
	if (options & NSEnumerationConcurrent) {
		[NSException raise:NSInvalidArgumentException format:@"%s is not ready for concurrent operation. Sorry.", __PRETTY_FUNCTION__];
	}
	
	__block id accumulatedValue;
	
	[self enumerateObjectsWithOptions:options usingBlock:^(id object, BOOL *stop) {
		if (accumulatedValue == nil) {
			accumulatedValue = object;
			return;
		}
		accumulatedValue = reductionBlock(accumulatedValue, object);
		
		if (accumulatedValue == nil) {
			[NSException raise:NSInvalidArgumentException format:@"@s - reductionBlock must not return nil."];
		}
	}];
	
	if (accumulatedValue == nil) {
		[NSException raise:NSInvalidArgumentException format:@"%s must not be called on an empty array.", __PRETTY_FUNCTION__];
	}
	
	return accumulatedValue;
}

@end

@protocol KMeansable <NSObject>

+ (id <KMeansable>)meanOfObjects:(NSSet *)objects;

- (double)distanceFromObject:(id <KMeansable>)other;

@end

@interface ColourCluster : NSObject


+ (instancetype)clusterWithMean:(id <KMeansable>)initialMean;

/// Designated initialiser
- (instancetype)initWithMean:(id <KMeansable>)initialMean __attribute((objc_designated_initializer));

@property (nonatomic, strong, readonly) id<KMeansable> mean;

- (void)updateMean;

- (void)addObject:(id <KMeansable>)object;
- (void)removeObject:(id <KMeansable>)object;

@end

@interface ColourCluster ()

@property (nonatomic, strong) id<KMeansable> mean;
@property (nonatomic, strong, readonly) NSMutableSet *members;

@end

@implementation ColourCluster

+ (instancetype)clusterWithMean:(id <KMeansable>)initialMean
{
	return [[self alloc] initWithMean:initialMean];
}

- (instancetype)init
{
	return [self initWithMean:nil];
}

- (instancetype)initWithMean:(id <KMeansable>)initialMean
{
	self = [super init];
	if (self == nil) return nil;
	
	if (initialMean == nil) {
		[NSException raise:NSInvalidArgumentException format:@"%@ can not be initialised without a mean.", [self class]];
	}
	
	_mean = initialMean;
	_members = [NSMutableSet set];
	
	return self;
}

- (void)updateMean
{
	[self setMean:[[[self mean] class] meanOfObjects:[self members]]];
}

- (void)addObject:(id<KMeansable>)object
{
	[[self members] addObject:object];
}

- (void)removeObject:(id<KMeansable>)object
{
	[[self members] removeObject:object];
}

- (NSString *)description
{
	return [NSString stringWithFormat:@"%@, Mean: %@, Members: %@", [super description], [self mean], [self members]];
}

@end

@interface PixelColour : NSObject <KMeansable>

+ (instancetype)colourWithRed:(unsigned char)red green:(unsigned char)green blue:(unsigned char)blue;
+ (instancetype)colourWithRGBColour:(ZORGBColour)colour;

/// Designated initialiser
- (instancetype)initWithRGBColour:(ZORGBColour)colour __attribute((objc_designated_initializer));

@property (nonatomic, readonly) ZORGBColour colour;

@end

@implementation PixelColour

+ (instancetype)colourWithRed:(unsigned char)red green:(unsigned char)green blue:(unsigned char)blue
{
	return [self colourWithRGBColour:(ZORGBColour){red, green, blue}];
}

+ (instancetype)colourWithRGBColour:(ZORGBColour)colour
{
	return [[self alloc] initWithRGBColour:colour];
}

- (instancetype)init
{
	NSLog(@"Incorrect initialiser “%s” sent to %@", __PRETTY_FUNCTION__, [self class]);
	return [self initWithRGBColour:(ZORGBColour){}];
}

- (instancetype)initWithRGBColour:(ZORGBColour)colour
{
	self = [super init];
	if (self == nil) return nil;
	
	_colour = colour;
	
	return self;
}

+ (id<KMeansable>)meanOfObjects:(NSSet *)objects
{
	int redSum = 0, greenSum = 0, blueSum = 0;
	
	for (PixelColour *colour in objects) {
		ZORGBColour rgb = [colour colour];
		redSum += rgb.red;
		greenSum += rgb.green;
		blueSum += rgb.blue;
	}
	
	double const count = [objects count]; // Deliberately a double so the divisions are floating point calculations.
	
	return [[PixelColour alloc] initWithRGBColour:(ZORGBColour){
		.red   = round(  redSum / count),
		.green = round(greenSum / count),
		.blue  = round( blueSum / count),
	}];
}

- (double)distanceFromObject:(id<KMeansable>)other
{
	if ([other isKindOfClass:[self class]] == NO) {
		[NSException raise:NSInvalidArgumentException format:@"Can not find distance to a %@", [other class]];
		return HUGE;
	}
	
	PixelColour *otherColour = other;
	
	int const redDiff = (int)([self colour].red) - (int)([otherColour colour].red);
	int const greenDiff = (int)([self colour].green) - (int)([otherColour colour].green);
	int const blueDiff = (int)([self colour].blue) - (int)([otherColour colour].blue);
	
	return sqrt(redDiff * redDiff + greenDiff * greenDiff + blueDiff * blueDiff);
}

- (NSString *)description
{
	return [NSString stringWithFormat:@"RGB(%u %u %u)", [self colour].red, [self colour].green, [self colour].blue];
}

@end

@implementation NSSet (ZOMap)

- (NSSet *)zo_setByMappingObjectsUsingMap:(id (^)(id object))map
{
	NSMutableSet *mappedSet = [NSMutableSet setWithCapacity:[self count]];
	
	for (id object in self) {
		id mappedObject = map(object);
		if (mappedObject) [mappedSet addObject:mappedObject];
	}
	
	return mappedSet;
}

@end

static NSSet *updateClusters(NSSet *means, NSSet *objects)
{
	NSSet *const clusters = [means zo_setByMappingObjectsUsingMap:^id(id <KMeansable>object) {
		return [ColourCluster clusterWithMean:object];
	}];
	
	for (PixelColour *pixelColour in objects) {
		// assign data points to cluster with nearest center
		// let sn = arg mink |xn − mk|
		
		double minimumDistance = HUGE;
		ColourCluster *nearestCluster;
		
		for (ColourCluster *cluster in clusters) {
			double const distance = [pixelColour distanceFromObject:[cluster mean]];
			if (distance < minimumDistance) {
				minimumDistance = distance;
				nearestCluster = cluster;
			}
		}
		
		[nearestCluster addObject:pixelColour];
	}
	
	for (ColourCluster *cluster in clusters) {
		// re-compute means
		// let mk = mean{xn : sn = k}
		[cluster updateMean];
		
		NSLog(@" %@", cluster);
	}
	
	return clusters;
}

static void something(NSURL *sourceImageURL)
{
	CGImageRef imageRef = createThumbnailFromURL(sourceImageURL, maxThumbnailSize);
	
	NSMutableSet *allPixels = [NSMutableSet set];
	
	enumeratePixels(imageRef, ^(ZORGBColour pixelColour, NSUInteger x, NSUInteger y) {
		NSLog(@"(%lu, %lu) %u %u %u", x, y, pixelColour.red, pixelColour.green, pixelColour.blue);
		[allPixels addObject:[[PixelColour alloc] initWithRGBColour:pixelColour]];
	});
	
	CGImageRelease(imageRef);
	imageRef = NULL;
	
	NSSet *const initalMeans = [NSSet setWithObjects:
								[PixelColour colourWithRed:0 green:0 blue:0],
								[PixelColour colourWithRed:255 green:255 blue:255],
								nil];
	
	NSSet *clusters;
	NSSet *means = initalMeans;
	
	for (NSUInteger count = 0; count < 5; ++count) {
		clusters = updateClusters(means, allPixels);
		means = [clusters zo_setByMappingObjectsUsingMap:^id(ColourCluster *cluster) {
			return [cluster mean];
		}];
	}
	
}

int main(int argc, char *argv[]) {
	@autoreleasepool {
		
		something([NSURL fileURLWithPath:@"/Users/dhill/Desktop/step_25437 3.jpg"]);
		
	}
	
	return 0;
}