Article

Creating an Image Classifier Model

Train a machine learning model to classify images, and add it to your Core ML app.

Overview

An image classifier is a machine learning model that recognizes images. When you give it an image, it responds with a label for that image.

Diagram showing how an image classifier predicts the label "Giraffe" from an image of a giraffe.

You train an image classifier by showing it many examples of images you’ve already labeled. For example, you can train an image classifier to recognize safari animals by showing it a variety of photos of elephants, giraffes, lions, and so on.

Diagram showing how Create ML trains a model using collections of labeled images.

After the image classifier finishes training, assess its accuracy and if it performs well enough, save it as a Core ML model file. In your Xcode project, you import the Core ML model file to use your image classifier in your app.

Prepare Your Data

Start by preparing the data that you’ll use to train and evaluate the classifier. Create a training data set from about 80% of your images for each label. Create a testing data set from the remaining 20% of images. Make sure that any given image appears in only one data set or the other.

Next, organize your data on disk to be compatible with one of the MLImageClassifier.DataSource types. One way is to create two folders named Training Data and Testing Data. In each folder, you create subfolders using your labels as names. Then sort the images into the appropriate subfolders for each data set.

Diagram showing a folder called Training Data with subfolders that are named using the label corresponding to the class of images they contain. For example, all the cheetah images go into a subfolder called "Cheetah".

The exact label strings aren’t important, as long as they make sense to you. For example, you might use the label Cheetah for all the images of cheetahs. You don’t have to name the image files in any particular way or add metadata to them. You only need to put them into the folder with the right label.

Use at least 10 images per label for the training set, but more is always better. Also, balance the number of images for each label. Don’t use 10 images for Cheetah and 1000 images for Elephant.

The images can be in any format whose uniform type identifier conforms to public.image. This includes common formats like JPEG and PNG. The images don’t have to be the same size as each other, nor do they have to be any particular size. However, it’s best to use images that are at least 299x299 pixels. If possible, train with images that you collect in a way that’s similar to how you collect the images for prediction. That is, if your app classifies images taken from the device’s camera in an outdoor setting, train your model with images taken outdoors with a similar camera.

Provide images with variety. For example, use images that show animals from many different angles and in different lighting conditions. A classifier trained on nearly identical images for a given label tends to have poorer performance than one trained with a more diverse image set.

Create an Image Classifier Project

Use Create ML to create an image classification project. With Xcode open, Control-click Xcode’s Dock icon and choose Open Developer Tool > Create ML. Or, from the Xcode menu, choose Open Developer Tool > Create ML.

In Create ML, choose File > New to see the list of model templates. Select Image Classifier and click Next.

Screenshot of the Create ML project templates, with options for an Image Classifier, and a Sound Classifier.

Give the classifier a meaningful name. Change the default MyImageClassifier to be AnimalClassifier by changing the name in the UI. You can also add more information about the model, like the author and a short description.

Screenshot of the new project options window with text fields for Project Name, Author, License, and Description. The first text fields have the values, "Animal Classifier", "John Appleseed", and "no license provided", respectively. The description text field reads: "A machine learning model that has been trained to recognized safari animals."

Configure the Training Input

Drag your Training Data folder from the Finder onto the Training Data section in your project window.

Screenshot of the project window with the highlighted segmented control cell, labeled "Input",  with a Finder folder named "Training Data" being dragged and dropped onto its Training Data section.

Drag your Testng Data folder from the Finder onto the Testing Data section in your project window.

Screenshot of the project window with the highlighted segmented control cell, labeled "Input", with a Finder folder named "Testing Data" being dragged and dropped onto its Testing Data section.

Optionally, customize the training parameters or add image augmentations.

Screenshot of the Create ML app that highlights the training Parameters and Augmentation sections. The parameters section has one setting named "maximum iterations" which is set to 25. The augmentations section has 6 checkboxes (which are all unselected) named: "Add Noise", "Blur", "Crop", "Expose", "Flip", and "Rotate".

Train the Image Classifier

Click the Play button in the project window’s toolbar to start the training process. The image classifier shows its progress as it learns from your training images.

Screenshot of the project window with the highlighted segmented control cell, labeled "Training" during training. Below the segmented control shows a progress bar. The label above the pogress bar reads: "Step 1 of 2,  extracting features". The label below the pogress bar reads: "Processed 300 out of 685 images."

Understanding the Training and Validation Process

As a part of the training process, Create ML automatically splits your training data into a training set and a validation set. The model iteratively learns from the training set and with each iteration, the model checks its own accuracy with the validation set. Because Create ML splits the data randomly into these two sets, you might get a different result each time you train the model.

Assess the Classifier’s Accuracy

When training finishes, Create ML shows training and validation accuracy scores. These report how well the trained model classifies images from the corresponding sets. Because the model trained on these images, it typically does a good job classifying them. In this example, the trained model correctly identified all the images in the training set and 97% of the images in the validation set.

Screenshot of the project window with the highlighted segmented control cell, labeled "Training", displaying the value 100%. Below the segmented controls is a line graph that starts at 50% at one iteration, increasing to 75% at five iterations, and ending near 100% at twenty iterations. In the accuracy section below the graph is a table with columns: "class", "item count", "precision", and "recall". The first row has the values: "Cheetah", 54, and 100% for precision and recall.

The UI shows a graph of the model’s accuracy progress with each training iteration, as well as the precision and recall details for each image class.

The precision is the number of true-positives divided by the sum of true-positives and false-positives. Recall is the number of true-positives divided by the sum of true-positives and false-negatives.

Test the Classifier’s Performance

Create ML tests your trained model’s performance with your testing data set, which contains images the model has never seen before.

The model processes all of the images and makes a prediction for each. Create ML compares each prediction with the correct label and presents the overall testing accuracy in the UI.

Screenshot of the project window with a highlighted segmented control cell, labeled "Testing", displaying the value 95%. Below the segmented controls is a section named "Testing Data" with a Pop Up Button for selecting a different data set and a "Retest" button below it. In the accuracy section at the bottom of the project window is a table with columns: "class", "item count", "precision", and "recall". The third row has the values: "Flamingo", 7, 88% precision, and 100% recall.

If the evaluation performance isn’t good enough, you may need to retrain with more varied training data. For example, you can enable image augmentation options for an image classifier. For information about how to do a more detailed model evaluation, as well as strategies for improving model performance, see Improving Your Model’s Accuracy.

Save the Model

When your model performs well enough, save it so that you can use it in your app. Drag the model icon from the Output tab to a directory in Finder to save the model to a file in the .mlmodel format.

Screenshot of the project window with the highlighted segmented control cell, labeled "Output". The cursor is dragging an ML model file from the Output cell to the Desktop.

Add the Model to Your App

Now add the trained model to an app that uses Core ML. You can use this model to replace the one that comes with the Classifying Images with Vision and Core ML sample code project. If you do this, the sample app works exactly the same as before, except that it recognizes and classifies animals according to the labels you’ve defined.

Open the sample code project in Xcode, and drag your model file into the navigation pane. Once the model is part of the project, Xcode shows you the model metadata, along with other information, like the model class.

Screenshot showing the animal classifier model integrated into the image classifier sample code project.

To use the new model in code, you change only one line. The project instantiates the MobileNet model in exactly one place in the ImageClassificationViewController class.

let model = try VNCoreMLModel(for: MobileNet().model)

Change this one line to use the new model class instead:

let model = try VNCoreMLModel(for: AnimalClassifier().model)

The models are interchangeable because both take an image as input and output a label. After your substitution, the sample app classifies images as before, except it uses your model and its associated labels.

Automate the Process of Building an Image Classifier

You can use Create ML to train a useful image classifier with very little code or machine learning expertise, as described in the sections above. However, you can also use an MLImageClassifier instance to script the process of training a model. The general tasks are the same: prepare data, train a model, evaluate performance, and save the result to a Core ML model file. The difference is that you do everything programmatically.

For example, instead of letting Create ML partition your training data to test the model’s performance, you partition the data yourself into its own directory. Then, you initialize an MLImageClassifier.DataSource instance with that directory and provide it to the image classifier’s evaluation(on:) method, which generates an MLClassifierMetrics instance.

See Also

Computer Vision

struct MLImageClassifier

A model you train to classify images.

struct MLObjectDetector

A model you train to detect objects within an image.