Machine Learning Blog Post: Snip 3

Demo.m

#pragma mark - Private Classifiers

- (void)classifyWithModelForImage:(UIImage *)image completion:(void (^)(NSDictionary <NSString *, NSString *> *))completion {
    dispatch_async(dispatch_get_global_queue(QOS_CLASS_UTILITY, 0), ^{
        NSMutableDictionary *returnValues = [NSMutableDictionary new];
        CVPixelBufferRef pixelBuffer = [self generatePixelBufferFromImage:image];

        if (pixelBuffer == nil) {
            completion(nil);
            return;
        }

        NSError *predictionError;
        Inceptionv3Output *results = [self.model predictionFromImage:pixelBuffer error:&predictionError];
        returnValues[BFR_CLASSIFICATION_KEY] = results.classLabel;

        // Tack on confidence level
        if (results.classLabelProbs[results.classLabel] != nil) {
            NSString *confidenceString = (NSString *)results.classLabelProbs[results.classLabel];
            CGFloat confidenceNumbericalValue = roundf(confidenceString.floatValue * 100);
            returnValues[BFR_CLASSIFICATION_CONFIDENCE_KEY] = @(confidenceNumbericalValue).stringValue;
        }

        dispatch_async(dispatch_get_main_queue(), ^ {
            completion(returnValues);
        });
    });
}

- (void)classifyWithVisionForImage:(UIImage *)image completion:(void (^)(NSDictionary<NSString *,NSString *> * _Nonnull))completion {
    if (self.visionModel == nil) {
        completion(@{BFR_CLASSIFICATION_KEY:@"",BFR_CLASSIFICATION_CONFIDENCE_KEY:@""});
    }
    
    VNCoreMLRequest *request = [[VNCoreMLRequest alloc] initWithModel:self.visionModel completionHandler:^(VNRequest *request, NSError *error) {
        dispatch_async(dispatch_get_main_queue(), ^ {
            if (error || request.results.count < 1) {
                completion(@{BFR_CLASSIFICATION_KEY:@"",BFR_CLASSIFICATION_CONFIDENCE_KEY:@""});
            }
            
            // Grab the highest level of confidence result
            VNClassificationObservation *result = request.results.firstObject;
            NSString *foundObject = result.identifier;
            CGFloat confidenceNumbericalValue = roundf(result.confidence * 100);
            completion(@{BFR_CLASSIFICATION_KEY:foundObject, BFR_CLASSIFICATION_CONFIDENCE_KEY:@(confidenceNumbericalValue).stringValue});
        });
    }];
    request.imageCropAndScaleOption = VNImageCropAndScaleOptionCenterCrop;
    
    CIImage *coreGraphicsImage = [[CIImage alloc] initWithImage:image];
    
    if (coreGraphicsImage == nil) {
        completion(@{BFR_CLASSIFICATION_KEY:@"",BFR_CLASSIFICATION_CONFIDENCE_KEY:@""});
    }
    
    CGImagePropertyOrientation cgOrientation = [self cgOrientationFromUIImageOrientation:image.imageOrientation];
    
    dispatch_async(dispatch_get_global_queue(QOS_CLASS_UTILITY, 0), ^{
        VNImageRequestHandler *handler = [[VNImageRequestHandler alloc] initWithCIImage:coreGraphicsImage orientation:cgOrientation options:@{}];
        [handler performRequests:@[request] error:nil];
    });
}