Title: Training a Custom Image Classification Network for OAK-D - PyImageSearch

URL Source: https://pyimagesearch.com/2023/02/06/training-a-custom-image-classification-network-for-oak-d/

Published Time: 2023-02-06T14:00:00+00:00

Markdown Content:
* * *

**Table of Contents**
---------------------

*   [Training a Custom Image Classification Network for OAK-D](#h2BPTitle)

*   [Configuring Your Development Environment](#h3Configuring)
*   [Having Problems Configuring Your Development Environment?](#h3Problems)
*   [Project Structure](#h3Structure)
*   [The Vegetable Image Dataset](#h3Vegetable)

*   [Data Analysis](#h4Analysis)

*   [Configuring the Prerequisites](#h3Prerequisites)
*   [Defining the Utilities](#h3Defining)
*   [The Image Classification Network](#h3Network)
*   [Training](#h3Training)

*   [Accuracy and Loss Plots: Training vs. Validation](#h4Accuracy)

*   [Testing](#h3Testing)

*   [Classification Report](#h4Report)
*   [Sample Test Set Predictions](#h4Sample)

*   [Summary](#h2Summary)

*   [Citation Information](#h3Citation)

[](#h3Citation)

* * *

In this tutorial, you will learn to train a custom image classification network for OAK-D using the TensorFlow framework. Furthermore, this tutorial aims to develop an image classification model that can learn to classify one of the 15 vegetables (e.g., tomato, brinjal, and bottle gourd).

[![Image 1](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20393'%3E%3C/svg%3E)](https://pyimagesearch.com/wp-content/uploads/2023/02/cicn-oak-d-featured.png)

If you are a regular PyImageSearch reader and have even basic knowledge of Deep Learning in Computer Vision, then this tutorial should be easy to understand. Furthermore, this tutorial acts as a foundation for the following tutorial, where we learn to deploy this trained image classification model on [OAK-D](https://shop.luxonis.com/?aff=10).

This is the 3rd lesson in our 4-part series on **OAK 101**:

1.  [_Introduction to OpenCV AI Kit (OAK)_](https://pyimg.co/dus4w)
2.  [_OAK-D: Understanding and Running Neural Network Inference with DepthAI API_](https://pyimg.co/8ynbk)
3.  [**_Training a Custom Image Classification Network for OAK-D_**](https://pyimg.co/ab1dj) **(today’s tutorial)**
4.  [**_Deploying a Custom Image Classifier on an OAK-D_**](https://pyimg.co/tnx0u)

**To learn how to train an image classification network for OAK-D,** **_just keep reading._**

![Image 2](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20128%20128'%3E%3C/svg%3E)

#### Looking for the source code to this post?

[Jump Right To The Downloads Section](#download-the-code)

* * *

[**Training a Custom Image Classification Network for OAK-D**](#TOC)
--------------------------------------------------------------------

Before we start data loading, analysis, and training the classification network on the data, we must carefully pick the suitable classification architecture as it would finally be deployed on the [OAK](https://shop.luxonis.com/?aff=10). Although [OAK](https://shop.luxonis.com/?aff=10) can process 4 trillion operations per second, it is still an edge device. It has its bottlenecks, so we cannot expect it to run a heavyweight model like EfficientNet-B7, achieving real-time inference speed.

In today’s experiment, we would use MobileNetV2 as a feature extractor pre-trained on ImageNet data. MobileNetV2 is a classification model developed by Google. It provides real-time classification capabilities under computing constraints in edge devices. Hence, the obvious choice for this tutorial.

We would add custom layers on top of the MobileNetV2 pre-trained feature extractor, and these layers would be trained to perform image classification while the feature extractor (base model) weights would be frozen (non-trainable).

In short, we would perform Transfer learning that takes features learned on one problem and leverage them on a new, similar problem. For example, features from the base model that has learned to identify 1000-ImageNet classes would be helpful to give a head start to a model meant to identify vegetables.

* * *

### [**Configuring Your Development Environment**](#TOC)

To follow this guide, you need to have `tensorflow`, `scikit-learn`, and `matplotlib` installed on your system.

Luckily, all these libraries are pip-installable:

$ pip install tensorflow==2.9.2
$ pip install scikit-learn==0.23.2
$ pip install matplotlib==3.3.2

* * *

### [**Having Problems Configuring Your Development Environment?**](#TOC)

[![Image 3](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20500%20334'%3E%3C/svg%3E)](https://pyimagesearch.com/pyimagesearch-university/)

**Figure 1:** Having trouble configuring your dev environment? Want access to pre-configured Jupyter Notebooks running on Google Colab? Be sure to join [PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/) — you’ll be up and running with this tutorial in a matter of minutes.

All that said, are you:

*   Short on time?
*   Learning on your employer’s administratively locked system?
*   Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
*   **Ready to run the code** **_right now_** **on your Windows, macOS, or Linux system?**

Then join [PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/) today!

**Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides that are** **_pre-configured_** **to run on Google Colab’s ecosystem right in your web browser!** No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

* * *

### [**Project Structure**](#TOC)

We first need to review our project directory structure.

Start by accessing the **_“Downloads”_** section of this tutorial to retrieve the source code and example images.

From there, take a look at the directory structure:

├── output
│   ├── accuracy\_loss\_plot.png
│   ├── test\_prediction\_images.png
│   └── vegetable\_classifier
│       ├── keras\_metadata.pb
│       ├── saved\_model.pb
│       └── variables
│           ├── variables.data-00000-of-00001
│           └── variables.index
├── pyimagesearch
│   ├── \_\_init\_\_.py
│   ├── config.py
│   ├── network.py
│   └── utils.py
├── test.py
└── train.py

In the `pyimagesearch` directory, we have the following:

*   `config.py`: The configuration file for the experiments
*   `utils.py`: The utilities for the image classification training
*   `network.py`: Houses the implementation of the end-to-end classification model (i.e., the base MobileNetV2 model) and the classification layers

In the core directory, we have 3 files:

*   `requirements.txt`: The Python packages that are required for this tutorial
*   `train.py`: It contains the script for training the image classification model
*   `test.py`: Script that runs an inference of the trained classification model on test data

* * *

### [**The Vegetable Image Dataset**](#TOC)

Before we go ahead and load the data, it’s good to look at what you’ll exactly be working with! The [Vegetable Image Dataset](https://www.kaggle.com/datasets/misrakahmed/vegetable-image-dataset) comprises 15 vegetable images, with 224×224 color images of 21,000 different vegetables from 15 categories and 1400 images per category. The training set has 15,000 images, and the validation and test sets have 3,000 images each. You can double-check this later when you have loaded your data!

The good news is that the image dimensions are already what our pretrained MobileNetV2 model expects (i.e., 224×224), so we do not need to worry about resizing, but we still do that as a sanity check step.

Before you load this data, you need to download it from Kaggle. We show how to download this dataset in our Google Colab in very few steps, so check that out!

**Figure 2** shows 9 of the 15 vegetable image categories from the dataset. At first glance, the dataset has a good amount of variation within the image. For example, brinjal is present at multiple scales in the image. But we would still apply data augmentation to ensure the model doesn’t overfit and generalize well on the test dataset.

[![Image 4](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20495%20500'%3E%3C/svg%3E)](https://pyimagesearch.com/wp-content/uploads/2023/02/vegetable-images-visualise.png)

**Figure 2:** Analyzing a few categories from the vegetable image classification dataset.

* * *

#### [**Data Analysis**](#TOC)

When working with data, especially supervised learning, it is often a best practice to check data imbalance. For example, **Figure 3** shows the percentage of each vegetable (category) across different splits (`train`, `valid`, and `test`), and it clearly shows that the data is perfectly balanced across all splits.

[![Image 5](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20271%20488'%3E%3C/svg%3E)](https://pyimagesearch.com/wp-content/uploads/2023/02/data-analysis.png)

**Figure 3:** Percentage of each vegetable category images across the `train`, `valid`, and `test` sets.

* * *

### [**Configuring the Prerequisites**](#TOC)

Before we start our implementation, let’s go over the configuration pipeline of our project. We will move on to the `config.py` script located in the `pyimagesearch` directory.

\# import the necessary packages
import os

# define the base path, paths to separate train, validation, and test splits
BASE\_DATASET\_PATH = "datasets/Vegetable Images"
TRAIN\_DATA\_PATH = os.path.join(BASE\_DATASET\_PATH, "train")
VALID\_DATA\_PATH = os.path.join(BASE\_DATASET\_PATH, "validation")
TEST\_DATA\_PATH = os.path.join(BASE\_DATASET\_PATH, "test")
OUTPUT\_PATH = "output"

# define the image size and the batch size of the dataset
IMAGE\_SIZE = 224
BATCH\_SIZE = 32

# number of channels, 1 for gray scale and 3 for color images
CHANNELS = 3

# define the classifier network learning rate
LR\_INIT = 0.0001

# number of epochs for training
NUM\_EPOCHS = 20

# number of categories/classes in the dataset
N\_CLASSES = 15

# define paths to store training plots, testing prediction and trained model
ACCURACY\_LOSS\_PLOT\_PATH = os.path.join("output", "accuracy\_loss\_plot.png")
TRAINED\_MODEL\_PATH = os.path.join("output", "vegetable\_classifier")
TEST\_PREDICTION\_OUTPUT = os.path.join("output", "test\_prediction\_images.png")

The `config.py` helps set up the file paths and parameters for the dataset and the classifier network that will train and test the image classification model.

We start by importing the `os` package to help us create datasets and output paths. Then, from **Lines 5-9**, we set the base path for the dataset and define the paths for the train, validation, and test splits of the dataset and the output path.

On **Lines 12 and 13**, we define the dataset’s image and batch sizes. Then determine the number of channels set to 3 for color images, the classifier network learning rate, the number of epochs for training, and the number of categories/classes in the dataset (**Lines 16-25**).

Finally, on **Lines 28-30**, we define the paths to store the training plots, testing predictions, and the trained model.

* * *

### [**Defining the Utilities**](#TOC)

Now that the configuration pipeline has been defined, we can define the training utilities. The `utils.py` script defines several functions:

*   prepare the dataset for training
*   apply data augmentation and normalization of the input images
*   define early stopping callbacks to stop the training if the validation loss does not improve

We will use these functions in the training process to make the model more robust.

from pyimagesearch import config
import tensorflow as tf

def prepare\_batch\_dataset(data\_path, img\_size, batch\_size, shuffle=True):
	return tf.keras.preprocessing.image\_dataset\_from\_directory(
		data\_path,
		image\_size=(img\_size, img\_size),
		shuffle=shuffle,
		batch\_size=batch\_size
	)

On **Lines 1 and 2**, we import `config` from the `pyimagesearch` module and the TensorFlow library.

From **Lines 4-10**, we define the function `prepare_batch_dataset` that takes in the following:

*   path of the dataset
*   image size
*   batch size
*   boolean value for shuffle

It then uses `tf.keras.preprocessing.image_dataset_from_directory` to load the images from the directory, resizes them to the given image size, shuffles them if shuffle is `True`, and returns the dataset in the form of batches.

def callbacks():
	# build an early stopping callback and return it
	callbacks = \[
		tf.keras.callbacks.EarlyStopping(
			monitor="val\_loss",
			min\_delta=0,
			patience=2,
			mode="auto",
		),
	\]
	return callbacks

On **Lines 12-22**, the function `callbacks` defines an early stopping callback and returns it. The callback monitors the validation loss and stops the training if it does not improve after two epochs.

def normalize\_layer(factor=1./127.5):
	# return a normalization layer
	return tf.keras.layers.Rescaling(factor, offset=-1)

def augmentation():
	# build a sequential model with augmentations
	data\_aug = tf.keras.Sequential(
		\[
			tf.keras.layers.RandomFlip("horizontal"),
			tf.keras.layers.RandomRotation(0.1),
			tf.keras.layers.RandomZoom(0.1),
		\]
	)
	return data\_aug

Then on **Lines 24-26**, the function `normalize_layer` returns a normalization layer. It normalizes the input data by dividing by `127.5` and subtracting `1`. In short, it normalizes or scales the input images from `[0, 255]` to `[-1, 1]`, which is what the MobileNetV2 network expects.

Finally, from **Lines 28-37**, the function `augmentation` creates a sequential model with image augmentations: random horizontal flipping, rotation, and zoom. The augmentation helps avoid overfitting and allows the model to generalize better.

* * *

### [**The Image Classification Network**](#TOC)

Now with the utilities defined, we head towards coding the most exciting part (i.e., the image classification network). If you haven’t coded an image classification network before, the section is definitely for you! So let’s get straight into the code.

\# import the necessary packages
from pyimagesearch.utils import augmentation
from pyimagesearch.utils import normalize\_layer
import tensorflow as tf

class MobileNet:
	@staticmethod
	def build(width, height, depth, classes):
		# initialize the pretrained mobilenet feature extractor model
		# and set the base model layers to non-trainable
		base\_model = tf.keras.applications.MobileNetV2(
			input\_shape=(height, width, depth),
			include\_top=False,
			weights="imagenet",
		)

We start by importing the necessary packages on **Lines 2-4**, including the `augmentation` and `normalize_layer` modules from the `pyimagesearch` library and TensorFlow as `tf`.

On **Line 6**, a class called `MobileNet` is defined. Within this class, a static method called `build` takes four arguments: `width`, `height`, `depth`, and `classes`.

From **Lines 11-15**, we define the `base_model` variable, which is the MobileNetV2 architecture with the input shape set to (`height`, `width`, `depth`), `include_top` set to `False`, and weights set to `imagenet`.

		base\_model.trainable = False

Then, on **Line 16**, the base model’s trainable attribute is set to `False` so that the model’s base layers are not trainable.

		# define the input to the classification network with
		# image dimensions
		inputs = tf.keras.Input(shape=(height, width, depth))

On **Line 20**, the input to the model is defined as a tensor with shape (height, width, depth) and named inputs.

		# apply augmentation to the inputs and normalize the batch
		# images to \[-1, 1\] as expected by MobileNet
		x = augmentation()(inputs)
		x = normalize\_layer()(x)

Then on **Lines 24 and 25**, an augmentation layer is applied to the inputs, and the normalized layer is applied to the augmented inputs, which scales the input pixels to `[-1, 1]`.

		# pass the normalized and augmented images to base model,
		# average the 7x7x1280 into a 1280 vector per image,
		# add dropout as a regularizer with dropout rate of 0.2
		x = base\_model(x, training=False)
		x = tf.keras.layers.GlobalAveragePooling2D()(x)
		x = tf.keras.layers.Dropout(0.2)(x)

The normalized and augmented inputs are then passed through the base model on **Line 30**. Next, the output of the base model is passed through a global average pooling layer (on **Line 31**). Then, on **Line 32**, a dropout layer is applied with a dropout rate of `0.2`.

		# apply dense layer to convert the feature vector into
		# a prediction of classes per image
		outputs = tf.keras.layers.Dense(classes)(x)

		# build the keras Model by passing input and output of the
		# model and return the model
		model = tf.keras.Model(inputs, outputs)
		return model

On **Line 36**, a dense layer with `classes` as the number of outputs is applied to the dropout layer’s output. This dense layer will convert the feature vector into predictions of classes per image.

Finally, on **Lines 40 and 41**, a Keras Model is built by passing the inputs and outputs to the Model class, and the model is returned to the calling function.

* * *

### [**Training**](#TOC)

With the configurations, helper functions, and, notably, the image classification network implemented, we can finally get into the code walkthrough of training the classification neural network.

\# USAGE
# python train.py

# import the necessary packages
from pyimagesearch.utils import prepare\_batch\_dataset
from pyimagesearch.utils import callbacks
from pyimagesearch import config
from pyimagesearch.network import MobileNet
from matplotlib import pyplot as plt
import tensorflow as tf
import os

# build the training and validation dataset pipeline
print("\[INFO\] building the training and validation dataset...")
train\_ds = prepare\_batch\_dataset(
	config.TRAIN\_DATA\_PATH, config.IMAGE\_SIZE, config.BATCH\_SIZE
)

val\_ds = prepare\_batch\_dataset(
	config.VALID\_DATA\_PATH, config.IMAGE\_SIZE, config.BATCH\_SIZE
)

The script starts by importing necessary packages such as the `prepare_batch_dataset`, `callbacks`, `config`, `MobileNet`, and `pyplot` from `matplotlib`. The script also imports TensorFlow as `tf` and `os` (**Lines 5-11**).

On **Lines 15-21**, a pipeline for the training and validation datasets is built. The function `prepare_batch_dataset` is used to create the dataset by passing in the path of the training data, image size, and batch size, as specified in the config file. The same is done for the validation dataset.

\# build the output path if not already exists
if not os.path.exists(config.OUTPUT\_PATH):
	os.makedirs(config.OUTPUT\_PATH)

# initialize the callbacks, model, optimizer and loss
print("\[INFO\] compiling model...")
callbacks = callbacks()
model = MobileNet.build(
	width=config.IMAGE\_SIZE,
	height=config.IMAGE\_SIZE,
	depth=config.CHANNELS,
	classes=config.N\_CLASSES
)
optimizer = tf.keras.optimizers.Adam(learning\_rate=config.LR\_INIT)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from\_logits=True)

# compile the model
model.compile(
	optimizer=optimizer,
	loss=loss,
	metrics=\["accuracy"\]
)

Then, on **Lines 24 and 25**, the output path is created if it doesn’t already exist.

From **Lines 29-44**,

*   The model is compiled by initializing the `callbacks`, `model`, `optimizer`, and `loss`.
*   The model is built by calling the build method of the MobileNet class and passing in the width, height, depth, and number of classes as specified in the config file.
*   The `Adam` optimizer is used with the initial learning rate specified in the config file, and the loss function used is sparse categorical cross-entropy.

\# evaluate the model initially
(initial\_loss, initial\_accuracy) = model.evaluate(val\_ds)
print("initial loss: {:.2f}".format(initial\_loss))
print("initial accuracy: {:.2f}".format(initial\_accuracy))

# train the image classification network
print("\[INFO\] training network...")
history = model.fit(
	train\_ds,
	epochs=config.NUM\_EPOCHS,
	validation\_data=val\_ds,
	callbacks=callbacks,
)

The initial loss and accuracy of the model are evaluated on the validation dataset on **Line 47**.

From **Lines 53-58**, the model is trained by calling the `fit` method on the model with the training dataset and passing in the number of epochs specified in the config file and the validation dataset. The callbacks are passed in as well during training.

\# save the model to disk
print("\[INFO\] serializing network...")
model.save(config.TRAINED\_MODEL\_PATH)

# save the training loss and accuracy plot
plt.style.use("ggplot")
acc = history.history\['accuracy'\]
val\_acc = history.history\['val\_accuracy'\]

loss = history.history\['loss'\]
val\_loss = history.history\['val\_loss'\]

plt.figure(figsize=(8, 8))
plt.subplot(2, 1, 1)
plt.plot(acc, label='Training Accuracy')
plt.plot(val\_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.ylabel('Accuracy')
plt.ylim(\[min(plt.ylim()),1\])
plt.title('Training and Validation Accuracy')

plt.subplot(2, 1, 2)
plt.plot(loss, label='Training Loss')
plt.plot(val\_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.ylabel('Cross Entropy')
plt.ylim(\[0,1.0\])
plt.title('Training and Validation Loss')
plt.xlabel('epoch')
plt.savefig(config.ACCURACY\_LOSS\_PLOT\_PATH)

We save the model to the disk on **Line 62** using the `save` method and the path specified in the config file.

Finally, from **Lines 65-89**, the training loss and accuracy are plotted using `matplotlib` and saved to a file using the path specified in the config file.

* * *

#### [**Accuracy and Loss Plots: Training vs. Validation**](#TOC)

From **Figure 4**, you can see that the validation loss and validation accuracy are almost in sync with the training loss and training accuracy. It’s a good validation that the model is balanced: the validation loss is decreasing and not increasing, and there is not much gap between training and validation accuracy. In fact, after `10` epochs, the training and validation curves coincide.

Therefore, your model’s generalization capability is good, but the test set could be challenging, and the model might not perform well.

[![Image 6](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20504%20496'%3E%3C/svg%3E)](https://pyimagesearch.com/wp-content/uploads/2023/02/training-validation-accuracy-loss-plots.png)

**Figure 4:** Training vs. Validation Accuracy and Loss Plots.

* * *

### [**Testing**](#TOC)

Finally, it’s time to put our trained model to the test and see how well it performs on the test set. In this last section of the tutorial, we will run inference on the test data and plot the classification report and a few test images with their respective predicted class labels.

\# USAGE
# python test.py

# import the necessary packages
from pyimagesearch.utils import prepare\_batch\_dataset
from pyimagesearch.utils import callbacks
from pyimagesearch import config
from pyimagesearch.network import MobileNet
from sklearn.metrics import classification\_report
from matplotlib import pyplot as plt
import tensorflow as tf
import numpy as np
import os

# build the test dataset pipeline with and without shuffling the dataset
print("\[INFO\] building the test dataset with and without shuffle...")
test\_ds\_wo\_shuffle = prepare\_batch\_dataset(
	config.TEST\_DATA\_PATH,
	config.IMAGE\_SIZE,
	config.BATCH\_SIZE,
	shuffle=False
)

test\_ds\_shuffle = prepare\_batch\_dataset(
	config.TEST\_DATA\_PATH, config.IMAGE\_SIZE, config.BATCH\_SIZE
)

We start by importing necessary packages such as the `prepare_batch_dataset`, `callbacks`, `config`, `MobileNet`, `classification_report` from `sklearn`, `pyplot` from `matplotlib`, `numpy`, and `tensorflow` on **Lines 5-13**.

Then on **Lines 17-26**, the test datasets are created with and without shuffling, using the `prepare_batch_dataset` function with the path of the test data, image size, and batch size, as specified in the config file.

We created without shuffling to have an association between the test data ground-truth labels and predicted labels for computing the classification report.

\# load the trained image classification convolutional neural network
model = tf.keras.models.load\_model(config.TRAINED\_MODEL\_PATH)
# print model summary on the terminal
print(model.summary())

# evaluate the model on test dataset and print the test accuracy
loss, accuracy = model.evaluate(test\_ds\_wo\_shuffle)
print("Test accuracy :", accuracy)

# fetch class names
class\_names = test\_ds\_wo\_shuffle.class\_names

# generate classification report by (i) predicting on test dataset
# (ii) take softmax of predictions (iii) for each sample in test set
# fetch index with max. probability (iv) create a vector of true labels
# (v) pass ground truth, prediction for test data along with class names
# to classification\_report method
test\_pred = model.predict(test\_ds\_wo\_shuffle)
test\_pred = tf.nn.softmax(test\_pred)
test\_pred = tf.argmax(test\_pred, axis=1)
test\_true\_labels = tf.concat(
	\[label for \_, label in test\_ds\_wo\_shuffle\], axis=0
)
print(
	classification\_report(
		test\_true\_labels, test\_pred, target\_names=class\_names
	)
)

On **Lines 29-31**, the pre-trained model is loaded from the path specified in the config file. Its summary is printed on the terminal.

Then on **Lines 34 and 35**, the model is evaluated on the test dataset (without shuffling). The accuracy is printed on the terminal.

On **Lines 45-55**, a classification report is generated by

*   predicting on the test dataset
*   applying softmax to the predictions
*   fetching the index with the maximum probability for each sample in the test set
*   creating a vector of true labels
*   passing that, along with the class names, to the `classification_report` method from `sklearn`

\# Retrieve a batch of images from the test set and run inference
image\_batch, label\_batch = test\_ds\_shuffle.as\_numpy\_iterator().next()
predictions = model.predict\_on\_batch(image\_batch)

# Apply a softmax
score = tf.nn.softmax(predictions)
print(score.shape)

# save the plot for model prediction along with respective test images
plt.figure(figsize=(10, 10))
for i in range(9):
	ax = plt.subplot(3, 3, i + 1)
	plt.imshow(image\_batch\[i\].astype("uint8"))
	plt.title(class\_names\[np.argmax(score\[i\])\])
	plt.axis("off")
plt.savefig(config.TEST\_PREDICTION\_OUTPUT)

On **Lines 58-62**, we retrieve a batch of images from the test set, run inference on it, and apply softmax.

Finally, from **Lines 66-72**, a plot is created of the model’s predictions along with their respective images from the test set, and it is saved to the path specified in the config file.

* * *

#### [**Classification Report**](#TOC)

**Figure 5** summarizes how the model performs on each of the 15 classes with help from the `precision` and `recall` metrics. **Figure 5** also shows that the model does exceptionally well in all the classes.

Apart from the categories: `Bean`, `Bitter_Gourd`, `Cauliflower`, `Cucumber`, `Papaya`, and `Tomato`, all other classes have 100% `f1-score` (harmonic average of `precision` and `recall`). This is an excellent validation for the model we trained on the vegetable classification dataset.

[![Image 7](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20470%20362'%3E%3C/svg%3E)](https://pyimagesearch.com/wp-content/uploads/2023/02/classification-report-test-set.png)

**Figure 5:** Classification report depicting `precision`, `recall`, `f1-score`, and `support` (number of samples) across each class with final weighted average accuracy.

* * *

#### [**Sample Test Set Predictions**](#TOC)

We have already seen the quantitative results in detail for each class through a classification report; in **Figure 6**, we show few qualitative results, the predictions made by the trained image classification model on the test set.

[![Image 8](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20495%20500'%3E%3C/svg%3E)](https://pyimagesearch.com/wp-content/uploads/2023/02/predictions-test-set.png)

**Figure 6:** Qualitative results showing the prediction on the test set by the trained classification model.

### What's next? We recommend [PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend).

**Course information:**  
84 total classes • 114+ hours of on-demand code walkthrough videos • Last updated: February 2024  
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

**I strongly believe that if you had the right teacher you could _master_ computer vision and deep learning.**

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s _not_ the case.

All you need to master computer vision and deep learning is for someone to explain things to you in _simple, intuitive_ terms. _And that’s exactly what I do_. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to _successfully_ and _confidently_ apply computer vision to your work, research, and projects. Join me in computer vision mastery.

**Inside PyImageSearch University you'll find:**

*   ✓ **84 courses** on essential computer vision, deep learning, and OpenCV topics
*   ✓ **84 Certificates** of Completion
*   ✓ **114+ hours** of on-demand video
*   ✓ **Brand new courses released _regularly_**, ensuring you can keep up with state-of-the-art techniques
*   ✓ **Pre-configured Jupyter Notebooks in Google Colab**
*   ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
*   ✓ Access to **centralized code repos for _all_ 536+ tutorials** on PyImageSearch
*   ✓ **Easy one-click downloads** for code, datasets, pre-trained models, etc.
*   ✓ **Access** on mobile, laptop, desktop, etc.

[Click here to join PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend)

* * *

[**Summary**](#TOC)
-------------------

In this tutorial, we trained an image classification neural network on a vegetable image classification dataset from Kaggle.

The image classification network consisted of MobileNetV2 as the base model for feature extraction pre-trained on the ImageNet dataset. The classification layers were added to the base model to convert the feature vectors into class predictions.

The model performed exceptionally well on the test dataset provided on Kaggle. Based on the classification report, the model achieved almost 100% accuracy across all classes.

As we mentioned earlier, this tutorial was an essential foundation for our following tutorial in the OAK-101 course, where we will learn to deploy this image classification model. In addition, we will learn to convert the TensorFlow framework model to the format [OAK](https://shop.luxonis.com/?aff=10) expects.

Also, we will learn to use DepthAI API to run inference on OAK with images and camera video streams.

* * *

### [**Citation Information**](#TOC)

**Sharma, A.** “Training a Custom Image Classification Network for OAK-D,” _PyImageSearch_, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, [https://pyimg.co/ab1dj](https://pyimg.co/ab1dj)

@incollection{Sharma\_2023\_CIC-OAK-D,
  author = {Aditya Sharma},
  title = {Training a Custom Image Classification Network for {OAK-D}},
  booktitle = {PyImageSearch},
  editor = {Puneet Chugh and Aritra Roy Gosthipaty and Susan Huot and Kseniia Kidriavsteva and Ritwik Raha and Abhishek Thanki},
  year = {2023},
  url = {https://pyimg.co/ab1dj},
}

* * *

![Image 9: Featured Image](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20500%20276'%3E%3C/svg%3E)

### Unleash the potential of computer vision with Roboflow - Free!

*   Step into the realm of the future by [signing up or logging into your Roboflow account](https://roboflow.com/?ref=pyimagesearch). Unlock a wealth of innovative dataset libraries and revolutionize your computer vision operations.
*   Jumpstart your journey by choosing from our broad array of datasets, or benefit from PyimageSearch’s comprehensive library, crafted to cater to a wide range of requirements.
*   Transfer your data to Roboflow in any of the 40+ compatible formats. Leverage cutting-edge model architectures for training, and deploy seamlessly across diverse platforms, including API, NVIDIA, browser, iOS, and beyond. Integrate our platform effortlessly with your applications or your favorite third-party tools.
*   Equip yourself with the ability to train a potent computer vision model in a mere afternoon. With a few images, you can import data from any source via API, annotate images using our superior cloud-hosted tool, kickstart model training with a single click, and deploy the model via a hosted API endpoint. Tailor your process by opting for a code-centric approach, leveraging our intuitive, cloud-based UI, or combining both to fit your unique needs.
*   Embark on your journey today with absolutely no credit card required. Step into the future with Roboflow.

[Join Roboflow Now](https://roboflow.com/?ref=pyimagesearch)

* * *

**To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), _simply enter your email address in the form below!_**

![Image 10](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20410%20567'%3E%3C/svg%3E)

#### Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a **FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning.** Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!