本文将介绍如何在具体的iOS项目中使用 Core ML模型。
主要步骤
一、准备一个Core ML模型
Core ML模型文件一般以.mlmodel作为后缀，我们可以通过很多途径获得一个已经训练好的模型，最直接的方式，就是从苹果官网上直接下载。
点开链接，我们选择MobileNet下载，获得一个MobileNet.mlmodel文件。
我们可以看一下MobileNet的介绍：
MobileNets are based on a streamlined architecture that have depth-wise separable convolutions to build lightweight, deep neural networks. Detects the dominant objects present in an image from a set of 1000 categories such as trees, animals, food, vehicles, people, and more.
这是一个图片分类的模型，可以识别像树木、动物、食物、车辆、人物等。由此可见，我们需要事先准备一些包含这些物体的场景或者图片。
二、准备一个iOS项目
在开始之前，我们需要准备一个iOS项目，来获取图片，加载模型，最终完成图片分类。作为一个iOS工程师，这部分将不再赘述，你可以直接从我的GitHub获取一个已经构建好的项目：
git clone [email protected]:yangchenlarkin/CoreML.git
项目的首页是一个UITableViewController，点击第一个“物体识别”，即可进入到我们的Demo页面ORViewController，代码位于ObjectRecognition文件加下。
ORViewController可以通过后置摄像头的拍摄来获取图片。
查看ORViewController最后的TODO部分，这里就是我们处理图片的代码了:
#pragma mark - predict
- (void)predict:(UIImage *)image {
//TODO
}
三、导入.mlmodel模型文件
我们将第一步获得的MobileNet.mlmodel文件，拖拽到项目中，勾选Copy items if needed。在项目中选中MobileNet.mlmodel，让我们来看一下这个模型。
可以看到该文件包含三个部分：
- 模型描述；
- 模型类MobileNet；
Xcode帮我们创建了一个叫MobileNet的类，我们点击旁边的箭头可以查看它的头文件，我们也可以在代码中直接引用、使用这个类。 - 模型参数；
这里主要包含了模型输入和输出的描述，在MobileNet中，入参是一个224x224的图片，出参有两个，classLabelProbs一个字典，key是分类名，value是入参图片命中该分类的概率，classLabel是入参图片最有可能（概率值最大）的分类名。
从文件中可以看到，我们需要将图片转化成224x224的尺寸大小（如果你有一定的机器学习基础，就可以理解这里图片大小固定的原因）。而模型会告诉我们这张图片最有可能包含哪一样物体。
四、查看MobileNet类
在写代码之前，我们来看一下MobileNet类。
点击MobileNet.mlmodel文件中MobileNet旁边的箭头，我们可以查看该类的头文件。暂时无视掉其他代码，我们关注一下这个方法：
/**
Make a prediction using the convenience interface
@param image Input image to be classified as color (kCVPixelFormatType_32ARGB) image buffer, 224 pixels wide by 224 pixels high:
@param error If an error occurs, upon return contains an NSError object that describes the problem. If you are not interested in possible errors, pass in NULL.
@return the prediction as MobileNetOutput
*/
- (nullable MobileNetOutput *)predictionFromImage:(CVPixelBufferRef)image error:(NSError * _Nullable * _Nullable)error;
入参：
从方法声明和注释中，我们发现，我们的输入图片需要是CVPixelBufferRef 类型的，格式需要是kCVPixelFormatType_32ARGB格式的。返回值：
该函数返回了一个MobileNetOutput类对象。我们跟进此类中去（其实也在这个文件里）。我们发现，他包含两个属性：
/// Probability of each category as dictionary of strings to doubles
@property (readwrite, nonatomic, strong) NSDictionary * classLabelProbs;
/// Most likely image category as string value
@property (readwrite, nonatomic, strong) NSString * classLabel;
这两个属性正好对应了MobileNet.mlmodel中的两个输出。
至此，我们脑海中大概有了一个使用模型的流程了：
1.将UIImage转化成CVPixelBufferRef，图片格式是kCVPixelFormatType_32ARGB，图片尺寸是224x224
2.创建模型对象，调用预测方法，获得MobileNetOutput
3.从MobileNetOutput的classLabel属性，就是我们需要的分类名
4.如果业务需要，我们也可以从classLabelProbs属性中，获得所有分类及其概率。
五、添加代码，实现图片分类
1.图片缩放
首先我们需要对图片做居中剪裁和缩放，将图片的中间的正方形部分，缩放到224x224大小。
我们定义一个缩放方法：
#pragma mark - predict
- (UIImage *)scaleImage:(UIImage *)image size:(CGFloat)size {
UIGraphicsBeginImageContextWithOptions(CGSizeMake(size, size), YES, 1);
CGFloat x, y, w, h;
CGFloat imageW = image.size.width;
CGFloat imageH = image.size.height;
if (imageW > imageH) {
w = imageW / imageH * size;
h = size;
x = (size - w) / 2;
y = 0;
} else {
h = imageH / imageW * size;
w = size;
y = (size - h) / 2;
x = 0;
}
[image drawInRect:CGRectMake(x, y, w, h)];
UIImage * scaledImage = UIGraphicsGetImageFromCurrentImageContext();
UIGraphicsEndImageContext();
return scaledImage;
}
然后调用他：
#pragma mark - predict
- (UIImage *)scaleImage:(UIImage *)image size:(CGFloat)size {
//...
}
- (void)predict:(UIImage *)image {
UIImage *scaledImage = [self scaleImage:image size:224];
//TODO
}
2.格式转换
直接上代码吧：
#pragma mark - predict
- (UIImage *)scaleImage:(UIImage *)image size:(CGFloat)size {
//...
}
- (CVPixelBufferRef)pixelBufferFromCGImage:(CGImageRef)image {
NSDictionary *options = @{
(NSString *)kCVPixelBufferCGImageCompatibilityKey : @YES,
(NSString *)kCVPixelBufferCGBitmapContextCompatibilityKey : @YES,
(NSString *)kCVPixelBufferIOSurfacePropertiesKey: [NSDictionary dictionary]
};
CVPixelBufferRef pxbuffer = NULL;
CGFloat frameWidth = CGImageGetWidth(image);
CGFloat frameHeight = CGImageGetHeight(image);
CVReturn status = CVPixelBufferCreate(kCFAllocatorDefault,
frameWidth,
frameHeight,
kCVPixelFormatType_32ARGB,
(__bridge CFDictionaryRef) options,
&pxbuffer);
NSParameterAssert(status == kCVReturnSuccess && pxbuffer != NULL);
CVPixelBufferLockBaseAddress(pxbuffer, 0);
void *pxdata = CVPixelBufferGetBaseAddress(pxbuffer);
NSParameterAssert(pxdata != NULL);
CGColorSpaceRef rgbColorSpace = CGColorSpaceCreateDeviceRGB();
CGContextRef context = CGBitmapContextCreate(pxdata,
frameWidth,
frameHeight,
8,
CVPixelBufferGetBytesPerRow(pxbuffer),
rgbColorSpace,
(CGBitmapInfo)kCGImageAlphaNoneSkipFirst);
NSParameterAssert(context);
CGContextConcatCTM(context, CGAffineTransformIdentity);
CGContextDrawImage(context, CGRectMake(0,
0,
frameWidth,
frameHeight),
image);
CGColorSpaceRelease(rgbColorSpace);
CGContextRelease(context);
CVPixelBufferUnlockBaseAddress(pxbuffer, 0);
return pxbuffer;
}
- (void)predict:(UIImage *)image {
UIImage *scaledImage = [self scaleImage:image size:224];
CVPixelBufferRef buffer = [self pixelBufferFromCGImage:scaledImage.CGImage];
//TODO
}
3.预测
我们需要生成MobileNet类对象，需要在文件开始处引入MobileNet.h
#import "MobileNet.h"
然后添加代码：
#pragma mark - predict
- (UIImage *)scaleImage:(UIImage *)image size:(CGFloat)size {
//...
}
- (CVPixelBufferRef)pixelBufferFromCGImage:(CGImageRef)image {
//...
}
- (void)predict:(UIImage *)image {
UIImage *scaledImage = [self scaleImage:image size:224];
CVPixelBufferRef buffer = [self pixelBufferFromCGImage:scaledImage.CGImage];
NSError *error = nil;
MobileNet *model = [[MobileNet alloc] init];//模型的创建和初始化，可以放在self属性中懒加载，这样不用每次重新创建
MobileNetOutput *output = [model predictionFromImage:buffer error:&error];
if (error) {
NSLog(@"%@", error);
return;
}
//TODO
}
4.显示预测
在ORViewController中有一个属性：
@property (nonatomic, copy) NSArray *> *array;
我们可以通过设置这个数组来控制一个tableview的显示：
#pragma mark - predict
- (UIImage *)scaleImage:(UIImage *)image size:(CGFloat)size {
//...
}
- (CVPixelBufferRef)pixelBufferFromCGImage:(CGImageRef)image {
//...
}
- (void)predict:(UIImage *)image {
UIImage *scaledImage = [self scaleImage:image size:224];
CVPixelBufferRef buffer = [self pixelBufferFromCGImage:scaledImage.CGImage];
NSError *error = nil;
MobileNet *model = [[MobileNet alloc] init];//模型的创建和初始化，可以放在self属性中懒加载，这样不用每次重新创建
MobileNetOutput *output = [model predictionFromImage:buffer error:&error];
if (error) {
NSLog(@"%@", error);
return;
}
NSMutableArray *result = [NSMutableArray arrayWithCapacity:output.classLabelProbs.count + 1];
[result addObject:@[@"这张图片可能包含：", output.classLabel]];
[output.classLabelProbs enumerateKeysAndObjectsUsingBlock:^(NSString * _Nonnull key, NSNumber * _Nonnull obj, BOOL * _Nonnull stop) {
NSString *title = [NSString stringWithFormat:@"%@的概率：", key];
[result addObject:@[title, obj.stringValue]];
}];
self.array = result;
}
5.测试
至此我们的代码就算写完了，是不是很简单！让我们运行一下，到Google上搜索一些图片看看结果吧：