Step-by-Step Guide to Training a Custom Image Classifier from Scratch

Training machine learning models to accurately classify images is a powerful capability with countless applications, from identifying plant species to detecting manufacturing defects to diagnosing medical conditions. And thanks to the advent of deep learning and convolutional neural networks (CNNs), the accuracy of image classifiers has increased dramatically in recent years.

However, many of the pre-trained state-of-the-art image classification models available today, like those trained on the ImageNet dataset, are focused on broad categories like identifying 1000 common objects. What if you need a model to classify a more niche image dataset, such as different types of cells under a microscope, or damage to car fenders? In those cases, you‘ll likely need to train a custom image classifier from scratch on your own dataset.

In this post, we‘ll walk through the key steps to building a robust image classification model for any custom dataset. While we‘ll provide code examples in Python with the Keras deep learning library, the high-level concepts are relevant regardless of your AI toolkit of choice.

The general process for training an image classifier is:

1. Collect and annotate a dataset of images
2. Preprocess the images
3. Split data into training, validation and test sets
4. Select a model architecture
5. Train the model
6. Evaluate performance on the test set
7. Iterate and refine the model

Let‘s dive into each of these in more detail.

1. Collecting and Annotating a Dataset

Having a high-quality, representative dataset is critical to the success of any machine learning model. For image classification, you‘ll need a labeled dataset, with each image assigned to one of your pre-defined categories or classes.

Collecting and annotating the dataset is often the most time-consuming part of an applied ML project. A few best practices:

• Collect a balanced dataset with roughly equal numbers of examples for each class. Imbalanced datasets can lead to a model that performs well on the overrepresented classes but poorly on others.

• Aim to collect at least 1000 example images per class if possible, though you may be able to get away with less for simpler problems. In general, more data leads to better performance.

• Ensure images are labeled consistently. Create clear annotation guidelines defining what belongs to each class. If using multiple human annotators, have them label a small subset to calculate inter-annotator agreement.

• Include a diverse variety of images in your dataset that are representative of what the model will encounter in the real world. Different angles, lighting conditions, backgrounds, etc.

If you don‘t already have an image dataset, there are a number of public repositories that can be good sources, including Kaggle Datasets and Google Dataset Search. Or for more customized datasets, you can collect images from sources like Google Images, Flickr, or by taking your own photographs.

Annotating a dataset can be tedious. A number of tools exist to streamline the process by providing labeling interfaces for teams of annotators, such as Labelbox, V7, and Supervisely.

2. Preprocessing the Images

With a dataset in hand, the next step is to preprocess the images to get them ready for training a model. The goal is to transform the raw images into a format that can be fed directly into the model.

Some common preprocessing steps for images include:

• Resizing: CNNs typically require images to be a fixed, square size, commonly 224×224 or 299×299 pixels. Aspect ratio is usually not preserved.

• Normalization: It‘s common to scale the pixel values from the original range of [0, 255] to [0, 1], which tends to make training converge faster. For some architectures like ResNet, it‘s also common to subtract the mean pixel value from each channel.

• Data augmentation: To increase the size and diversity of the training set, it‘s useful to programmatically generate new images by randomly transforming existing ones. Common augmentations include flips, rotations, crops, and color jittering.

Here‘s an example in Keras of resizing and normalizing an image before feeding it to a model:

from tensorflow.keras.preprocessing import image

img = image.load_img("image.jpg", target_size=(224, 224))
img_array = image.img_to_array(img)
img_array = img_array / 255.0
img_array = np.expand_dims(img_array, axis=0)  # add batch dimension

And an example of using Keras‘s ImageDataGenerator to create a data augmentation pipeline:

from tensorflow.keras.preprocessing.image import ImageDataGenerator

datagen = ImageDataGenerator(
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True)

datagen.fit(x_train)

3. Splitting the Dataset

With the images preprocessed, the next step is to split your dataset into three subsets:

• Training set: The largest portion of the data, used to train the model and adjust its weights. Usually around 70-80% of the total data.

• Validation set: A smaller subset, usually 10-15% of the data, used to tune hyperparameters and make model architecture decisions. Performance on the validation set is used as a proxy for generalization performance.

• Test set: The final 10-20% of data that is only used once at the very end to evaluate the model‘s performance on unseen data. It‘s important not to allow any information from the test set to "leak" into training the model.

The train_test_split function from scikit-learn is the easiest way to split a dataset:

from sklearn.model_selection import train_test_split

x_train, x_test, y_train, y_test = train_test_split(
    x_dataset, y_dataset, test_size=0.2, stratify=y_dataset)

x_train, x_val, y_train, y_val = train_test_split(
    x_train, y_train, test_size=0.15, stratify=y_train)

Setting stratify ensures the proportion of examples in each class remains consistent between splits.

4. Selecting a Model Architecture

Now we‘re ready to define the architecture of our image classifier. CNNs have become the go-to models for image-related tasks. They work by learning a hierarchy of feature detectors made up of convolution and pooling layers that capture patterns in the image data at different spatial scales.

Here‘s an example of a basic CNN defined with the Keras Sequential API:

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPool2D, Flatten, Dense 

model = Sequential([
    Conv2D(32, 3, activation="relu", input_shape=(224, 224, 3)),
    MaxPool2D(pool_size=2),
    Conv2D(64, 3, activation="relu"),
    MaxPool2D(),
    Conv2D(128, 3, activation="relu"),
    MaxPool2D(),
    Flatten(),
    Dense(64, activation="relu"),
    Dense(10, activation="softmax")
])

This toy example has three convolutional layers that learn 32, 64, and 128 filters respectively, with ReLU activations, each followed by max pooling layers to reduce spatial dimensions. The final feature maps are flattened and passed to two fully-connected dense layers to generate class probabilities.

In practice, it‘s common to leverage transfer learning by starting with a pre-trained CNN like VGG, Inception, or ResNet as a base model. The weights of the pre-trained layers are frozen while new final classification layers are added on top and trained on your custom dataset. Using a pre-trained model allows you to achieve higher accuracy with less data and compute.

Here‘s an example of using a pre-trained ResNet50 model as a base in Keras:

from tensorflow.keras.applications.resnet50 import ResNet50

base = ResNet50(include_top=False, input_shape=(224,224,3))
base.trainable = False  # freeze base layers

x = base.output
x = GlobalAveragePooling2D()(x)
x = Dense(1024, activation=‘relu‘)(x) 
predictions = Dense(10, activation=‘softmax‘)(x)

model = Model(inputs=base.input, outputs=predictions)

5. Training the Model

With our model architecture defined, the next step is to train its weights on our dataset. This involves iteratively showing the model batches of images, having it make predictions, and adjusting the weights to minimize the difference between predictions and ground truth labels.

Before training, we need to compile the model by specifying an optimizer like Adam or SGD to adjust the weights, a loss function like categorical cross-entropy to measure the difference between predicted and actual labels, and optional metrics like accuracy to monitor during training:

model.compile(optimizer=‘adam‘,
              loss=‘categorical_crossentropy‘,
              metrics=[‘accuracy‘])

Then we kick off training by calling fit() on the model, passing in our training images and labels, and specifying the number of epochs (full passes through the training set) and optionally batch size:

history = model.fit(x_train, y_train, 
                    epochs=50, batch_size=32,
                    validation_data=(x_val, y_val))

It‘s important to monitor your model‘s performance on the validation set during training. You want to see the validation loss decreasing and accuracy increasing over time, but stop training if you see it begin to stagnate or go back up, as this is a sign of overfitting.

Early stopping is a useful callback for automatically halting training when the model‘s performance on the validation set stops improving:

from tensorflow.keras.callbacks import EarlyStopping

early_stopping = EarlyStopping(monitor=‘val_loss‘, 
                               patience=5,
                               restore_best_weights=True)

history = model.fit(x_train, y_train,
                    epochs=50, batch_size=32,
                    validation_data=(x_val, y_val),
                    callbacks=[early_stopping])                      

This will stop training if the validation loss doesn‘t improve for 5 consecutive epochs, and restore the weights from the best-performing epoch.

6. Evaluating Test Set Performance

Once you have a trained model you‘re happy with, it‘s time to get an unbiased estimate of how it will perform on unseen data by evaluating it on the held-out test set.

The simplest metric is overall accuracy – the percentage of images in the test set that the model classified correctly:

loss, accuracy = model.evaluate(x_test, y_test)
print(f"Test accuracy: {accuracy:.3f}")

However, accuracy can be misleading, especially for imbalanced datasets. It‘s often more informative to look at a classification report showing precision, recall and f1 score for each class:

from sklearn.metrics import classification_report

y_pred = np.argmax(model.predict(x_test), axis=-1)
y_test = np.argmax(y_test, axis=-1)
print(classification_report(y_test, y_pred))

A confusion matrix can also provide insight into which classes are most often confused:

from sklearn.metrics import confusion_matrix

print(confusion_matrix(y_test, y_pred))

Examining example images that the model misclassified can be illuminating and suggest improvements to the model or dataset.

7. Iterating and Refining

Rarely is a model optimal after only one iteration. Especially for challenging image classification tasks, multiple rounds of refinement are often necessary to achieve the desired level of accuracy.

Some avenues to explore if your model isn‘t performing as well as you‘d like:

• Collect more training data, especially for underperforming classes
• Optimize preprocessing and augmentation steps
• Try a different model architecture or add more layers
• Fine-tune hyperparameters like learning rate, batch size, etc.
• Unfreeze more layers of the pre-trained base model during training
• Use techniques like learning rate scheduling or dropout to combat overfitting

Each time you make changes, retrain the model from scratch and evaluate its performance on the test set to gauge progress.

Summary

Let‘s recap the key steps for training an image classifier on a custom dataset:

1. Collect a labeled dataset with a variety of representative examples for each class
2. Preprocess the images by resizing, normalizing, and optionally augmenting the data
3. Split into train, validation, and test sets
4. Select a model architecture, often based on a pre-trained CNN
5. Train the model, monitoring performance on the validation set
6. Evaluate on the test set with metrics like accuracy, precision, recall
7. Iterate on the data, model, and training process to improve performance

While there‘s no universal recipe for the perfect image classifier, following these steps will put you on the path to training a highly accurate model for your specific domain. Thanks to the power of transfer learning, you can achieve impressive results even with a relatively small dataset.

The beauty of deep learning is that these same techniques can be readily applied to any image classification task, from identifying plant species to detecting cancerous cells to classifying types of damage to car fenders. We‘re excited to see what you‘ll build!