Cassava Disease Classification with TensorFlow

In this tutorial, we will learn how to build an image classification model using the TensorFlow 2.x framework and the Cassava Leaf Disease dataset. The Cassava Leaf Disease dataset is a collection of images of cassava leaves, which are labeled with one of five possible classes of disease. The goal of our model is to predict the disease class of a given cassava leaf image.

We will start by exploring the dataset and preparing it for use in our model. Then, we will build and train a deep learning model using TensorFlow 2.x, and evaluate its performance on a test dataset. Finally, we will use the trained model to make predictions on new images.

Note: You can find the full notebook for this project here for reference.

Table of Contents

Follow the below parts to fully understand the tutorial:

1. Prerequisites
2. Exploring the Dataset
3. Building the Model
4. Training and Evaluating the Model
5. Fine-tuning the Model
6. Conclusion

Prerequisites

Before diving into this tutorial, make sure you have the following:

• Basic knowledge of Python programming language and machine learning concepts. See how to install Python if you have not installed Python on your Windows machine yet.
• TensorFlow 2.x installed on your machine.
• Jupyter Notebook or another similar tool to run and edit the provided code.
• The Cassava Leaf Disease dataset, which you can download from the official website. In this tutorial, we will just download the train.zip and test.zip files.

Exploring the Dataset

Before building our image classification model, let's first explore the Cassava Leaf Disease dataset. The dataset contains two directories: one for the training set and one for the test set. Each directory contains five subdirectories, one for each class of disease.

• train - the set of training images for 5 categories of disease
• train/cbb - the set of training images for class cbb
• train/cmd - the set of training images for class cmd
• train/cbsd - the set of training images for class cbsd
• train/cgm - the set of training images for class cgm
• train/healthy - the set of training images for class healthy
• test - the set of test images to run your trained model on

We can use the tensorflow.keras.preprocessing.image.ImageDataGenerator class to load the images and their labels from the dataset directories. This class can also perform data augmentation techniques such as random rotations, zooming, and flipping, which can improve the performance of our model.

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Define data generators for training and test sets
train_datagen = ImageDataGenerator(
  rescale=1./255,
  rotation_range=20,
  width_shift_range=0.2,
  height_shift_range=0.2,
  shear_range=0.2,
  zoom_range=0.2,
  horizontal_flip=True,
  fill_mode='nearest',
  validation_split=0.2
)
test_datagen = ImageDataGenerator(
  rescale=1./255
)

# Load the training and test data
train_data = train_datagen.flow_from_directory(
  'data/train/',
  target_size=(224, 224),
  batch_size=32,
  class_mode='categorical',
  shuffle=True,
  subset="training"
)
val_data = train_datagen.flow_from_directory(
  'data/train/',
  target_size=(224, 224),
  batch_size=32,
  class_mode='categorical',
  shuffle=True,
  subset="validation"
)
test_data = test_datagen.flow_from_directory(
  'data/test/',
  target_size=(224, 224),
  batch_size=32,
  class_mode='categorical'
)

In the above code, we have defined two ImageDataGenerator objects for the training and test sets, respectively. We have also loaded the training and test data using the flow_from_directory method of the ImageDataGenerator class, which reads the images from their respective directories and performs data augmentation (in the case of the training set).

Also, the original train dataset is splitted into train_data and val_data for validation purpose with the ratio of 80:20 as we set validation_split=0.2 inside the ImageDataGenerator method.

Remember to replace /path/to/training/set and /path/to/test/set to your training and testing dataset path.

Note that we have also rescaled the pixel values of the images to be in the range [0, 1] by dividing them by 255. This is a standard preprocessing step in deep learning to help improve the stability and performance of the model.

Building the Model

Now that we have loaded and preprocessed the dataset, let's build our image classification model using TensorFlow. We will use a pre-trained convolutional neural network (CNN) as the base of our model, and then add some fully connected layers on top to perform the classification.

In this tutorial, we will use the VGG16 CNN architecture as our base model, which has shown strong performance on image classification tasks. We will download the pre-trained weights for this model using the tensorflow.keras.applications.vgg16 module:

from tensorflow.keras.applications.vgg16 import VGG16

# Load the pre-trained VGG16 model without the top (fully connected) layers
base_model = VGG16(
  weights='imagenet',
  include_top=False,
  input_shape=(224, 224, 3)
)

In the above code, we have loaded the VGG16 model with the pre-trained ImageNet weights, but without the top (fully connected) layers. We have also specified the input shape of our images to be (224, 224, 3), which is the default input shape for the VGG16 model.

Next, we will add some fully connected layers on top of the base model to perform the classification:

from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.models import Model

# Add our own fully connected layers on top of the base model
x = base_model.output
x = Flatten()(x)
x = Dense(512, activation='relu')(x)
x = Dense(256, activation='relu')(x)
predictions = Dense(5, activation='softmax')(x)

# Create our final model
model = Model(inputs=base_model.input, outputs=predictions)

In the above code, we have added a flatten layer and two fully connected layers with ReLU activation functions on top of the base model. We have also added a final output layer with 5 units (one for each class of disease) and a softmax activation function, which will output the predicted probabilities for each class.

Finally, we create our final model by specifying the input and output layers of the model using the Model class from Keras.

Training and Evaluating the Model

Now that we have built our model, let's train and evaluate it using the preprocessed Cassava Leaf Disease dataset.

First, we need to compile the model and specify the loss function, optimizer, and metrics for evaluation:

from tensorflow.keras.optimizers import Adam

opt = Adam(learning_rate=5e-5)

model.compile(
  loss='categorical_crossentropy',
  optimizer=opt,
  metrics=['accuracy']
)

In the above code, we have specified the categorical cross-entropy loss function, the Adam optimizer with learning_rate set to 5e-5, and the accuracy metric for evaluation.

Next, we can train the model on our preprocessed training dataset:

history = model.fit(
  train_data,
  validation_data=val_data,
  epochs=10,
  verbose=1
)

In the above code, we have used the fit method of our model to train it on the preprocessed training dataset for 10 epochs, with the preprocessed test dataset as the validation data. The verbose argument is set to 1 to display the progress during training.

The training progress will be like this, just wait for it to complete.

At the last epoch (epoch 10), we get the results:

Epoch 10/10
142/142 [==============================] - 93s 656ms/step - loss: 0.4024 - accuracy: 0.8655 - val_loss: 0.5313 - val_accuracy: 0.8308

After training, we can evaluate our model on the preprocessed test dataset:

val_loss, val_acc = model.evaluate(val_data, verbose=0)
print('Validation accuracy: ', test_acc)

Output:

Test accuracy:  0.8246235847473145

In the above code, we have used the evaluate method of our model to evaluate it on the preprocessed test dataset and print the test accuracy of ~0.825, pretty good for this model.

We can also plot the training and validation accuracy and loss over epochs using the matplotlib library:

import matplotlib.pyplot as plt

acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
loss = history.history['loss']
val_loss = history.history['val_loss']

epochs_range = range(1, len(acc) + 1)

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')
plt.plot(epochs_range, val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')

plt.show()

In the above code, we have plotted the training and validation accuracy and loss over epochs using two subplots. The figure method is used to specify the figure size, and the subplot method is used to create two subplots side-by-side. The plot method is used to plot the accuracy and loss for the training and validation datasets, and the legend and title methods are used to add labels to the plots. Finally, the show method is used to display the plots.

Fine-tuning the Model

After training and evaluating our model, we can try to improve its performance by fine-tuning it on the preprocessed Cassava Leaf Disease dataset.

Fine-tuning involves unfreezing some or all of the layers of our model and retraining it on the preprocessed dataset with a lower learning rate, so that the model can learn more specific features for our dataset.

First, let's unfreeze the last few layers of our model and freeze the rest:

base_model.trainable = True

for layer in base_model.layers[:-10]:
    layer.trainable = False

In the above code, we have set the trainable attribute of the base model to True and frozen all but the last 10 layers of the model by setting their trainable attribute to False.

Next, we need to compile the model again and specify a lower learning rate for fine-tuning:

model.compile(
  loss='categorical_crossentropy',
  optimizer=opt,
  metrics=['accuracy']
)

In the above code, we have used the Adam optimizer with a learning rate of 1e-5 for fine-tuning.

We can then continue training the model on the preprocessed training dataset for a few more epochs:

history_fine = model.fit(
  train_data,
  validation_data=val_data,
  epochs=5,
  initial_epoch=history.epoch[-1],
  verbose=1
)

In the above code, we have used the fit method of our model to continue training it on the preprocessed training dataset for 5 more epochs, starting from the last epoch of the previous training session.

After fine-tuning, we can evaluate our model on the preprocessed test dataset again and see if its performance has improved:

val_loss_fine, val_acc_fine = model.evaluate(val_data, verbose=0)
print('Fine-tuned validation accuracy: ', val_acc_fine)

In the above code, we have used the evaluate method of our model to evaluate it on the preprocessed test dataset again and print the fine-tuned test accuracy.

Conclusion

In this tutorial, we have explored how to build an image classification model using TensorFlow and the Cassava Leaf Disease dataset. We have started by preprocessing the dataset, then building and training a transfer learning model based on the VGG16 architecture.

We have also discussed how to fine-tune the model by unfreezing some of its layers and retraining it on the preprocessed dataset with a lower learning rate.

With this tutorial, you should have a good understanding of how to build and fine-tune image classification models using TensorFlow and transfer learning, as well as how to apply these techniques to a real-world dataset.

We hope this tutorial has been helpful to you and wish you good luck with your future machine learning projects!