Table of Contents

CycleGAN: Unpaired Image-to-Image Translation (Part 3)

In this tutorial, we will dive deeper into the training process of our unpaired image-to-image translation model. Specifically, we will train our CycleGAN model using Keras and TensorFlow and also learn how we can use it to perform unpaired image translation on novel unseen images.

This lesson is the last in a 3-part series on GANs 301:

1. CycleGAN: Unpaired Image-to-Image Translation (Part 1)
2. CycleGAN: Unpaired Image-to-Image Translation (Part 2)
3. CycleGAN: Unpaired Image-to-Image Translation (Part 3) (this tutorial)

To learn to train and use the CycleGAN model in real-time, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

CycleGAN: Unpaired Image-to-Image Translation (Part 3)

In the first tutorial of this series on unpaired image-to-image translation, we introduced the CycleGAN model. We also discussed the formulation and principles that allow it to perform image-to-image translation from unpaired data. Furthermore, in the previous tutorial of this series, we discussed the Apples2Oranges Dataset and implemented the CycleGAN architecture from scratch in Keras and TensorFlow.

In this tutorial, we will continue this discussion and discuss in detail the training process of our CycleGAN model. Specifically, we will develop our data pipeline, implement the loss functions discussed in Part 1 and write our own code to train the CycleGAN model end-to-end using Keras and TensorFlow. We will also see how we can use our trained CycleGAN model to perform inference and translate images in real time.

Configuring Your Development Environment

To follow this guide, you need to have the TensorFlow library installed on your system.

Luckily, TensorFlow is pip-installable:

$ pip install tensorflow

Need Help Configuring Your Development Environment?

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 now on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides 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

We first need to review our project directory structure.

Start by accessing this tutorial’s “Downloads” section to retrieve the source code and example images.

From there, take a look at the directory structure:

├── inference.py
├── outputs
│   ├── images
│   └── models
├── pyimagesearch
│   ├── CycleGANTraining.py
│   ├── __init__.py
│   ├── config.py
│   ├── data_preprocess.py
│   ├── model.py
│   └── train_monitor.py
└── train.py

In the previous tutorial of this series, we discussed the function of each file in our project directory. Furthermore, we discussed the config file details, the model architecture implementation (i.e., the model.py file), and the data preprocess procedure (i.e., the data_preprocess.py file).

In this part, we will discuss the training process of our image translation model. Specifically, we will discuss the CycleGANTraining.py and train.py files along with the implementation of the training callback, which will allow us to monitor the training process (i.e., the train_monitor.py file). Furthermore, we will also look into the inference stage of our trained CycleGAN model and discuss in detail the inference.py file.

Implementing CycleGAN Training

We start by implementing the CycleGANTraining class, which implements the training procedure of our CycleGAN model. With the help of this class, we implement the loss functions which we need to train our model, define one training iteration and call the optimizers to update our parameters after backpropagation.

So, let us open the CycleGANTraining.py file and discuss the code line by line to understand the training procedure better.

# import the necessary packages
from tensorflow.keras import Model
import tensorflow as tf

class CycleGANTraining(Model):
    def __init__(self, generatorG, discriminatorX, generatorF, 
        discriminatorY, **kwargs):
        super().__init__(**kwargs)
        # initialize the generators and discriminators
        self.generatorG = generatorG
        self.discriminatorX = discriminatorX
        self.generatorF = generatorF
        self.discriminatorY = discriminatorY

    def compile(self, gOptimizerG, dOptimizerX, gOptimizerF, 
        dOptimizerY, bceLoss):
        super().compile()
        # initialize the optimizers for the generator 
        # and discriminator
        self.gOptimizerG = gOptimizerG
        self.dOptimizerX = dOptimizerX
        
        self.gOptimizerF = gOptimizerF
        self.dOptimizerY = dOptimizerY

        # initialize the loss functions
        self.bceLoss = bceLoss

    def train_step(self, images):
        # grab the input images and target images
        (inputImage, targetImage) = images
    
        # initialize gradient tapes for both generator and discriminator
        with  tf.GradientTape() as genG_tape, tf.GradientTape() as discY_tape, tf.GradientTape() as genF_tape, tf.GradientTape() as discX_tape:
            # generate fake target images and cycle input images
            genImagesY = self.generatorG(inputImage, training=True)
            cycledImageX = self.generatorF(genImagesY, training=True)

            # generate fake input images and cycle-real target images
            genImagesX = self.generatorF(targetImage, training=True)
            cycledImageY = self.generatorG(genImagesX, training=True)

            # identity mapping
            samegenX = self.generatorF(inputImage, training=True)
            samegenY = self.generatorG(targetImage, training=True)

            # discriminator output for real target images
            discRealOutputY = self.discriminatorY([targetImage],
                training=True
            )

            # discriminator output for fake target images
            discFakeOutputY = self.discriminatorY([genImagesY], 
                training=True
            )

            # discriminator output for real input images
            discRealOutputX = self.discriminatorX([inputImage],
                training=True
            )

            # discriminator output for fake input images
            discFakeOutputX = self.discriminatorX([genImagesX], 
                training=True
            )

            # calculate cycle loss
            lossA = 10 * (tf.reduce_mean(tf.abs(targetImage - cycledImageY)))   
            lossB = 10 * (tf.reduce_mean(tf.abs(inputImage - cycledImageX)))
            totalCycleLoss = lossA + lossB 

            # calculate identity mapping
            idenityLossG = 10 * 0.5 * (tf.reduce_mean(tf.abs(targetImage - samegenY)))
            identityLossF = 10 * 0.5 * (tf.reduce_mean(tf.abs(inputImage - samegenX)))

            # calculate generator loss 
            ganLossG = self.bceLoss(tf.ones_like(discFakeOutputY), discFakeOutputY)
            ganLossF = self.bceLoss(tf.ones_like(discFakeOutputX), discFakeOutputX)

            # calculate all discriminator losses
            realDiscLossY = self.bceLoss(tf.ones_like(discRealOutputY),
                discRealOutputY)
            fakeDiscLossY = self.bceLoss(tf.zeros_like(discFakeOutputY),
                discFakeOutputY)    
            realDiscLossX = self.bceLoss(tf.ones_like(discRealOutputX),
                discRealOutputX)
            fakeDiscLossX = self.bceLoss(tf.zeros_like(discFakeOutputX),
                discFakeOutputX)

            # calculate total discriminator loss
            totalDiscLossY = 0.5*( realDiscLossY + fakeDiscLossY)
            totalDiscLossX = 0.5*( realDiscLossX + fakeDiscLossX)

            # calculate total generator loss
            totalGenLossG = ganLossG + totalCycleLoss + idenityLossG
            totalGenLossF = ganLossF + totalCycleLoss + identityLossF

        # calculate the generator and discriminator gradients
        generatorGradientsG = genG_tape.gradient(totalGenLossG, 
            self.generatorG.trainable_variables)
        generatorGradientsF = genF_tape.gradient(totalGenLossF, 
            self.generatorF.trainable_variables)
        discriminatorXGradients = discX_tape.gradient(totalDiscLossX, 
            self.discriminatorX.trainable_variables)
        discriminatorYGradients = discY_tape.gradient(totalDiscLossY, 
            self.discriminatorY.trainable_variables)

        # apply the gradients to both generators and discriminators
        self.gOptimizerG.apply_gradients(zip(generatorGradientsG,
            self.generatorG.trainable_variables))
        self.gOptimizerF.apply_gradients(zip(generatorGradientsF,
            self.generatorF.trainable_variables))
        self.dOptimizerX.apply_gradients(zip(discriminatorXGradients,
            self.discriminatorX.trainable_variables))
        self.dOptimizerY.apply_gradients(zip(discriminatorYGradients,
            self.discriminatorY.trainable_variables))

        # return the generator and discriminator losses
        return {"dLossX_input": totalDiscLossX, "gLossG": ganLossG+totalCycleLoss,
            "dLossY_output": totalDiscLossY, "gLossF": ganLossF+totalCycleLoss}

Let us review the pipeline of our CycleGAN model, which we discussed in detail in Part 1 of this series. We have 2 Generators (G and F) and 2 Discriminators (X and Y), which allow us to perform image-to-image translation between the apples and oranges domains without paired images.

Generator G takes images from one domain (say, apples) and translates them to images of the other domain (say, oranges), as discussed in Part 1 of this series. Furthermore, Generator F performs the reverse mapping and takes as input images from the oranges domain and outputs images in the apples domain.

On the other hand, Discriminator X takes the outputs from Generator G and ensures that they match the real images in the oranges domain with the help of adversarial loss. Similarly, Discriminator Y takes the outputs from Generator F and ensures that they match the real images in the apples domain.

Now that we have discussed an overview of the pipeline, let us start implementing it.

We start by importing the Model module from tensorflow.keras and the tensorflow library, as shown on Lines 2 and 3.

Now, we start defining our CycleGANTraining class which implements the training procedure for our image-to-image translation model.

We first define our __init__ constructor, which takes as input the components of our CycleGAN model, that is, the two generators (i.e., generatorG and generatorF) and discriminators (i.e., discriminatorX and discriminatorY), as shown on Lines 6 and 7.

In the __init__ function, we initialize the generator attributes (i.e., self.generatorG and self.generatorF) and discriminator attributes (i.e., self.discriminatorX and self.discriminatorY) of the class with the generator and discriminator arguments (Lines 10-13).

Now that we have defined our constructor, we implement the compile function (Lines 15-27), which takes as input the corresponding generator optimizers (i.e., gOptimizerG and gOptimizerF), the discriminator optimizers (i.e., dOptimizerX and dOptimizerY) and the loss function (i.e., bceLoss).

The compile function simply initializes the optimizers attributes for the generators (self.gOptimizerG and self.gOptimizerF) and the optimizers attributes for discriminators (self.dOptimizerX and self.dOptimizerY) with the optimizers in the arguments of the function (Lines 20-24). Finally, we initialize the loss function attribute (i.e., self.bceLoss) with the bceLoss argument.

Now that we have defined our helper functions, it is time to implement the training function (i.e., train_step()), which takes as arguments the input images.

On Line 31, we grab the input and target images (i.e., inputImage and targetImage) from the images argument.

Next, we initialize the gradient tapes for both generator and discriminator (Line 34) since, during training, we want TensorFlow to track gradients so we can backpropagate through them later.

We first pass the inputImage in Domain X through our Generator G to get images in Domain Y (i.e., genImagesY, Line 36). Then we pass this output through our Generator F (Line 37) to get back our image in Domain X (i.e., cycledImageX). Notice that this forms our forward cyclic consistency cycle, as explained in Part 1 of this series.

Similarly, we pass the targetImage in Domain Y through Generator F to get images in Domain X (i.e., genImagesX, Line 40). Then we pass this output through our Generator G (Line 41) to get back our image in Domain X (i.e., cycledImageY). Notice that this forms our backward cyclic consistency cycle, as explained in Part 1 of this series.

Furthermore, we also get samegenX and samegenY by passing the inputImage and targetImage through generatorF and generatorG, respectively, as shown on Lines 44 and 45).

Next, we pass the real targetImage, which belongs to Domain Y through the discriminatorY to get discRealOutputY (Line 48). Furthermore, we pass the generated or fake images in Domain Y genImagesY through the discriminatorY to get discFakeOutputY (Line 53).

Similarly, we pass the real inputImage, which belongs to Domain X, through discriminatorX to get discRealOutputX (Line 58). Furthermore, we pass the generated or fake images in Domain X genImagesX through discriminatorX to get discFakeOutputX (Line 63).

Now that we have the outputs of the generators and discriminators, it is time to compute the adversarial and cyclic consistency losses.

As discussed above, we start with the backward cycle consistency loss where we want cycledImageY to be close to the targetImage and impose the mean absolute error loss, as shown on Line 68. Similarly, we apply the forward cycle consistency loss where we want cycledImageX to be close to the inputImage and impose the mean absolute error loss, as shown on Line 69.

Note that the coefficients (which is 10.0 here) for both losses are hyperparameters, allowing us to weigh the different losses. Finally, our total cyclic consistency loss is the sum of the forward and backward cyclic consistency loss, as shown on Line 70.

Next, we compute and impose the identity loss. Note that this loss simply tries to regularize the generators to be near an identity mapping if an image of Domain Y is passed as input to Generator G or an image of Domain X is passed through Generator F.

The idea behind this loss is that if Generator G (whose task is to translate images from Domain X to Domain Y) gets as input an image that is already in Domain Y (i.e., targetImage in our case) it should not change it, and its output should be the same as the input.

Thus, to ensure this, we impose a mean absolute error-based loss such that targetImage is close to samegenY (Line 73). Similarly, we also regularize Generator F with this loss by ensuring that inputImage is close to samegenX (Line 74).

Now that we have defined our cyclic consistency and identity mapping losses, it is time to compute the adversarial loss for training the generators and discriminators. Let us start with the generators.

For Generator G, we want it to make genImagesY close to real images from Domain Y such that when genImagesY is passed through the Discriminator Y the output that is discFakeOutputY is a high probability score (i.e., close to 1). Thus, we apply our self.bceLoss between discFakeOutputY and a set of ones (i.e., tf.ones_like(discFakeOutputY)), as shown on Line 77.

Similarly, for Generator F, we want it to make genImagesX close to real images from Domain X such that when genImagesX is passed through the Discriminator X, the output that is discFakeOutputX is a high probability score (i.e., close to 1). Thus, we apply our self.bceLoss between discFakeOutputX and a set of ones (i.e., tf.ones_like(discFakeOutputX)), as shown on Line 78.

Next, we want to train the discriminators to give high probability scores (i.e., close to 1) to real images in Domains X and Y and low probability scores to fake or generated images (i.e., close to 0).

This implies that for Discriminator Y, we want it to give a high probability to discRealOutputY and a low probability to discFakeOutputY. Thus, we apply our self.bceLoss between discRealOutputY and a set of ones (i.e., tf.ones_like(discRealOutputY)) (Line 81) and between discFakeOutputY and a set of zeros (i.e., tf.zeros_like(discFakeOutputY)) (Line 83).

Similarly, for Discriminator X, we want it to give a high probability to discRealOutputX and a low probability to discFakeOutputX. Thus, we apply our self.bceLoss between discRealOutputX and a set of ones (i.e., tf.ones_like(discRealOutputY)) (Line 85) and between discFakeOutputX and a set of zeros (i.e., tf.zeros_like(discFakeOutputX)) (Line 87).

Finally, we define the total loss for Discriminator Y (i.e., totalDiscLossY), which is the sum of realDiscLossY and fakeDiscLossY, as shown on Line 91. Similarly, the total loss for Discriminator X (i.e., totalDiscLossX) is the sum of realDiscLossX and fakeDiscLossX, as shown on Line 92. Note that the coefficient 0.5 in both losses is a hyperparameter which allows us to weigh the losses.

Furthermore, we define the total loss for Generator G (i.e., totalGenLossG), which is the sum of ganLossG, the cycle consistency loss computed above (i.e., totalCycleLoss), and the identity mapping loss for Generator G (i.e., idenityLossG) (Line 95).

Similarly, we define the total loss for Generator F (i.e., totalGenLossF), which is the sum of ganLossF, the cycle consistency loss computed above (i.e., totalCycleLoss) and the identity mapping loss for Generator G (i.e., idenityLossF) (Line 96).

Now that we have computed all our losses, it is time to backpropagate through our model and compute the gradients for both the generators and discriminators.

On Lines 99 and 100, we compute the gradients of the total loss for Generator G (i.e., totalGenLossG) w.r.t. its trainable parameters (i.e., self.generatorG.trainable_variables) with the help of the gradient() functionality. Then, we repeat the same process for Generator F (Lines 101 and 102), Discriminator X (Lines 103 and 104), and Discriminator Y (Lines 105 and 106).

Next, we move toward optimizing our model using respective optimizers. We first zip together the gradients and their corresponding parameters and apply all the computed gradients to the parameters of the respective generators and discriminators using the apply_gradients functionality.

On Lines 109 and 110, we apply the gradients to Generator G with the help of the optimizer self.gOptimizerG. Similarly, we apply the respective gradients to Generator F (Lines 111 and 112), Discriminator X (Lines 113 and 114), and Discriminator Y (Lines 115 and 116).

Finally, we return the computed discriminator and generator losses, as shown on Lines 119 and 120.

Implementing Training Callback

Now that we have defined the class that implements our training procedure, it is time to implement our training callback which will allow us to monitor the CycleGAN training process.

We open the train_monitor.py file and get started.

# import the necessary packages
from tensorflow.keras.preprocessing.image import array_to_img
from tensorflow.keras.callbacks import Callback
from matplotlib.pyplot import subplots
import matplotlib.pyplot as plt
import tensorflow as tf

def get_train_monitor(testInput, testOutput, imagePath, batchSize, epochInterval):
    # grab the input image and target image
    inputImage = next(iter(testInput))
    outputImage = next(iter(testOutput))

    class TrainMonitor(Callback):
        def __init__(self, epochInterval=None):
            self.epochInterval = epochInterval

        def on_epoch_end(self, epoch, logs=None):
            if self.epochInterval and epoch % self.epochInterval == 0:
                # get the CycleGAN prediction
                preds = self.model.generatorG.predict(inputImage)

                # initialize the subplots
                (fig, axes) = subplots(nrows=batchSize, ncols=3,
                    figsize=(50, 50))

                # plot the predicted images 
                for (ax, inp, pred, tgt) in zip(axes, inputImage,
                    preds, outputImage):
                    # plot the input image
                    ax[0].imshow(array_to_img(inp))
                    ax[0].set_title("Input Image")

                    # plot the predicted CycleGAN image
                    ax[1].imshow(array_to_img(pred))
                    ax[1].set_title("CycleGAN Prediction")

                    # plot the ground truth
                    ax[2].imshow(array_to_img(tgt))
                    ax[2].set_title("Output Image")

                plt.savefig(f"{imagePath}/{epoch:03d}.png")
                plt.close()
    
    # instantiate a train monitor callback
    trainMonitor = TrainMonitor(epochInterval=epochInterval)

    # return the train monitor
    return trainMonitor

We start by importing the necessary packages on Lines 2-6, which include the important functionalities like array_to_img and Callback (Lines 2 and 3), the packages from matplotlib for visualization (Lines 4 and 5), and the tensorflow library (Line 6).

Now that we have imported the important modules, we start with the definition of our get_train_monitor() function (Lines 8-48), which implements the TrainMonitor() class.

The get_train_monitor() function takes as arguments the input and target images (i.e., testInput and testOutput, respectively), the imagePath, the batchSize, and the epochInterval parameter (Line 8).

Furthermore, we create an iterator for testInput and testOutput using the iter() method and grab the input and output images (i.e., inputImage and outputImage) using the next() method, as shown on Lines 10 and 11.

Next, we define the TrainMonitor class (Lines 13-42), which inherits from the Callback module, as shown on Line 13. We start by defining the init method, which takes the epochInterval parameter as an argument and initializes the self.epochInterval attribute, as shown on Line 15.

Now we define the on_epoch_end() function, which takes the current epoch and the logs parameter as input arguments, as shown on Line 17.

On Line 18, we check if the current epoch is divisible by the epochInterval, and if this is true, we execute the function. But, first, we pass the inputImage through the CycleGAN generator using the self.model.generatorG.predict() function and store the output as preds, as shown on Line 20.

To visualize the input and output predictions, we first use matplotlib subplots to initialize subplots, as shown on Lines 23 and 24. Note that the subplot function takes as argument the number of rows and columns and the size of the figure to be plotted as shown.

We first instantiate a for loop to plot the results (Line 27). Then, we plot the input image by first converting the input (i.e., inp) to an image using the array_to_img() function (Line 30) and setting the image title to “Input Image” (Line 31). Next, we plot the corresponding CycleGAN output (i.e., pred) (Line 34) and set the image title to “CycleGAN Prediction” (Line 35). Similarly, we plot the ground truth image (i.e., tgt) (Lines 38 and 39).

Finally, we save our visualization using plt.savefig() at the given imagePath (Line 41) and finish our plotting task with plt.close() (Line 42). With this, we finish the definition of our TrainMonitor class.

Now we instantiate a train monitor callback on Line 45 and return it on Line 48.

Implementing Data Pipeline and Model Training

Now that we have defined the callbacks and the class that implements our training procedure, it is time to build our data pipeline and call our CycleGANTraining class to train our end-to-end image translation model.

Let us open the train.py file and get started.

# USAGE
# python train.py 

# import tensorflow and fix the random seed for better reproducibility
import tensorflow as tf
tf.random.set_seed(42)

# import the necessary packages
from pyimagesearch import config
from pyimagesearch.model import CycleGAN
from pyimagesearch.CycleGANTraining import CycleGANTraining
from pyimagesearch.data_preprocess import read_train_example
from pyimagesearch.data_preprocess import  read_test_example
from pyimagesearch.train_monitor import get_train_monitor
from tensorflow.keras.optimizers import Adam
import tensorflow_datasets as tfds
import pathlib
import os

# define the module level autotune
AUTO = tf.data.AUTOTUNE

# downloading the apple to orange dataset using tensorflow datasets
print("[INFO] downloading the apple 2 orange dataset...")
dataset = tfds.load("cycle_gan/apple2orange")
(trainInput, trainOutput)= (dataset["trainA"], dataset["trainB"]) 

# prepare the data using data processing functions 
print("[INFO] pre-processing the training dataset...")
trainInput = trainInput.map(
    read_train_example, num_parallel_calls=AUTO).shuffle(
    config.TRAIN_BATCH_SIZE).batch(config.TRAIN_BATCH_SIZE).repeat()
trainOutput = trainOutput.map(
    read_train_example, num_parallel_calls=AUTO).shuffle(
    config.TRAIN_BATCH_SIZE).batch(config.TRAIN_BATCH_SIZE).repeat()

# load the test data and pre-process it
(testInput, testOutput) = (dataset["testA"], dataset["testB"])
testInput = testInput.map(read_test_example,
    num_parallel_calls=AUTO).shuffle(
    config.INFER_BATCH_SIZE).batch(config.INFER_BATCH_SIZE)
testOutput = testOutput.map(read_test_example,
    num_parallel_calls=AUTO).shuffle(
    config.INFER_BATCH_SIZE).batch(config.INFER_BATCH_SIZE)

# build the training dataset
trainDataset = tf.data.Dataset.zip((trainInput, trainOutput))

# initialize the binary cross entropy loss function
loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)

# instantiate CycleGAN object
print("[INFO] initializing the CycleGAN model...")
model = CycleGAN(config.IMG_HEIGHT, config.IMG_WIDTH)

# initialize the generator and discriminator networks
discriminatorX = model.discriminator()
discriminatorY = model.discriminator()
generatorG = model.generator()
generatorF = model.generator()

# check whether output images directory exists
# if it doesn't, then create it
if not os.path.exists(config.BASE_IMAGES_PATH):
    os.makedirs(config.BASE_IMAGES_PATH)

# build the CycleGAN training model and compile it
print("[INFO] building and compiling the CycleGAN training model...")
cycleGAN = CycleGANTraining(
    generatorG=generatorG,
    discriminatorX=discriminatorX,
    generatorF=generatorF,
    discriminatorY=discriminatorY)
cycleGAN.compile(
    gOptimizerG=Adam(learning_rate=config.LR),
    dOptimizerX=Adam(learning_rate=config.LR),
    gOptimizerF=Adam(learning_rate=config.LR),
    dOptimizerY=Adam(learning_rate=config.LR),
    bceLoss=loss
)

# train the CycleGAN model
print("[INFO] training the cycleGAN model...")
callbacks = [get_train_monitor(testInput, testOutput, epochInterval=10,
    imagePath=config.BASE_IMAGES_PATH,
    batchSize=config.INFER_BATCH_SIZE)]
cycleGAN.fit(trainDataset, epochs=config.EPOCHS, callbacks=callbacks,
    steps_per_epoch=config.STEPS_PER_EPOCH)

# save the CycleGAN generator to disk
print("[INFO] saving cycleGAN generator to {}...".format(
    config.GENERATOR_MODEL))
cycleGAN.generatorG.save(config.GENERATOR_MODEL)

We start by importing the tensorflow library (Line 5) and setting the seed for the training so we can reproduce the training process later (Line 6).

Now we import the necessary modules and packages that we will need to train our CycleGAN model. First, we import the config file and the CycleGAN model, which we discussed above (Lines 9 and 10). Next, we also import the CycleGANTraining module (Line 11), the read_train_example (Line 12) and read_test_example (Line 13) functions which we have defined and discussed above.

In addition, we import the get_train_monitor function, which implements the TrainMonitor callback (Line 14), and the Adam optimizer, which we will use to train our CycleGAN model (Line 15). Furthermore, we import the tensorflow_datasets module, the pathlib package, and the os module, as shown on Lines 16-18.

On Line 21, we define the module-level autotune parameter AUTO using the tf.data.AUTOTUNE functionality.

Now that we have imported the important modules and set up the configurations, let us load our Apples2Oranges Dataset, which we will use for this tutorial. The tensorflow_datasets module provides a simple API that lets us download and load the Apples2Oranges Dataset, as shown on Line 25.

Next, we grab the two parts of the training dataset, which are dataset["trainA"] (the input to our model) and the dataset["trainB"] (the desired output from our model) and store them as trainInput and trainOutput, respectively (Line 26).

Now that we have our training data, we apply the pre-processing functions defined earlier to our trainInput. Therefore, we first pre-process the trainInput data (Lines 30-32) using the following functionalities provided by TensorFlow.

• map() functionality allows us to apply the read_train_example function to the images in the data (Lines 30 and 31).
• shuffle() functionality as shown (which takes as argument buffer size which is config.TRAIN_BATCH_SIZE) to randomly sample elements from a buffer of elements (Lines 31 and 32).
• batch() functionality (which takes as argument config.TRAIN_BATCH_SIZE) allows us to sample batches of data samples with the number of elements per batch defined by the config.TRAIN_BATCH_SIZE argument (Line 32).
• repeat() allows us to repeat the dataset samples/entries multiple times to draw samples from the dataset continuously (Line 32).

Similarly, we preprocess the trainOutput data, as shown on Lines 33-35.

Now that our training data is ready, it is time to load the test data. We will use the same process discussed above for the training data. We first load and divide the 2 parts of the test dataset (i.e., testInput, testOutput) (Line 38) and preprocess them similarly to what we did above for the training data (Lines 39-44). Finally, we combine and consolidate our entire training data (i.e., trainDataset) by zipping together trainInput and trainOutput using the tf.data.Dataset.zip() function (Line 47).

Now that we have created our data pipeline, we are ready to initialize our model and the corresponding loss functions.

On Line 50, we initialize the binary cross-entropy loss using tf.keras.losses.BinaryCrossentropy(from_logits=True). Next, on Line 54, we instantiate our CycleGAN model with the (config.IMG_HEIGHT, config.IMG_WIDTH) as arguments. Finally, on Lines 57-60, we initialize our 2 discriminators (i.e., discriminatorX and discriminatorY) and our 2 generators (i.e., generatorG and generatorF).

We then check whether the output image directory exists, and if not, we create it (Lines 64 and 65).

Once we have built the data pipeline and initialized our model, it is time to create the training pipeline for our model. For this, we use the CycleGANTraining module, which takes as input the components of our model, that is, the two generators and discriminators (i.e., generatorG, generatorF and discriminatorX, discriminatorY) (Lines 69-73).

Finally, we compile our model by defining the optimizer that will be used (Adam optimizer in our case) to optimize the two generators and discriminators and the loss function (i.e., loss), as shown on Lines 74-80.

We then define our TrainMonitor callback (which, as discussed above, allows us to visualize the results and monitor training and certain intervals of epochs) using the get_train_monitor function. The function takes as input the testInput and testOutput, the epochInterval=10 at which the callback should be called, the imagePath where the visualizations will be stored, and also the batchSize (Lines 84-86).

Finally, we call the .fit() functionality of Keras with the trainDataset as input along with the number of epochs epochs=config.EPOCHS, the callback that we defined, and finally the steps_per_epoch (Lines 87 and 88).

In the end, we save our trained generator model using the save() functionality which takes as input the path (i.e., config.GENERATOR_MODEL) and saves the weights of our generator, which we will need later for inference to translate images in real-time (Line 93).

Perform Image-to-Image Translation

It is now time to implement the inference stage of our unpaired image-to-image translation pipeline and see our trained CycleGAN model in action.

We open the inference.py file and get started.

# USAGE
# python inference.py 

# import tensorflow and fix the random seed for better reproducibility
import tensorflow as tf
tf.random.set_seed(42)

# import the necessary packages
from pyimagesearch import config
from pyimagesearch.data_preprocess import  read_test_example
from tensorflow.keras.preprocessing.image import array_to_img
from tensorflow.keras.models import load_model
from matplotlib.pyplot import subplots
import tensorflow_datasets as tfds
import pathlib
import os

# define the module level autotune
AUTO = tf.data.AUTOTUNE

# load the test data
print("[INFO] loading the test data...")
dataset = tfds.load("cycle_gan/apple2orange")
(testInput, testOutput) = (dataset["testA"], dataset["testB"])

# pre-process the test data
print("[INFO] pre-processing the test data...")
testInput = testInput.map(read_test_example,
    num_parallel_calls=AUTO).shuffle(
    config.INFER_BATCH_SIZE, seed=18).batch(config.INFER_BATCH_SIZE)
testOutput = testOutput.map(read_test_example,
    num_parallel_calls=AUTO).shuffle(
    config.INFER_BATCH_SIZE, seed=18).batch(config.INFER_BATCH_SIZE)
    
# get the first batch of testing images
sampleInput = next(iter(testInput))
sampleOutput = next(iter(testOutput))

# load the CycleGan model
print("[INFO] loading the CycleGAN model...")
model = load_model(config.GENERATOR_MODEL, compile=False)
    
# predict using CycleGan 
print("[INFO] making predictions with the CycleGAN model...")
preds = model.predict(sampleInput)

# plot the respective predictions
print("[INFO] plotting the CycleGan predictions...")
(fig, axes) = subplots(nrows=config.INFER_BATCH_SIZE, ncols=3,
    figsize=(50, 50))

# plot the predicted images 
for (ax, inp, pred, tar) in zip(axes, sampleInput,
    preds, sampleOutput):
    # plot the input mask image
    ax[0].imshow(array_to_img(inp))
    ax[0].set_title("Input Image")

    # plot the predicted CycleGan image
    ax[1].imshow(array_to_img(pred))
    ax[1].set_title("CycleGan prediction")

    # plot the ground truth
    ax[2].imshow(array_to_img(tar))
    ax[2].set_title("Target label")

# check whether output image directory exists
# if it doesn't then create it
if not os.path.exists(config.BASE_IMAGES_PATH):
    os.makedirs(config.BASE_IMAGES_PATH)

# serialize the results to disk
print("[INFO] saving the CycleGan predictions to disk...")
fig.savefig(config.GRID_IMAGE_PATH)

We start by importing the tensorflow library (Line 5) and setting the seed so we can reproduce the training process later (Line 6).

Next, we import the config file (Line 9) and the important functions for inference like read_test_example, array_to_img, and load_model (Lines 10-12). We also import the subplots module from matplotlib for visualizing our results (Line 13) and the tensorflow_datasets, pathlib package, and the os module (Lines 14-16).

On Line 19, we define the module-level autotune parameter AUTO using the tf.data.AUTOTUNE functionality.

Now that we have imported the important modules and set up the configurations, let us load our apple2orange test dataset using the tensorflow_datasets API, which allows us to directly download and load the apple2orange dataset, as shown on Line 23.

We divide our dataset into 2 parts (i.e., testInput, testOutput) similar to what we had seen earlier in the train.py file.

Similar to how we processed our data during the training phase, we will use the map(), shuffle(), and batch() functionalities to pre-process our testInput and testOutput data and create batches of data samples (Lines 28-33).

Now that we have our test data, we can get the test images and perform inference.

We use the iter() method to create iterators for the testInput and testOutput data and use the next() function to get a batch of samples from each of them (i.e., sampleInput and sampleOutput) (Lines 36 and 37).

Next, we load our trained CycleGAN generator model that we saved above at path config.GENERATOR_MODEL using the load_model functionality from Keras (Line 41).

We can now forward pass our test inputs through our trained CycleGAN generator using the model.predict() function and save the outputs in preds, as shown on Line 45.

To visualize the predictions, we first use the matplotlib subplots to initialize subplots, as shown on Lines 49 and 50. Note that the subplot function takes as an argument the number of rows and columns and the size of the figure to be plotted, as shown.

We then instantiate a for loop to plot the results (Line 53). Next, we plot the input image by first converting the input (i.e., inp) to an image using the array_to_img() function (Line 56) and setting the image title to “Input Image” (Line 57). Then, we plot the corresponding CycleGAN output (i.e., pred) (Line 60) and set the image title to “CycleGAN prediction” (Line 61). Similarly, we plot the ground truth image (i.e., tar) (Lines 64 and 65).

Now that we have visualized our results, we check whether the output image directory where we will save our results exists, and if it does not, we create it (Lines 69 and 70).

Finally, we save our visualization using fig.savefig at the given path, which is config.GRID_IMAGE_PATH (Line 74).

What's next? I recommend PyImageSearch University.

Course information:
76 total classes • 90 hours of on-demand code walkthrough videos • Last updated: May 2023
★★★★★ 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:

• ✓ 76 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 76 Certificates of Completion
• ✓ 90 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 500+ 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

Summary

In this tutorial, we continued our discussion on building an unpaired image translation model and looked into the training process of our CycleGAN pipeline.

Specifically, we developed our data pipeline and implemented the CycleGAN training pipeline from scratch in Keras and TensorFlow. Furthermore, we looked into the inference stage of our CycleGAN model and discussed how we can use the trained model for translating images from one domain to another in real-time.

Citation Information

Chandhok, S. “CycleGAN: Unpaired Image-to-Image Translation (Part 3),” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/b1qon

@incollection{Chandhok_2023_CycleGAN-Part3,
  author = {Shivam Chandhok},
  title = {{CycleGAN}: Unpaired Image-to-Image Translation (Part 3)},
  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/b1qon},
}

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. 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

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!

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!

Download the code!

Website

The post CycleGAN: Unpaired Image-to-Image Translation (Part 3) appeared first on PyImageSearch.

Table of Contents

CycleGAN: Unpaired Image-to-Image Translation (Part 2)

In this tutorial, we will implement our CycleGAN model for unpaired image-to-image translation tasks using TensorFlow and Keras. We will dive into the details of the CycleGAN model architecture and discuss the Apples2Oranges Dataset, which we will use for our unpaired image translation task.

This lesson is the 2nd in a 3-part series on GANs 301:

1. CycleGAN: Unpaired Image-to-Image Translation (Part 1)
2. CycleGAN: Unpaired Image-to-Image Translation (Part 2) (this tutorial)
3. CycleGAN: Unpaired Image-to-Image Translation (Part 3)

To learn how to implement CycleGAN for Unpaired Image-to-Image Translation, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

CycleGAN: Unpaired Image-to-Image Translation (Part 2)

In the previous tutorial of this series, we discussed the task of unpaired image-to-image translation and got a high-level intuition of the CycleGAN model. Furthermore, we delved deeper into the mechanism and loss functions used by CycleGAN to seamlessly perform image-to-image translations from a dataset of unpaired images.

In this tutorial, we will continue our discussion and implement the architecture of our CycleGAN model from scratch using Keras and TensorFlow. Furthermore, we look closer at the Apples2Oranges Dataset and discuss dataset preprocessing techniques, allowing us to process our input data and build our end-to-end image translation model.

Apples2Oranges Dataset

As discussed briefly in the previous tutorial of this series, we will be using the Apples2Oranges Dataset in this tutorial to perform image translation. The Apples2Oranges Dataset was officially introduced in Unpaired Image-to-Image Translation Using Cycle-Consistent Adversarial Networks. It was used to show the unpaired image-to-image translation performance and capabilities of CycleGAN. It can be easily found and downloaded from Roboflow.

We can easily download and get quick access to this dataset from Roboflow universe Apples2Oranges Dataset section. Roboflow provides easy and quick access to many curated computer vision datasets for diverse tasks like Single or Multi-Label Classification, Object Detection, Instance Segmentation etc. It also provides an amazing API that seamlessly allows you to upload your own datasets and apply data augmentation techniques and transformations on your images in real-time. 

If you are interested in experiencing the amazing Roboflow universe and the features it provides quickly head over to the Roboflow website and get your free account now. Go ahead and get started with the tutorials which will help you get started and fully enjoy the features of the roboflow universe.

This dataset consists of 1261 Apples’ Photos & 1267 Oranges’ Photos which form the two domains between which image translation is performed. Notice that the dataset consists of approximately equal numbers of apple and orange data samples.

Furthermore, the dataset is split into a train set which consists of 80% of the data, and a test set which consists of 20% of the data. Figure 1 shows some example images for apples (bottom) and oranges (top) from this dataset.

Configuring Your Development Environment

To follow this guide, you need to have the TensorFlow library installed on your system.

Luckily, TensorFlow is pip-installable:

$ pip install tensorflow

Need Help Configuring Your Development Environment?

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 now on your Windows, macOS, or Linux system?

Then join 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

We first need to review our project directory structure.

Start by accessing this tutorial’s “Downloads” section to retrieve the source code and example images.

From there, let us take a look at the directory structure:

├── inference.py
├── outputs
│   ├── images
│   └── models/generator
├── pyimagesearch
│   ├── CycleGANTraining.py
│   ├── __init__.py
│   ├── config.py
│   ├── data_preprocess.py
│   ├── model.py
│   └── train_monitor.py
└── train.py

The inference.py file implements the code we will use during the inference stage to translate images in real-time and see our model in action. Finally, the outputs folder is where we store the output images and save our trained CycleGAN model.

The pyimagesearch folder contains the main components of our CycleGAN pipeline. In addition, this folder includes the CycleGANTraining.py file, which implements the training procedure for our model.

Furthermore, the config.py file contains the parameter configurations we will use while implementing our image translation pipeline, and the data_preprocess.py file contains the dataset preprocessing code.

The model.py file implements the architecture of our CycleGAN model, and the train_monitor.py file implements a callback which will allow us to visualize and monitor the training process.

Finally, the train.py file implements the code to train our end-to-end CycleGAN model.

In this part, we will discuss the config file, the implementation of the model architecture (i.e., model.py file), and the data preprocess procedure (i.e., data_preprocess.py file).

In the next part of this blog series, we will dive deeper into the training process of our image translation model. Specifically, we will discuss the CycleGANTraining.py and train.py files along with the callback implementation, which will help us monitor the training process (i.e., the train_monitor.py file). Furthermore, we will also look into the inference stage of our trained CycleGAN model and discuss the inference.py file in detail.

Creating Our Configuration File

We start by opening the config.py file, which contains the parameters and initial configurations we will use to implement our CycleGAN model.

# import the necessary packages
import os

# define the batch size for training and inference
TRAIN_BATCH_SIZE = 1
INFER_BATCH_SIZE = 8

# dataset specs
IMG_WIDTH = 256
IMG_HEIGHT = 256
IMG_CHANNELS = 3

# training specs
LR = 2e-4
EPOCHS = 50
STEPS_PER_EPOCH = 800

# path to the base output directory
BASE_OUTPUT_PATH = "outputs"

# path to the cycle gan generator
GENERATOR_MODEL = os.path.join(BASE_OUTPUT_PATH, "models",
    "generator")

# path to the inferred images and to the grid image
BASE_IMAGES_PATH = os.path.join(BASE_OUTPUT_PATH, "images")
GRID_IMAGE_PATH = os.path.join(BASE_IMAGES_PATH, "grid.png")

On Line 2, we import the os module for file system functionalities. Next, we define our batch size for training (i.e., TRAIN_BATCH_SIZE) and inference stage (i.e., INFER_BATCH_SIZE) on Lines 5 and 6, respectively.

Next, we define our data parameters, such as the dimensions (i.e., IMG_WIDTH and IMG_HEIGHT) of the image and the number of channels (i.e., IMG_CHANNELS) (Lines 9-11).

Furthermore, we define the specifications for our training process, such as the learning rate (i.e., LR), the total number of epochs (i.e., EPOCHS), and the number of iterations or steps per epoch (i.e., STEPS_PER_EPOCH), as shown on Lines 14-16.

On Line 19, we define the parent output directory (i.e., BASE_OUTPUT_PATH), and on Line 22, we define the path where the CycleGAN generator will be saved after training (i.e., GENERATOR_MODEL).

Finally, we define the paths where our visualizations from the inference stage will be stored (i.e., BASE_IMAGES_PATH and GRID_IMAGE_PATH) on Lines 26 and 27.

Preprocessing Our Dataset

Now that we have discussed the config file and defined the initial parameter configurations, we are ready to discuss the code that will allow us to preprocess our input data during the training and testing or inference stage.

We open the data_preprocess.py file, which implements this code, and start.

# import the necessary packages
import tensorflow as tf

def preprocess_image(image):
    # convert both images to float32 tensors and
    # convert pixels to the range of -1 and 1
    image = tf.cast(image, tf.float32) / 127.5 - 1

    # return the image
    return (image)
    
def random_jitter(image):
    # upscale the image and randomly crop them
    image = tf.image.resize(image, [286, 286],
        method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    cropped = tf.image.random_crop(image, size=[256, 256, 3])

    # randomly flip the cropped image
    image = tf.image.random_flip_left_right(cropped)

    # return the image
    return image

def read_train_example(data):
    # pre-process the image
    image = preprocess_image(data["image"])
    image = random_jitter(image)

    # reshape the input image
    image = tf.image.resize(image, [256, 256])

    # return the input image
    return (image)

def read_test_example(data):
    # pre-process the image and resize it
    image = preprocess_image(data["image"])
    image = tf.image.resize(image, [256, 256])

    # return the image
    return (image)

We start by importing the tensorflow library on Line 2.

Next, on Lines 4-10, we write the preprocess_image() function, which will allow us to preprocess our images. The function takes the image as input to pre-process, as shown on Line 4.

Then, on Line 7, we convert the image to tf.float32 format using the tf.cast() function and convert its pixels to the range [-1,1]. Since the pixel values are in the range [0,255], we can achieve this by dividing the pixel values by 127.5 and subtracting the value 1. Finally, we return the pre-processed image on Line 10.

Then, on Lines 12-22, we define the random_jitter() function, which will apply data augmentations to our input images. The function takes the image as input to apply data augmentations, as shown on Line 12. Next, on Line 14, we use the tf.image.resize() function to upscale the image to the [286,286] size using the NEAREST_NEIGHBOR interpolation technique, as shown.

Then, on Line 16, we randomly crop the image to the desired size [256, 256, 3] using the tf.image.random_crop function. Finally, we use the tf.image.random_flip_left_right() function on our upscaled and cropped image to apply random flip augmentation and return our final output image on Line 22.

On Lines 24-33, we define the read_train_example() function, which takes as input the data and allows us to apply pre-processing and data augmentation operations to our data during training. Next, we use the preprocess_image() function and the random_jitter() function we defined above on Lines 26 and 27, respectively. Then, on Line 30, we resize our image to the desired [256, 256] dimension and finally return our image on Line 33.

Similar to the read_train_example() function, we now define the read_test_example function, which allows us to pre-process our data during test time. However, since we need to apply data augmentations only during training and not during test time, we only pre-process our image using the preprocess_image function on Line 37 and resize our image to the desired [256, 256] dimension and finally return our image on Line 41.

Implementing the CycleGAN Architecture

We are now ready to dive into the details of our CycleGAN architecture and implement it from scratch using Keras and TensorFlow.

As discussed in the first part of this series, the CycleGAN model consists of two generators and two discriminators. We will create a CycleGAN class that implements our generator architecture and discriminator architecture.

Note that the generator of our CycleGAN follows a structure similar to U-Net with a succession of downsampling layers, the U-shaped bend followed by a succession of upsampling layers with skip connections.

If you are unfamiliar with the U-Net architecture or wish to brush up on the concepts, we have an amazing tutorial that offers an in-depth explanation of the U-Net architecture (U-Net: Training Image Segmentation Models in PyTorch).

We open the model.py file containing the code to implement our CycleGAN model definition and get started.

# import the necessary packages
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.layers import Conv2DTranspose
from tensorflow.keras.layers import LeakyReLU
from tensorflow.keras.layers import concatenate
from tensorflow.keras.layers import MaxPool2D
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import Dropout
from tensorflow.keras import Model
from tensorflow.keras import Input

class CycleGAN():
    def __init__(self, imageHeight, imageWidth):
        # initialize the image height and width
        self.imageHeight = imageHeight
        self.imageWidth = imageWidth

    def generator(self):
        # initialize the input layer
        inputs = Input([self.imageHeight, self.imageWidth, 3])
  
        # down Layer 1 (d1) => final layer 1 (f1)
        d1 = Conv2D(32, (3, 3), activation="relu", padding="same")(
            inputs)
        d1 = Dropout(0.1)(d1)
        f1 = MaxPool2D((2, 2))(d1)

        # down Layer 2 (l2) => final layer 2 (f2)
        d2 = Conv2D(64, (3, 3), activation="relu", padding="same")(f1)
        f2 = MaxPool2D((2, 2))(d2)

        #  down Layer 3 (l3) => final layer 3 (f3)
        d3 = Conv2D(96, (3, 3), activation="relu", padding="same")(f2)
        f3 = MaxPool2D((2, 2))(d3)

        # down Layer 4 (l3) => final layer 4 (f4)
        d4 = Conv2D(96, (3, 3), activation="relu", padding="same")(f3)
        f4 = MaxPool2D((2, 2))(d4)

        # u-bend of the u-bet
        b5 = Conv2D(96, (3, 3), activation="relu", padding="same")(f4)
        b5 = Dropout(0.3)(b5)
        b5 = Conv2D(256, (3, 3), activation="relu", padding="same")(b5)

        # upsample Layer 6 (u6)
        u6 = Conv2DTranspose(128, (2, 2), strides=(2, 2),
            padding="same")(b5)
        u6 = concatenate([u6, d4])
        u6 = Conv2D(128, (3, 3), activation="relu", padding="same")(
            u6)

        # upsample Layer 7 (u7)
        u7 = Conv2DTranspose(96, (2, 2), strides=(2, 2),
            padding="same")(u6)
        u7 = concatenate([u7, d3])
        u7 = Conv2D(128, (3, 3), activation="relu", padding="same")(
            u7)

        # upsample Layer 8 (u8)
        u8 = Conv2DTranspose(64, (2, 2), strides=(2, 2),
            padding="same")(u7)
        u8 = concatenate([u8, d2])
        u8 = Conv2D(128, (3, 3), activation="relu", padding="same")(u8)

        # upsample Layer 9 (u9)
        u9 = Conv2DTranspose(32, (2, 2), strides=(2, 2),
            padding="same")(u8)
        u9 = concatenate([u9, d1])
        u9 = Dropout(0.1)(u9)
        u9 = Conv2D(128, (3, 3), activation="relu", padding="same")(u9)

        # final conv2D layer
        outputLayer = Conv2D(3, (1, 1), activation="tanh")(u9)
    
        # create the generator model
        generator = Model(inputs, outputLayer)

        # return the generator
        return generator

We start by importing the important layers to help us build our CycleGAN model on Lines 2-10. Then, on Lines 12-79, we define our CycleGAN() class.

We start by defining the __init__ constructor first (Lines 13-16), which takes as input the imageHeight and imageWidth as shown on Line 13 and initializes the self.imageHeight and self.imageWidth attributes of the class on Lines 15 and 16.

Next, on Lines 18-79, we implement the definition of our CycleGAN generator. We start by initializing the Input layer with the desired dimensions of our input, which is [self.imageHeight, self.imageWidth, 3] on Line 20. Then we begin downsampling our input with a sequence of Conv2D → Dropout → MaxPool2D layers, as shown on Lines 23-26. We further downsample using a sequence of Conv2D → MaxPool2D operations as shown on Lines 29-38.

Then we build the U-shaped bend of our generator using a Conv2D → Dropout → Conv2D sequence of layers, as shown on Lines 41-43, and finally, we use a succession of Conv2DTranspose → concatenate → Conv2D layers to upsample our feature maps, as shown on Lines 46-70.

Note that the concatenate operation implements the skip connections of the U-Net by concatenating the features from the downsampling part of the U shape to the upsampling part of the U shape.

Finally, we have our 1×1 Conv2D output layer (Line 73). We create our generator model using the Model functionality of Keras with the inputs and outputLayer layers as the input and output of our generator model. We then return our generator model on Line 79.

   def discriminator(self):
        # initialize input layer according to PatchGAN
        targetImage = Input(
            shape=[self.imageHeight, self.imageWidth, 3], 
            name="target_image"
        )

        # add four conv2D convolution layers
        x = Conv2D(64, 4, strides=2, padding="same")(targetImage)  
        x = LeakyReLU()(x)
        x = Conv2D(128, 4, strides=2, padding="same")(x)
        x = LeakyReLU()(x)  
        x = Conv2D(256, 4, strides=2, padding="same")(x)
        x = LeakyReLU()(x)   
        x = Conv2D(512, 4, strides=1, padding="same")(x)  

        # add a batch-normalization layer => LeakyReLU => zeropad
        x = BatchNormalization()(x)
        x = LeakyReLU()(x)

        # final conv layer
        last = Conv2D(1, 3, strides=1)(x) 
  
        # create the discriminator model
        discriminator = Model(inputs=[targetImage],
            outputs=last)

        # return the discriminator
        return discriminator

Now that we have defined our generator, we are ready to define the discriminator of our CycleGAN (Lines 81-109). We initialize the input layer using the Input functionality with shape [self.imageHeight, self.imageWidth, 3] and layer name as target_image.

Then we use a succession of Conv2D → LeakyReLU layers, as shown on Lines 89-95, to build our discriminator. We then add a BatchNormalization layer → LeakyReLU sequence (Lines 98 and 99) and finally use a Conv2D layer as the last layer of our discriminator (Line 102).

Similar to what we did in the generator case, we then create our discriminator model using the Model functionality of Keras with the [targetImage] and last layer as the model input and output (Lines 105 and 106). We then return our discriminator model on Line 109.

This completes the implementation of our CycleGAN model, which consists of the generator and discriminator architectures, as discussed in detail above.

What's next? I recommend PyImageSearch University.

Course information:
76 total classes • 90 hours of on-demand code walkthrough videos • Last updated: May 2023
★★★★★ 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:

• ✓ 76 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 76 Certificates of Completion
• ✓ 90 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 500+ 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

Summary

In this tutorial, we continued our discussion on CycleGAN and unpaired image-to-image translation, which we started in the previous post of this series.

Specifically, we implemented the CycleGAN architecture in Keras and TensorFlow from scratch. We implemented the code to pre-process our input data during the training and testing stages of our CycleGAN pipeline.

In the next tutorial of this series, we will dive deeper into the training and inference details of our CycleGAN and see how we can use it to perform the unpaired image-to-image translation in real-time.

Citation Information

Chandhok, S. “CycleGAN: Unpaired Image-to-Image Translation (Part 2),” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/jnael

@incollection{Chandhok_2023_CycleGAN-Part2,
  author = {Shivam Chandhok},
  title = {{CycleGAN}: Unpaired Image-to-Image Translation (Part 2)},
  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/jnael},
}

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. 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

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!

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!

Download the code!

Website

The post CycleGAN: Unpaired Image-to-Image Translation (Part 2) appeared first on PyImageSearch.

Table of Contents

DETR Breakdown Part 1: Introduction to DEtection TRansformers

In this tutorial, we’ll learn about DETR, an end-to-end trainable deep learning architecture for object detection that utilizes a transformer block.

This lesson is the 1st of a 3-part series on DEtection TRansformers (DETR) Explained: A Comprehensive Guide to DETR in Computer Vision:

1. DETR Breakdown Part 1: Introduction to DEtection TRansformers (today’s tutorial)
2. DETR Breakdown Part 2: Methodologies and Algorithms
3. DETR Breakdown Part 3: Architecture and Details

To learn about Detection Transformers, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

DETR Breakdown Part 1: Introduction to DEtection TRansformers

In this blog post, we explore the revolution in object detection with DETR (the entire architecture is presented in the interactive Figure shown below), a unique approach employing Transformers and set prediction for parallel decoding that reimagines the problem statement, bringing an alternative to traditional methods.

We also uncover how Carion et al. (2020) present a niche solution that transcends the old days of object detection.

As we progress through the tutorial series, we aim to understand the components of the above interactive Figure. But, for now, let us focus on the topics at hand:

• why DETR
• where it came from
• what’s the big deal about it

Why DETR and What’s So Special? 🤔

In a world plagued by hand-crafted features and time-consuming processes, a new architecture emerged from the lab: DETR, the end-to-end object detector that could battle the almighty CNNs in its own field! By combining the strengths of two unlikely parents CNNs and transformers, DETR streamlined the training process, bidding farewell to laborious manual work.

With its unique superpower of attention mechanisms, DETR could identify objects and their relationships with remarkable accuracy.

As news of this extraordinary architecture spread, labs turned to newer possibilities inspired by components from DETR, leading to better image segmentation and object detection models. And so, the legacy of DETR, the end-to-end super object detector, was born.

Today we shall delve into its past, the factors that led to its birth, and how this miracle was even achieved in the first place.

What Is Object Detection? 😯

Alright, so object detection is basically guessing where objects are in an image by drawing boxes around them and figuring out what each object is called.

Previously object detection methods solved the problem of predicting object locations in an image by creating related sub-problems. These sub-problems involve estimating object properties using the image’s numerous predefined regions, points, or centers.

So What’s New in DETR? 🎁

The authors want us to see that object detection is pretty straightforward, but they also point out that people used to tackle it more roundaboutly.

The authors view object detection as a set prediction problem. A set prediction problem is when you try to guess a group of items based on some information. Think of it like trying to figure out which movies your friends might like based on the movies they’ve already watched.

Life Is Not a Bed of Roses 🌹: Past Challenges

Just a heads up: Thinking of object detection as a set prediction problem has some challenges. The biggest one is getting rid of duplicate predictions. Keep this in mind as you go through the rest of the tutorial. It’ll help you better understand why the authors made certain choices and decisions.

The Key Players 🤾: Main Components of DETR

Their paper primarily consists of two main components:

1. A unique set-based global loss, which utilizes bipartite matching (an algorithm that facilitates set prediction) to encourage distinct set predictions.
2. A transformer-based architecture, which includes both an encoder and a decoder.

In a Nutshell 🔩

By treating object detection as a set prediction problem, the need for manually designed components previously required in object detection tasks to incorporate prior knowledge is eliminated. This approach simplifies the process and streamlines the task.

The Problem Statement Reimagined 🧠

It is important to understand the contribution of this paper in the field of Computer Vision. To make sense of this, we first reimagine the problem statement by understanding object detection as it was in the old days. Next, we look at the two new methodologies that the authors proposed to solve End-to-End Object Detection. And finally, we look at the simple solution provided by the authors.

1. Object Detection from the Old Days (problem statement)
2. Set Prediction for Object Detection (methodology 1)
3. Transformers and Parallel Decoding for Object Detection (methodology 2)
4. A Simple Solution (solution)

Object Detection from the Old Days 🧓

Previously object detection methods made predictions relative to some initial guesses.

• Two-stage detectors: predict boxes with respect to proposals (shown in Figure 1)
  

• Single-stage detectors: predict with respect to anchors or a grid of possible object centers (shown in Figure 2)
  

The Reimagined Problem Statement: The quality of the predictions relied a lot on how the initial guesses were made. So, instead of using handcrafted methods, we need a way to directly predict the set of detections.

Set Prediction for Object Detection 🕵️

In image processing applications, the basic set prediction task is multilabel classification. But first, let us quickly recap what multilabel and multiclass mean.

Multiclass: An image belonging to one class out of several possible classes (an example of Multiclass object detection is shown in Figure 3).

Multilabel: An image with multiple labels assigned to it simultaneously (an example of Multilabel object detection is shown in Figure 4).

The baseline approach for multilabel classification is called one-vs-rest. Here we train a binary classification model for each label and then try to predict which labels are present in the image and which are not. Of course, this technique works well only when the labels are mutually exclusive.

In the object detection pipeline, near-identical boxes are quite common. This fails the above-explained one-vs-rest approach. Instead, the authors use direct set prediction with a twist to achieve set prediction in such a scenario.

Note: In direct set prediction, we need a way to consider the relationships between all the predicted elements to avoid duplicates.

This is where the authors use the Hungarian algorithm (covered in upcoming tutorials of this series), which helps them match the ground truth (actual data) with their predictions smartly.

Transformers and Parallel Decoding for Object Detection ⚡

Attention mechanisms aggregate information from the entire input sequence. These are more suitable for long sequence representations.

To gain confidence and a better understanding of Transformers and their underlying principles, we recommend watching Video 1, which provides a comprehensive overview and intuitive explanation of Transformers.

Let’s pause and ponder on this for a second. What could go wrong if we input an image, and the object detection model predicts the objects (bounding boxes) in parallel?

While a model outputs all the objects in parallel, each prediction does not have the context of all the other predictions, unlike in autoregressive methods. This can create two problems:

• The parallel model outputs duplicates.
• The parallel model can output objects in a different order each time.

A Niche Solution 🥳

DETR addresses its challenges by combining bipartite matching loss and transformers. The matching loss function helps to pair each prediction with a unique ground truth object, so duplicates aren’t a concern.

Additionally, the bipartite matching loss doesn’t care about the order of predicted objects, meaning we don’t need to worry about the sequence in which predictions are made.

Quiz Time! 🤓

What's next? I recommend PyImageSearch University.

Course information:
76 total classes • 90 hours of on-demand code walkthrough videos • Last updated: May 2023
★★★★★ 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:

• ✓ 76 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 76 Certificates of Completion
• ✓ 90 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 500+ 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

Summary

Let’s take a moment to recap what we’ve learned in this tutorial so far:

1. We explored the challenges and pain points of object detection before the advent of DETR.
2. We briefly introduced DETR’s innovative features and main components, including Set Prediction Loss and the Transformer-based architecture.
3. We reformulated the problem statement to better align with DETR’s approach.
4. We familiarized ourselves with the two key components of DETR and how they address the problem statement.
   

In the next two tutorials of this series, we’ll delve into the tools, methodologies, and architecture of DETR in greater detail, providing a more comprehensive understanding of this groundbreaking approach.

However, valuable research is about more than just reading papers and admiring the brilliance of their authors. It’s also about asking ourselves what we would do in their position and how we would design an architecture to tackle these challenges if DETR didn’t exist.

Stay tuned for the upcoming parts of this series, and be sure to share your thoughts and learnings from DETR and what you’re looking forward to. Connect with us on Twitter by tagging @pyimagesearch.

Citation Information

A. R. Gosthipaty and R. Raha. “DETR Breakdown Part 1: Introduction to DEtection TRansformers,” PyImageSearch, P. Chugh, S. Huot, K. Kidriavsteva, and A. Thanki, eds., 2023, https://pyimg.co/fhm45

@incollection{ARG-RR_2023_DETR-breakdown-part1,
  author = {Aritra Roy Gosthipaty and Ritwik Raha},
  title = {{DETR} Breakdown Part 1: Introduction to DEtection TRansformers},
  booktitle = {PyImageSearch},
  editor = {Puneet Chugh and Susan Huot and Kseniia Kidriavsteva and Abhishek Thanki},
  year = {2023},
  url = {https://pyimg.co/fhm45},
}

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. 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.
• Join our ever-growing community of over 250,000 developers and machine learning engineers from esteemed organizations like Cardinal Health, Walmart, USG, Rivian, Intel, and Medtronic. Experience firsthand the transformative power of Roboflow in constructing efficient computer vision pipelines. Embark on your journey today with absolutely no credit card required. Step into the future with Roboflow.

Join Roboflow Now

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!

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!

Download the code!

Website

The post DETR Breakdown Part 1: Introduction to DEtection TRansformers appeared first on PyImageSearch.

Table of Contents

Hand Gesture Recognition with YOLOv8 on OAK-D in Near Real-Time

In this tutorial, you will learn to perform hand gesture recognition using YOLOv8 on the OAK-D platform. In our previous OAK-102 series tutorial:

• We trained the YOLOv8 object detection model in the PyTorch framework using the Ultralytics repository.
• We will take it a step further by deploying the model on the OAK-D device.
• To perform hand gesture recognition on the OAK device, we would optimize the PyTorch model weights into the MyriadX blob file format using the Luxonis toolkit.
• We will then utilize the DepthAI API to run a real-time hand gesture recognition application on the OAK-D device.

By the end of this tutorial, you will have a good understanding of the process involved in deploying an object detection model on the OAK-D platform, as well as the skills to recognize hand gestures using OAK-D’s camera with the help of the DepthAI API in Python.

This lesson is the 2nd in our 3-part series on OAK 102:

1. Training the YOLOv8 Object Detector for OAK-D
2. Hand Gesture Recognition with YOLOv8 on OAK-D in Near Real-Time (today’s tutorial)
3. OAK 102 (lesson 3)

To learn how to run a hand gesture recognition application on OAK-D, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Hand Gesture Recognition with YOLOv8 on OAK-D in Near Real-Time

Introduction

Hand gesture recognition is an emerging field in computer vision focusing on identifying and interpreting human hand gestures using computer vision and deep learning. It has many applications like sign language recognition, human-computer interaction, virtual reality, and gaming.

Computer vision techniques like object detection are used to analyze and interpret images or videos of human movements to recognize hand gestures. However, hand gesture recognition often requires temporal information, which means analyzing the movement of the human body over time for the algorithm to accurately predict the form of the hand gesture.

For instance, if you want to control the volume of your car’s infotainment system through hand gesture recognition, the algorithm must distinguish between volume up and volume down hand gestures. The deep learning model would require temporal information instead of just spatial information to understand the volume up and volume down hand gestures.

Figure 1 shows sample images from the hand gesture recognition dataset with ground-truth bounding boxes annotated in red, belonging to classes four, five, two, and three. In today’s tutorial, we will learn to recognize one of these five hand gestures using the OAK-D platform.

If you’re interested in the capabilities of an OAK device and the computer vision applications it can run, check out the Introduction to OpenCV AI Kit (OAK) tutorial on PyImageSearch. Specifically, you should look at the Applications on OAK section of the tutorial.

If you followed our previous tutorial on the OAK-102 series, you would be familiar with training a state-of-the-art YOLOv8 object detector using PyTorch with Ultralytics repository, which can recognize hand gestures like one, two, and three. The YOLOv8 model was pretrained on the MS COCO dataset and fine-tuned on only 587 hand gesture images. Despite the small number of images used for training, the model performed surprisingly well, achieving 0.824 mAP@0.5 IoU with the Nano variant and 0.887 mAP@0.5 IoU with the Small variant.

This tutorial will take it further and show you how to deploy the hand gesture recognition model onto the OAK-D device. We will first use the Luxonis toolkit to convert and optimize the PyTorch model for hand gesture recognition into the MyriadX blob file format to accomplish this. Next, we will run the converted model on the OAK-D device using still images and the camera stream.

How would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free. Head over to Roboflow and get a free account to grab these hand gesture images.

Let’s start with today’s tutorial and delve into deploying and running hand gesture recognition applications on the OAK device without further ado.

Configuring Your Development Environment

To follow this guide, you need to have depthai, opencv, and imutils libraries installed on your system.

Luckily, all these libraries are pip-installable:

$ pip install depthai
$ pip install opencv-python
$ pip install imutils

If you need help configuring your development environment for OpenCV, we highly recommend that you read our pip install OpenCV guide — it will have you up and running in minutes.

Need Help Configuring Your Development Environment?

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 on your Windows, macOS, or Linux system now?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides 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

We first need to review our project directory structure.

Start by accessing this tutorial’s “Downloads” section to retrieve the source code and example images.

From there, take a look at the directory structure:

$ tree .
.
├── gesture_recognition_model
│   ├── yolov8n
│   │   ├── yolov8ntrained-simplified.onnx
│   │   ├── yolov8ntrained.bin
│   │   ├── yolov8ntrained.json
│   │   ├── yolov8ntrained.xml
│   │   └── yolov8ntrained_gesture_recog.blob
│   └── yolov8s
│       ├── yolov8strained-simplified.onnx
│       ├── yolov8strained.bin
│       ├── yolov8strained.json
│       ├── yolov8strained.xml
│       └── yolov8strained_gesture_recog.blob
├── pyimagesearch
│   ├── __init__.py
│   ├── config.py
│   └── utils.py
├── recognize_camera.py
├── recognize_images.py
├── results
│   ├── gesture_camera_v8n.mp4
│   ├── gesture_camera_v8s.mp4
│   ├── gesture_pred_images_v8n
│   │   ├── five.jpg
│   │   ├── five_2.jpg
│   │   ├── four.jpg
│   │   ├── one.jpg
│   │   ├── three.jpg
│   │   └── two.jpg
│   └── gesture_pred_images_v8s
│       ├── five.jpg
│       ├── five_2.jpg
│       ├── four.jpg
│       ├── one.jpg
│       ├── three.jpg
│       └── two.jpg
├── test_data
│   ├── five.jpg
│   ├── five_2.jpg
│   ├── four.jpg
│   ├── one.jpg
│   ├── three.jpg
│   └── two.jpg
└── test_images

9 directories, 35 files

In the pyimagesearch directory, we have the following files:

• config.py: The configuration file for the task
• utils.py: The utilities for running the hand gesture recognition on OAK (e.g., creating images and camera pipelines and a few other helper functions)

In the core directory, we have the following:

• gesture_recognition_model: Houses the hand gesture recognition trained model files converted to OpenVINO format (.blob) as required by OAK hardware for both YOLOv8n and YOLOv8s
• test_data: It contains a few hand gesture images from the test set, which the recognize_images.py script will use
• recognize_images.py: The inference script to leverage OAK’s neural accelerator for recognizing hand gestures in images
• recognize_camera.py: The inference script to run hand gesture recognition with OAK’s color camera

In the results directory, we have:

• gesture_pred_images_v8n: Hosts the prediction results performed on the test_data images with the YOLOv8n model
• gesture_pred_images_v8s: Hosts the prediction results performed on the test_data images with the YOLOv8s model
• gesture_camera_v8n.mp4: The prediction output file when inference is performed with OAK’s 4K color camera leveraging the YOLOv8n model
• gesture_camera_v8s.mp4: The prediction output file when inference is performed with OAK’s 4K color camera leveraging the YOLOv8s model

YOLOv8 Model Export to OAK

In this section, we will discuss the steps required to convert the YOLOv8n and YOLOv8s hand gesture recognition models (in PyTorch framework) to MyriadX blob file format.

For a better understanding of this and the following sections of this tutorial, we highly recommend you check out our previous tutorial on Training the YOLOv8 Object Detector for OAK-D.

By now, you already know that the OAK device is not directly compatible with the popular deep learning frameworks, necessitating the conversion of the model into the MyriadX blob file format, which the device supports.

The underlying reason for this specific model format requirement, and the incompatibility of the widely used deep learning frameworks with the OAK device, is the presence of a visual processing unit within the hardware. This unit is built on Intel’s MyriadX processor, which mandates using the blob file format for models.

To export a YOLOv8 model for deployment on an OAK device, you need to follow several steps to convert the model into the MyriadX blob file format. Here’s a step-by-step guide:

1. Model Conversion: First, you need to convert the YOLOv8 model from its native format (typically PyTorch or TensorFlow) to an intermediate format called ONNX (Open Neural Network Exchange).
2. Optimization: To ensure optimal performance on the OAK device, you need to optimize the ONNX model using OpenVINO’s Model Optimizer. This step will generate an Intermediate Representation (IR) consisting of two files: an XML file containing the model architecture and a BIN file containing the model weights.
3. Blob Conversion: Use OpenVINO’s compile_tool to convert the optimized IR model to a MyriadX blob file format.

Luckily, Luxonis offers a convenient tool that streamlines converting YOLOv5-v8 PyTorch models to the MyriadX blob file format. For example, in this tutorial, we aim to deploy a YOLOv8 hand gesture recognition model developed in the PyTorch framework on an OAK device for recognizing hand gestures. Hence, this tool is valuable in assisting us with the YOLOv8 PyTorch model conversion to the MyriadX blob file format.

The following steps will guide you through generating MyriadX blob files for YOLOv8n and YOLOv8s models:

1. Assuming you’ve completed the previous tutorial, you should have fine-tuned YOLOv8n and YOLOv8s PyTorch model weights (using the hand gesture recognition dataset) saved on your local machine.
2. Navigate to the Automatic Yolo exports for OAKs page.
3. As demonstrated in the GIF, set the Yolo Version to “YoloV8 (detection only).”
4. In the File section, upload the yolov8ntrained.pt weights.
5. Set the Input shape to 416.
6. Repeat the above steps for the YOLOv8s model, but this time upload the yolov8strained.pt weights in the File parameter.

Please click the image below to follow the steps to generate the MyriadX blob files for both the YOLOv8n and YOLOv8s models.

Deploy on OAK device: Having generated the blob file, we can proceed to the following section. We will deploy the YOLOv8 model on an OAK device (either OAK-1 or OAK-D) utilizing the DepthAI library.

Configuring the Prerequisites

Before we start our implementation, let’s review our project’s configuration. For that, we will move on to the config.py script located in the pyimagesearch directory.

The config.py script sets up the necessary paths for the YOLOv8n and YOLOv8s models, their configurations, test data, and output directories for the resulting images and videos. It also defines the camera preview dimensions and label names for the hand gestures.

# import the necessary packages
import os
import glob

# define path to the model, test data directory and results
YOLOV8N_MODEL = os.path.join(
    "gesture_recognition_model","yolov8n","yolov8ntrained_gesture_recog.blob"
)
YOLOV8N_CONFIG = os.path.join(
    "gesture_recognition_model","yolov8n","yolov8ntrained.json"
)


YOLOV8S_MODEL = os.path.join(
    "gesture_recognition_model","yolov8s","yolov8strained_gesture_recog.blob"
)
YOLOV8S_CONFIG = os.path.join(
    "gesture_recognition_model","yolov8s","yolov8strained.json"
)

TEST_DATA = glob.glob("test_data/*.jpg")
OUTPUT_IMAGES_YOLOv8n = os.path.join("results", "gesture_pred_images_v8n")
OUTPUT_IMAGES_YOLOv8s = os.path.join("results", "gesture_pred_images_v8s")
OUTPUT_VIDEO_YOLOv8n = os.path.join("results", "gesture_camera_v8n.mp4")
OUTPUT_VIDEO_YOLOv8s = os.path.join("results", "gesture_camera_v8s.mp4")

# define camera preview dimensions same as YOLOv8 model input size
CAMERA_PREV_DIM = (416, 416)

# define the class label names list
LABELS = ["Five", "Four", "One", "Three", "Two"]

On Lines 2 and 3, we import the os and glob modules. Then, from Lines 6-19, we define the following:

• path to the YOLOv8n and YOLOv8s hand gesture recognition models
• test data directory
• output locations for saving inference results of YOLOv8n and YOLOv8s on images and videos

From Lines 28-31, we also define the camera preview dimension and a list of class label names to help decode class predictions to human-readable class names.

Defining the Utilities

Now that the configuration has been defined, we can determine the utilities for creating images, camera DepthAI pipelines, and a few helper functions. The utils.py script defines several functions:

• Helps create the pipeline for hand gesture recognition on OAK with images
• Pipeline for hand gesture recognition on OAK with color camera stream
• Define a few helper functions for loading configuration files, annotating camera frames, and normalizing the predictions

Creating the Images Pipeline

# import the necessary packages
from pyimagesearch import config
import json
import numpy as np
import cv2
from pathlib import Path
import depthai as dai

def create_image_pipeline(config_path, model_path):
   # initialize a depthai pipeline
   pipeline = dai.Pipeline()

   # load model config file and fetch nn_config parameters
   print("[INFO] loading model config...")
   configPath = Path(config_path)
   model_config = load_config(configPath)
   nnConfig = model_config.get("nn_config", {})

On Lines 2-7, we import the necessary packages like config from the pyimagesearch module, json for loading model configuration files, cv2 for image-related operations, and most importantly, the depthai module along with the numpy and Path modules.

We define the function create_images_pipeline() that takes config_path and model_path as input parameters on Line 9. Then, a depthai pipeline is initialized on the host, which helps define the nodes, the flow of data, and communication between the nodes (Line 11).

On Lines 15 and 16, we load the model configuration file using the provided config_path and extract the nn_config object from the model configuration on Line 17.

   print("[INFO] extracting metadata from model config...")
   # using nnConfig extract metadata like classes,
   # iou and confidence threshold, number of coordinates
   metadata = nnConfig.get("NN_specific_metadata", {})
   classes = metadata.get("classes", {})
   coordinates = metadata.get("coordinates", {})
   anchors = metadata.get("anchors", {})
   anchorMasks = metadata.get("anchor_masks", {})
   iouThreshold = metadata.get("iou_threshold", {})
   confidenceThreshold = metadata.get("confidence_threshold", {})

On Line 22, we extract the relevant metadata NN_specific_metadata using the nnConfig variable. With the help of metadata, we extract information like classes, coordinates, anchors, iouThreshold, and confidenceThreshold. We would pass this information to the detection network node later.

   print("[INFO] configuring inputs and output...")
   # configure inputs for depthai pipeline
   # since this pipeline is dealing with images an XLinkIn node is created
   detectionIN = pipeline.createXLinkIn()
   # create a Yolo detection node
   detectionNetwork = pipeline.create(dai.node.YoloDetectionNetwork)
   # create a XLinkOut node for fetching the neural network outputs to host
   nnOut = pipeline.create(dai.node.XLinkOut)

   print("[INFO] setting stream names for queues...")
   # set stream names used in queue to fetch data when the pipeline is started
   nnOut.setStreamName("nn")
   detectionIN.setStreamName("detection_in")

Next, we create inputs for the depthai pipeline. First, we create an XLinkIn() node since this pipeline deals with images on Line 33. Then, a YoloDetectionNetwork node is created for performing detection and an XLinkOut node for sending the neural network output to the host on Lines 35-37.

On Lines 41 and 42, we set the output stream names, which would later be used in the queue to fetch image and neural network data when the pipeline is started.

   print("[INFO] setting YOLO network properties...")
   # network specific settings - parameters read from config file
   # confidence and iou threshold, classes, coordinates are set
   # most important the model .blob file is used to load weights
   detectionNetwork.setConfidenceThreshold(confidenceThreshold)
   detectionNetwork.setNumClasses(classes)
   detectionNetwork.setCoordinateSize(coordinates)
   detectionNetwork.setAnchors(anchors)
   detectionNetwork.setAnchorMasks(anchorMasks)
   detectionNetwork.setIouThreshold(iouThreshold)
   detectionNetwork.setBlobPath(model_path)
   detectionNetwork.setNumInferenceThreads(2)
   detectionNetwork.input.setBlocking(False)

With the detectionNetwork node defined, we set the YOLO network properties based on the metadata extracted from the model configuration file from Lines 48-56. One important thing to note is that we load the model weights (i.e., setBlobPath) on Line 54 by passing the model_path as a parameter for the YOLO model.

   print("[INFO] creating links...")
   # linking the nodes - image node output is linked to detection node
   # detection network node output is linked to XLinkOut input
   detectionIN.out.link(detectionNetwork.input)
   detectionNetwork.out.link(nnOut.input)

   # return the pipeline to the calling function
   return pipeline

On Lines 61 and 62, we link the detectionIN ( XLinkIn node) output to the YOLO detection network node (detectionNetwork) input and the YOLO detection network node output to the nnOut (XLinkOut node) input.

Finally, the pipeline object is returned to the calling function on Line 65.

Creating the Camera Pipeline

def create_camera_pipeline(config_path, model_path):
   # initialize a depthai pipeline
   pipeline = dai.Pipeline()

   # load model config file and fetch nn_config parameters
   print("[INFO] loading model config...")
   configPath = Path(config_path)
   model_config = load_config(configPath)
   nnConfig = model_config.get("nn_config", {})

   print("[INFO] extracting metadata from model config...")
   # using nnConfig extract metadata like classes,
   # iou and confidence threshold, number of coordinates
   metadata = nnConfig.get("NN_specific_metadata", {})
   classes = metadata.get("classes", {})
   coordinates = metadata.get("coordinates", {})
   anchors = metadata.get("anchors", {})
   anchorMasks = metadata.get("anchor_masks", {})
   iouThreshold = metadata.get("iou_threshold", {})
   confidenceThreshold = metadata.get("confidence_threshold", {})

   # output of metadata - feel free to tweak the threshold parameters
   #{'classes': 5, 'coordinates': 4, 'anchors': [], 'anchor_masks': {},
   # 'iou_threshold': 0.5, 'confidence_threshold': 0.5}
   print(metadata)

We define the function create_camera_pipeline() that takes config_path and model_path as input parameters on Line 67.

The rest of the code from Lines 69-86 is the same as we defined in the create_images_pipeline() function, so we skip the explanation here. But generally, we initialize a pipeline, load the configuration file, and extract metadata from the configuration file.

  print("[INFO] configuring source and outputs...")
   # define sources and outputs
   # since OAK's camera is used in this pipeline
   # a color camera node is defined
   camRgb = pipeline.create(dai.node.ColorCamera)
   # create a Yolo detection node
   detectionNetwork = pipeline.create(dai.node.YoloDetectionNetwork)
   xoutRgb = pipeline.create(dai.node.XLinkOut)
   # create a XLinkOut node for getting the detection results to host
   nnOut = pipeline.create(dai.node.XLinkOut)

   print("[INFO] setting stream names for queues...")
   # set stream names used in queue to fetch data when the pipeline is started
   xoutRgb.setStreamName("rgb")
   nnOut.setStreamName("nn")

We now establish input and output nodes for the camera pipeline, which differs from the previously defined image pipeline. The input source is an OAK device camera rather than images in this case.

On Lines 97-102, we create the following nodes:

• ColorCamera: an input to the detection network
• YoloDetectionNetwork: for performing inference
• XLinkOut: for displaying the detection results on the host computer

Then on Lines 106 and 107, we set stream names from respective xoutRgb and nnOut nodes to rgb and nn, which would be passed as a parameter to the OutputQueue for fetching the color frames and detections.

   print("[INFO] setting camera properties...")
   # setting camera properties like the output preview size,
   # camera resolution, color channel ordering and FPS
   camRgb.setPreviewSize(config.CAMERA_PREV_DIM)
   camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
   camRgb.setInterleaved(False)
   camRgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
   camRgb.setFps(40)

On Lines 112-116, we configure the camRgb (ColorCamera node) properties, such as setting the preview window resolution, selecting 1080P as the camera resolution, defining the color channel order as BGR, and establishing a camera frame rate of 40 FPS.

   print("[INFO] setting YOLO network properties...")
   # network specific settings - parameters read from config file
   # confidence and iou threshold, classes, coordinates are set
   # most important the model .blob file is used to load weights
   detectionNetwork.setConfidenceThreshold(confidenceThreshold)
   detectionNetwork.setNumClasses(classes)
   detectionNetwork.setCoordinateSize(coordinates)
   detectionNetwork.setAnchors(anchors)
   detectionNetwork.setAnchorMasks(anchorMasks)
   detectionNetwork.setIouThreshold(iouThreshold)
   detectionNetwork.setBlobPath(model_path)
   detectionNetwork.setNumInferenceThreads(2)
   detectionNetwork.input.setBlocking(False)

   print("[INFO] creating links...")
   # linking the nodes - camera stream output is linked to detection node
   # RGB frame is passed through detection node linked with XLinkOut
   # used for annotating the frame with detection output
   # detection network node output is linked to XLinkOut input
   camRgb.preview.link(detectionNetwork.input)
   detectionNetwork.passthrough.link(xoutRgb.input)
   detectionNetwork.out.link(nnOut.input)

   # return the pipeline to the calling function
   return pipeline

We can omit the explanation of Lines 122-130 since they are identical to the steps we performed while constructing the image pipeline.

Lines 137-139 create connections between the nodes in the DepthAI pipeline:

• camRgb.preview.link(detectionNetwork.input): This line connects the output of the camRgb.preview node (camera RGB preview) to the input of the detectionNetwork node (YoloDetectionNetwork). This connection provides the camera’s RGB preview data as input to the object detection network.
• detectionNetwork.passthrough.link(xoutRgb.input): This line links the passthrough output of the detectionNetwork node to the input of the xoutRgb node (XLinkOut). The passthrough output provides the original input image (RGB preview) to be sent to the host and the object detection results. This passthrough connection also helps maintain the synchronization between the RGB frames and the detection results.
• detectionNetwork.out.link(nnOut.input): This line connects the output of the detectionNetwork node (YoloDetectionNetwork) to the input of the nnOut node (XLinkOut). This node forwards the object detection results to the host for further processing or display.

The above connections establish a pipeline where the camera RGB preview data is input for the YOLO object detection network. The original input image and the detection results are sent to the host.

Finally, the pipeline object is returned to the calling function on Line 142.

Defining a Few Helper Functions

def load_config(config_path):
   # open the config file and load using json module
   with config_path.open() as f:
       config = json.load(f)
       return config

def annotateFrame(frame, detections, model_name):
    # loops over all detections in a given frame
    # annotates the frame with model name, class label,
    # confidence score, and draw bounding box on the object
    color = (0, 0, 255)
    for detection in detections:
        bbox = frameNorm(frame, (detection.xmin, detection.ymin, detection.xmax, detection.ymax))
        cv2.putText(frame, model_name, (20, 40), cv2.FONT_HERSHEY_TRIPLEX, 1,
                    color)
        cv2.putText(frame, config.LABELS[detection.label], (bbox[0] + 10, bbox[1] + 25), cv2.FONT_HERSHEY_TRIPLEX, 1,
                    color)
        cv2.putText(frame, f"{int(detection.confidence * 100)}%", (bbox[0] + 10, bbox[1] + 60),
                    cv2.FONT_HERSHEY_TRIPLEX, 1, color)
        cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
    return frame

On Lines 145-149, we define the load_config() function that takes config_path as an argument, opens the configuration file using context managers, and reads its contents. The file’s contents are loaded as a JSON object using the json module and returned to the calling function.

Next, on Lines 151-165, the annotateFrame() method is defined, which takes the following three arguments as an input:

• frame: The input image/frame.
• detections: A list of detected objects in the frame provided by the YOLOv8 object detection network.
• model_name: The name of the model used for object detection (YOLOv8n/YOLOv8s).

Then the function iterates over each detection in the detections list on Line 156. For each detection, it does the following:

• normalizes the bounding box coordinates using the frameNorm function (not shown in the provided code discussed later)
• annotates the frame with the model_name in the top-left corner
• annotates the frame with the class label of the detected object using the corresponding label from the config.LABELS list
• annotates the frame with the detected object’s confidence score (percentage)
• draws a bounding box around the detected object using the cv2.rectangle function

Finally, on Line 165, the annotateFrame function returns the annotated frame after adding the model name, class labels, confidence scores, and bounding boxes for each detected object. This annotated frame can be displayed, saved, or used for further processing.

def to_planar(arr: np.ndarray, shape: tuple) -> np.ndarray:
   # resize the image array and modify the channel dimensions
   resized = cv2.resize(arr, shape)
   return resized.transpose(2, 0, 1)

def frameNorm(frame, bbox):
   # nn data, being the bounding box locations, are in <0..1> range
   # normalized them with frame width/height
   normVals = np.full(len(bbox), frame.shape[0])
   normVals[::2] = frame.shape[1]
   return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int)

On Lines 168-171, we define the to_planar() method; it accepts two parameters as input

• arr: a numpy array, which represents the input image and
• shape: a tuple representing the desired shape for the output image

The purpose of the to_planar() method is to resize a given numpy array (image) to the desired shape and modify the channel dimensions using the transpose function.

On Line 173, the frameNorm() function accepts input parameters

• frame: a numpy array representing input frame/image
• bbox: list or numpy array containing the bounding box locations in the normalized range from 0 to 1

The function normalizes the bounding box locations with respect to the frame dimensions. On Lines 176-178:

• A normVals numpy array is initialized with the same length as the input bounding box array and is filled with the frame’s height.
• Then, every alternate value in normVals is set to the frame’s width.
• The input bounding box values are clipped to the range [0, 1] and multiplied by the normVals to scale the bounding box coordinates to the frame dimensions.
• Finally, the scaled bounding box coordinates are converted to integers using the astype(int) method.

Hand Gesture Recognition in Images

With the configurations and utilities implemented, we can finally get into the code walkthrough of recognizing hand gestures in images with OAK-D.

# import the necessary packages
from pyimagesearch import config
from pyimagesearch import utils
import argparse
import cv2
import depthai as dai

# parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("-m", "--model", help="Provide model name for inference",
                   default='yolov8n', type=str)
args = parser.parse_args()

We start by importing the necessary packages on Lines 6-10:

• config module from the pyimagesearch package, which contains configuration settings or constants
• utils from the pyimagesearch package, which provides utility functions for creating image and camera pipelines
• argparse for parsing command line arguments
• cv2 for performing image and video processing tasks
• depthai module for creating DepthAI pipelines for OAK-D

Then, on Lines 13-16, we create a command line argument -m or --model where users can provide the model name for hand gesture recognition which can be either YOLOv8n or YOLOv8s.

# initialize a depthai camera pipeline
print("[INFO] initializing a depthai images pipeline...")
model_name = args.model
# if the model for inference is nano variant create image pipeline
# with nano variant config file and model weights
if model_name.lower() == "yolov8n":
    pipeline = utils.create_image_pipeline(config_path=config.YOLOV8N_CONFIG,
                                        model_path=config.YOLOV8N_MODEL)
    output_image_path = config.OUTPUT_IMAGES_YOLOv8n

# if the model for inference is small variant create image pipeline
# with small variant config file and model weights
else:
    pipeline = utils.create_image_pipeline(config_path=config.YOLOV8S_CONFIG,
                                            model_path=config.YOLOV8S_MODEL)
    output_image_path = config.OUTPUT_IMAGES_YOLOv8s

On Line 21, we assign the model name provided via command line arguments to the model_name variable.

Then from Lines 24-34:

• If the model name equals yolov8n, we create an image pipeline using the utility function utils.create_image_pipeline with the YOLOv8 Nano variant configuration file and model weights (specified in the config module). Also, via the config module, set the output_image_path variable to the appropriate output directory for the YOLOv8 Nano variant.
• If the model name is not yolov8n, then we perform the same steps as above, but only this time, we initialize the image pipeline with YOLOv8 Small variant configuration and model weights file.

In summary, the code initializes an image pipeline using the DepthAI library. Then, it sets up the appropriate configuration, model weights, and output image path based on the model name provided by the user.

# pipeline defined, now the device is assigned and pipeline is started
with dai.Device(pipeline) as device:
   # define the queues that will be used in order to communicate with
   # depthai and then send our input image for predictions
   detectionIN = device.getInputQueue("detection_in")
   detectionNN = device.getOutputQueue("nn")

   print("[INFO] loading image from disk...")
   for img_path in config.TEST_DATA:
       # load the input image and then resize it
       image = cv2.imread(img_path)
       image_res = cv2.resize(image, config.CAMERA_PREV_DIM)
       # create a copy of image for inference
       image_copy = image.copy()
       # initialize depthai NNData() class which is fed with the
       # image data resized and transposed to model input shape
       nn_data = dai.NNData()
       nn_data.setLayer(
           "input",
           utils.to_planar(image_copy, config.CAMERA_PREV_DIM)
       )
       # send the image to detectionIN queue further passed
       # to the detection network for inference as defined in pipeline
       detectionIN.send(nn_data)

Next, on Lines 37-59, we set up the DepthAI device and process a list of test images for recognizing hand gestures using the previously initialized pipeline. The input images are resized and preprocessed before being fed to the YOLO network for inference. The processed images are then sent to the DepthAI device for hand gesture recognition.

On Line 37, we create a context for the DepthAI device using the with statement, which ensures that the device is properly closed after use. The device is initialized with the previously created pipeline.

Two queues for communicating with the OAK device are defined on Lines 40 and 41:

• detectionIN: input queue for sending images to the device for detection
• detectionNN: output queue for receiving the detection results from the device

On Lines 44-59, we iterate over the test images specified in the config.TEST_DATA variable. For each test image, we

• load the image using the cv2.imread function
• resize the loaded image to the dimensions specified in config.CAMERA_PREV_DIM using the cv2.resize method
• create a copy of the loaded image for preprocessing
• initialize a dai.NNData() object to store the preprocessed image data
• preprocess the copied image to the model input shape and set the preprocessed image data as the input layer of the dai.NNData() object
• send the preprocessed image data to the detectionIN queue, which is further passed to the YOLO detection network for inference as defined in the image pipeline in utils.py

       print("[INFO] fetching neural network output for {}".
           format(img_path.split('/')[1]))
       # fetch the neural network output
       inDet = detectionNN.get()
       # if detection is available for given image, fetch the detections
       if inDet is not None:
           detections = inDet.detections

           # if object detected, annotate the image
           image_res = utils.annotateFrame(image_res, detections, args.model.lower())

       # finally write the image to the output path
       cv2.imwrite(
           output_image_path +"/"+img_path.split('/')[1],
           image_res
       )

With the image loaded from the disk, preprocessed, and inferred with the YOLO detection network, as a final step, we now fetch the neural network predictions (hand gesture recognition) for each test image, annotate the image with the detections, and finally save the annotated image to the disk.

On Line 64, we fetch the neural network output for the current image using the detectionNN.get() method and store the result in the inDet variable.

Next, on Lines 66-70, we check if the fetched neural network output is not None, meaning detections are available for the given image:

• If detections are available, extract them from the inDet object using the detections attribute and store them in the detections variable.
• Annotate the resized image (image_res) with the detections using the utils.annotateFrame function. The function, as discussed before, draws bounding boxes and labels for each detected object on the image.

Finally, on Lines 73-76, the annotated image is written to the output path using the cv2.imwrite function.

Results

Great! Now that we have completed implementing the code for hand gesture recognition in images, it’s time to examine the results.

Figure 3 displays the hand gesture recognition outcomes on several test images using YOLOv8n when supplied to the OAK device. The figure shows that the converted (.blob) and optimized hand gesture recognition model identifies hand gestures in the test images well. However, among the six samples:

• one False Positive (2nd row, 3rd image classified one as two)
• one False Negative (1st row, 2nd image failed to recognize hand gesture five though it was tough).

Figure 4 showcases the hand gesture recognition results using the YOLOv8s variant on the same sample test images when executed on the OAK device. As observed in the figure, the YOLOv8s variant performed slightly better than the YOLOv8n variant, as it detected the hand gesture five (in the 1st row, 2nd image). However, this time the model failed to detect the hand gesture one (in the 2nd row, 3rd image), resulting in one False Negative.

Nonetheless, the overall performance was impressive. It would be intriguing to compare the performance of YOLOv8s and YOLOv8n when running hand gesture recognition on a camera stream.

We can safely conclude that our DepthAI implementation for hand gesture recognition in images is effective. However, the minor inaccuracies observed in the test images depend on the model, which could potentially be improved for even better performance.

Hand Gesture Recognition in Camera Stream

Fantastic! We have now acquired the skill of identifying hand gestures in images using an OAK device. Next, it’s time to take things to the next level and utilize the OAK’s color camera to categorize the frames. This transition to leveraging OAK’s camera stream is where we believe the OAK module truly shines in practical application.

# import the necessary packages
from pyimagesearch import config
from pyimagesearch import utils
from imutils.video import FPS
import argparse
import time
import cv2
import depthai as dai

# parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("-m", "--model", help="Provide model name for inference",
                   default='yolov8n', type=str)
args = parser.parse_args()

We start by importing the necessary packages on Lines 6-12

• config and utils modules from the pyimagesearch package, which contains configuration settings/constants and utility functions for creating OAK pipelines
• FPS module from imutils.video helps us compute how long the hand gesture recognition takes to perform inference on each frame
• argparse for parsing command line arguments
• cv2 for performing image and video processing tasks
• depthai module for creating DepthAI pipelines for OAK

Then, on Lines 15-18, we create a command line argument -m or --model where users can provide the model name for hand gesture recognition which can be either YOLOv8n or YOLOv8s.

When the script is run, args will be a namespace object containing the values of any arguments passed to the script. If the -m argument is not provided, args.model will have the default value of yolov8n. If the -m argument is provided, its value will be stored in args.model.

# initialize a depthai camera pipeline
print("[INFO] initializing a depthai camera pipeline...")
model_name = args.model
# if the model for inference is nano variant create camera pipeline
# with nano variant config file and model weights
if model_name.lower() == "yolov8n":
   pipeline = utils.create_camera_pipeline(config_path=config.YOLOV8N_CONFIG,
                                       model_path=config.YOLOV8N_MODEL)
   output_video = config.OUTPUT_VIDEO_YOLOv8n

# if the model for inference is small variant create camera pipeline
# with small variant config file and model weights
else:
   pipeline = utils.create_camera_pipeline(config_path=config.YOLOV8S_CONFIG,
                                           model_path=config.YOLOV8S_MODEL)
   output_video = config.OUTPUT_VIDEO_YOLOv8s

# set the video codec to use with video writer
fourcc = cv2.VideoWriter_fourcc(*'MJPG')
# create video writer object with parameters: output video path,
# video codec, frame rate of output video, and dimensions of video frame
out = cv2.VideoWriter(
   output_video,
   fourcc,
   20.0,
   config.CAMERA_PREV_DIM
)

On Lines 22-36, we initialize a depthai camera pipeline for hand gesture recognition using the YOLOv8 Nano or Small model.

Based on the model variant (determined by args.model), we select the appropriate configuration file and model weights for the YOLOv8 Nano or Small variant on Lines 26-36:

• If the model variant is YOLOv8 Nano (yolov8n), it creates a camera pipeline with the configuration file and model weights specified for the Nano variant.
• If the model variant is not YOLOv8 Nano, it assumes the model is the Small variant (yolov8s). It creates a camera pipeline with the configuration file and model weights specified for the Small variant.

We also set the output_video path where the inference results will be stored as a video.

Next, on Line 39, we set the video codec to MJPG using the cv2.VideoWriter_fourcc function.

Finally, on Lines 42-47, we create an OpenCV VideoWriter object that will be used to save the processed video with hand gesture recognition annotations. The VideoWriter is initialized with the following:

• The output video path differs for the Nano and Small variants.
• The MJPG video codec.
• The frame rate of the output video is set to 20 frames per second (FPS).
• The dimensions of the video frame are specified by config.CAMERA_PREV_DIM.

# pipeline defined, now the device is assigned and pipeline is started
with dai.Device(pipeline) as device:

   # output queues will be used to get the rgb frames
   # and nn data from the outputs defined above
   qRgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
   qDet = device.getOutputQueue(name="nn", maxSize=4, blocking=False)

   # initialize variables like frame, start time for NN FPS
   # also start the FPS module timer, define color pattern for FPS text
   frame = None
   startTime = time.monotonic()
   fps = FPS().start()
   counter = 0
   color2 = (255, 255, 255)

Then, we set up and start the depthai camera pipeline using the previously created pipeline configuration. We also initialize variables for processing video frames and calculating FPS.

On Line 50, the dai.Device class is used as a context manager to ensure the device is properly initialized and cleaned up. The pipeline configuration is passed to the device when it is created.

Next, on Lines 54 and 55, two output queues are created using the device.getOutputQueue() method, which is used to retrieve the RGB video frames and neural network (NN) inference results. These queues are named “rgb” and “nn” (remember we named them while creating the camera pipeline using .setStreamName()), with a maximum size of 4, and are set to non-blocking mode. In non-blocking mode, the oldest data will be overwritten with new data if the queue is full.

On Lines 59-63, several variables are initialized, such as

• frame: for storing the current video frame being processed,
• startTime: the starting time for calculating the neural network FPS, using the time.monotonic() function, which returns the current time in seconds,
• fps: An instance of an FPS utility class is created and started to measure the frames per second of the video processing,
• counter: A counter variable, initialized to 0, which may be used to count frames or other events during processing,
• color2: A color tuple (255, 255, 255) representing white in BGR color space, which will be used to annotate video frames with FPS information.

    print("[INFO] starting inference with OAK camera...")
    while True:
        # fetch the RGB frames and YOLO detections for the frame
        inRgb = qRgb.get()
        inDet = qDet.get()

        if inRgb is not None:
            # convert inRgb output to a format OpenCV library can work
            frame = inRgb.getCvFrame()
            # annotate the frame with FPS information
            cv2.putText(frame, "NN fps: {:.2f}".format(counter / (time.monotonic() - startTime)),
                        (2, frame.shape[0] - 4), cv2.FONT_HERSHEY_TRIPLEX, 0.8, color2)
            # update the FPS counter
            fps.update()

        if inDet is not None:
            # if inDet is not none, fetch all the detections for a frame
            detections = inDet.detections
            counter += 1

        if frame is not None:
            # annotate frame with detection results
            frame = utils.annotateFrame(frame, detections, args.model.lower())
            # display the frame with gesture output on the screen
            cv2.imshow(args.model.lower(), frame)


        # write the annotated frame to the file
        out.write(frame)

        # break out of the while loop if `q` key is pressed
        if cv2.waitKey(1) == ord('q'):
            break

# stop the timer and display FPS information
fps.stop()
print("[INFO] elapsed time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
# do a bit of cleanup
out.release()
cv2.destroyAllWindows()

With the pipeline started and queues defined, we start the infinite while loop, which continuously fetches RGB frames and YOLOv8 detections from two different input queues (qRgb and qDet) on Lines 66-69.

On Lines 71-78, if an RGB frame is available:

• It is converted to a format that the OpenCV library can process.
• Then annotated with the FPS (Frames Per Second) information. The FPS information is calculated using a counter, and the time elapsed since the start of the inference process.
• The FPS counter is updated every time an RGB frame is processed.

Next, on Lines 80-83, if a YOLO detection is available, the detections for that frame are fetched, and the counter is incremented.

Suppose a frame is available, and the detections for that frame have been fetched. In that case, the frame is annotated with the detection results using the annotateFrame() function from the utils module on Lines 85-87. The annotated frame is then displayed on the screen and written to an output file on Lines 89 and 93.

Finally, on Lines 96 and 97, we check if the user pressed the ‘q’ key, at which point the while loop is broken, and the inference process is stopped. The FPS information is then printed to the console, the output file is released, and the OpenCV windows are destroyed on Lines 100-105.

Results

We have done an amazing job implementing hand gesture recognition on the camera stream using OAK! And we are thrilled to share the fantastic inference results we achieved with both the YOLOv8n and YOLOv8s variants. Can you believe that YOLOv8s performed exceptionally well in recognizing all five hand gestures with barely any false positives or negatives? It is undoubtedly the clear winner here!

In terms of FPS, YOLOv8n achieved a speed of approximately 25 FPS, which is impressive. But wait, YOLOv8s is no slouch either, with a commendable speed of around 13 FPS.

The video below displays the hand gesture recognition inference results with the YOLOv8 Nano variant:

The video below displays the hand gesture recognition inference results with the YOLOv8 Small variant:

What's next? I recommend PyImageSearch University.

Course information:
76 total classes • 90 hours of on-demand code walkthrough videos • Last updated: May 2023
★★★★★ 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:

• ✓ 76 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 76 Certificates of Completion
• ✓ 90 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 500+ 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

Summary

Congratulations! You have now completed this tutorial, and we hope it has provided valuable insights into creating a hand gesture recognition application on OAK-D using YOLOv8n and YOLOv8s object detectors.

We began by outlining the project structure and proceeded with a detailed walkthrough of the process, which included the following:

• Exporting the YOLOv8 PyTorch model to OAK
• Configuring prerequisites
• Defining utilities
• Creating image and camera pipelines
• Defining helper functions

Following this, we implemented hand gesture recognition in images and camera streams using the previously defined pipelines and utilities.

To conclude the tutorial, we compared the results of the YOLOv8n and YOLOv8s variants, emphasizing the superior performance of the YOLOv8s variant in camera streams.

Additionally, we highlighted the near real-time performance achieved with both YOLOv8n and YOLOv8s variants for hand gesture recognition in camera streams, demonstrating their effectiveness in practical applications.

Citation Information

Sharma, A. “Hand Gesture Recognition with YOLOv8 on OAK-D in Near Real-Time,” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/92by6

@incollection{Sharma_2023_Hand-Gesture-Recognition-YOLOv8-OAK-D,
  author = {Aditya Sharma},
  title = {Hand Gesture Recognition with {YOLOv8} on {OAK-D} in Near Real-Time},
  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/92by6},
}

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!

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!

Download the code!

Website

The post Hand Gesture Recognition with YOLOv8 on OAK-D in Near Real-Time appeared first on PyImageSearch.

Table of Contents

Introduction to Causality in Machine Learning

In this tutorial, you will learn about causality and how to correctly interpret and use it for your machine learning projects.

This lesson is the 1st in a 4-part series on Causality in Machine Learning:

1. Introduction to Causality in Machine Learning (this tutorial)
2. Tools and Methodologies for Studying Causal Effect
3. A Brief Introduction to Do-Calculus
4. Studying Causal Effect with Microsoft’s Do-Why Library

It is a known phenomenon that Nobel laureates are avid chocolate lovers. This means, as we look at Nobel laureates’ food preferences, many of them love eating chocolates. But what if we reverse the relation?

Does eating chocolates help us get a Nobel prize? As avid chocolate lovers, we can confirm this is not the case. Why? What is the reason for such injustice, and how can we exploit that in machine learning? Let’s find out.

To learn how to understand and correctly interpret causality, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Introduction to Causality in Machine Learning

So, what does causal inference mean? What does it mean to cause? Is causality as simple as understanding cause and effect? Let us first understand a few ground rules before playing the game of causality. There is a lot of noise surrounding causality and its neighbor correlation.

For example, the following quote from François Chollet (as shown in Figure 1) makes the topic even more confusing.

Correlation and Causation

“Correlation does not mean Causation”: If you had a penny for every time you heard that line, you would probably be a millionaire by now. But what does it mean? Why does the correlation between two events do not imply causality? What is up?

Correlation is a relationship or connection between two variables where whenever one changes, the other is likely to also change. But a change in one variable doesn’t cause the other to change. That’s a correlation, but it’s not causation.

For example, if you see a lot of birds in the sky and it starts raining, it doesn’t mean that the birds caused the rain. They just happened at the same time.

Mathematically, correlation is measured by the correlation coefficient between `-1` and `1`. A correlation coefficient of `1` indicates a perfect positive correlation between two variables. In contrast, a correlation coefficient of `-1` means a perfect negative correlation between two variables. A correlation coefficient of `0` means no correlation between two variables.

Conversely, causation is a relationship between two variables where one variable causes the other variable to change.

At this point, you’re probably wondering: I have heard about causality a few hundred times, but do I need to care?

Let’s ask you a different question, then.

Case Study 1: A “Marvelous” Problem

How many times have you looked at the result of your model and wondered what-if the data was something other than what it trained on? You could write an algorithm that predicts the sales of comic books, and your model works well and produces high-accuracy predictions, but you need to know why. Or maybe it’s the opposite, your algorithm predicts completely wrong sales figures, and you need to figure out why.

Confused? Let’s look at an example.

Suppose Marvel hires you as a data scientist. There has been a recent rise in comic book sales, and you need to figure out why so the company can maintain the sales figures. After some data analysis, you conclude that there is a direct correlation between comic book sales and Disney+ subscriptions. But you still don’t have an exact reason, so you develop some scenarios to form a hypothesis.

Scenario 1: A Direct Cause

The comics display ads from Disney+, so increased comic book sales cause more Disney+ subscriptions.

Scenario 2: Reversing the Cause and Effect

Disney+ has shows about comic book characters. New fans would like to consume more of such content. So an increase in Disney+ subscriptions led to more Marvel Comics being bought.

Scenario 3: Investigating a Hidden Cause

Or it’s something quite different. All Marvel Cinematic Universe movies are on Disney+. When a new Marvel movie comes out, the hype around these characters leads to more comic book sales (and, in turn, more Disney+ subscriptions).

Causal Thinking

The possibility presented in Scenario 3 is a hidden cause.

Each node is a variable, and each arrow shows the direction of the causal connection.

Suppose we build an accurate model to predict when the user will buy more comic books. But that is all that our model would do. The problem statement could be better expressed by a question like: What would the user have done if we had done something differently?

So you might think: I get it; the model doesn’t always consider hidden causes, but it still works. The predictions are still accurate.

Having our observed and causal effects in the same direction is considered lucky. More often than not, we see the opposite in real-world data, as shown in Figure 5.

Case Study 2: Food App Conundrum

Let’s look at another example to understand this better: Suppose you have now been employed at a food ordering app as a Data Scientist. The product designer wants to introduce a new User Interface (UI), which she believes will engage more users. She has introduced the new UI to a select few people and wants you to make sense of the data that has been gathered.

You decide to judge the old UI and the new one on a common criterion: whether the user orders after opening the app. Let’s call this the opened-and-ordered criteria. Table 1 describes this.

Based on this data, it seems pretty straightforward. The new UI is the clear winner. But something strange happens when you condition the data on high-activity vs. low-activity users, as shown in Table 2.

This data discrepancy modeled on a different subset of the original data is called “Simpson’s paradox.”

“But hey, why does this discrepancy occur?”

Well, maybe, since the old UI is already tried and tested, it is shown to users only at a particular time of the day, around the evening, when app activity is generally high. The new UI might be shown at other times of the day when the number of lower-activity users who are likely to order food is less.

The causal diagram may look like the following in Figure 6.

Quiz Time!

The debate on what it means to cause something and how cause precedes effect has been going on forever. Most prominent contestants include but are not limited to:

• Aristotle
• Homer
• Einstein
• That French guy from Matrix 3
• Judea Pearl, and many more

However, sadly, until recently, the science of causal inference was not considered formal mathematics or even something that could be quantified and studied.

This was remedied hugely by the contributions of Judea Pearl (Graphical Causal Models) and Donald Rubin (Rubin causal model).

What's next? I recommend PyImageSearch University.

Course information:
76 total classes • 90 hours of on-demand code walkthrough videos • Last updated: May 2023
★★★★★ 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:

• ✓ 76 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 76 Certificates of Completion
• ✓ 90 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 500+ 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

Summary

So, in conclusion, we learned here that simply because two variables are correlated does not necessarily mean that one is the direct cause of the other. We can do more harm through our model if we do not look for hidden causes and their effects.

Therefore the idea of what would happen had something in the observation process been changed becomes an imperative question to ask.

This is called Counterfactual Thinking.

Counterfactual thinking, specifically about what-if scenarios, is the heart of causal inference.

Some ways we can answer a counterfactual question are:

• Randomization
• Natural Experiments
• Conditioning

Interestingly, the methods become much harder to use as we go down the list. But as we go up, the methods are more foolproof and give better validation results. We will closely examine these and the various underlying methodologies in the next blog post of this series. Till then, keep learning, and don’t stop asking the question: What-If?

References

• Causal Inference in Online Systems by Amit Sharma
• An awesome playlist for Causal Learning
• The Book of Why: The New Science of Cause and Effect

Citation Information

A. R. Gosthipaty and R. Raha. “Introduction to Causality in Machine Learning,” PyImageSearch, P. Chugh, S. Huot, K. Kidriavsteva, and A. Thanki, eds., 2023, https://pyimg.co/3sa8x

@incollection{ARG-RR_2023_CausalityIntroML,
  author = {Aritra Roy Gosthipaty and Ritwik Raha},
  title = {Introduction to Causality in Machine Learning},
  booktitle = {PyImageSearch},
  editor = {Puneet Chugh and Susan Huot and Kseniia Kidriavsteva and Abhishek Thanki},
  year = {2023},
  note = {https://pyimg.co/3sa8x},
}

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!

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!

Download the code!

Website

The post Introduction to Causality in Machine Learning appeared first on PyImageSearch.

Table of Contents

Training the YOLOv8 Object Detector for OAK-D

In this tutorial, you will learn to train a YOLOv8 object detector to recognize hand gestures in the PyTorch framework using the Ultralytics repository. The goal would be to train a YOLOv8 variant that can learn to recognize 1 of 5 hand gestures (e.g., one, two, three, four, and five) with good mean average precision (mAP). Furthermore, since this tutorial acts as a strong base for an upcoming tutorial, the trained YOLOv8 variant should be able to run inference in near real-time on the OpenCV AI Kit (OAK) that comes powered with the Intel MyriadX neural hardware accelerator.

This lesson is the first in our series on OAK 102:

1. Training the YOLOv8 Object Detector for OAK-D (this tutorial)
2. Gesture Recognition with YOLOv8 on OAK-D in Near Real-Time

To learn how to train a YOLOv8 object detector on a hand gesture dataset for OAK-D, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Training the YOLOv8 Object Detector for OAK-D

Introduction

Object detection is one of the most exciting problems in the computer vision domain. The progress in this domain has been significant; every year, the research community achieves a new state-of-the-art benchmark. And, of course, all of this wouldn’t have been possible without the power of Deep Neural Networks (DNNs) and the massive computation by NVIDIA GPUs.

It all started when Redmon et al. (2016) published the YOLO research community gem, “You Only Look Once: Unified, Real-Time Object Detection,” at the CVPR (Computer Vision and Pattern Recognition) Conference. YOLO, or YOLOv1, was the first single-stage object detection model. It quickly gained popularity due to its high speed and accuracy.

The authors continued from there. Redmon and Farhadi (2017) published YOLOv2 at the CVPR Conference and improved the original model by incorporating batch normalization, anchor boxes, and dimension clusters.

And then came the YOLO model wave. In 2023, we arrived at Ultralytics YOLOv8. Yes, you read it right! From the day YOLOv1 was out, a new version of YOLO was published every year with improvements in both speed and accuracy.

Today, YOLO is the go-to object detection model in the computer vision community since it is the most practical object detector focusing on speed and accuracy.

Figure 1 shows the progression in YOLO models from YOLOv1 to PP-YOLOv2. One interesting aspect in the figure is the YOLOv5 model by Ultralytics, published in the year 2020, and this year, they released yet another state-of-the-art object detection model, YOLOv8. And today’s tutorial is all about experimenting with YOLOv8 but for OAK-D.

If you would like to learn about the entire history of the YOLO family, we highly recommend you check out our series on YOLO!

A Primer on YOLOv8

YOLOv8 is the latest version of the YOLO object detection, classification, and segmentation model developed by Ultralytics. While writing this tutorial, YOLOv8 is a state-of-the-art, cutting-edge model. Like previous versions built and improved upon the predecessor YOLO models, YOLOv8 also builds upon previous YOLO versions’ success. The new features and improvements in YOLOv8 boost performance and accuracy, making it the most practical object detection model.

One key feature of YOLOv8 is its extensibility. It is designed as a framework that supports all previous versions of YOLO, making it easy to switch between versions and benchmark their performance. This makes YOLOv8 an ideal choice for users who want to take advantage of the latest YOLO technology while still being able to use their existing YOLO models.

Table 1 shows the performance (mAP) and speed (frames per second (FPS)) benchmarks of five YOLOv8 variants on the MS COCO (Microsoft Common Objects in Context) validation dataset at 640×640 image resolution on Ampere 100 GPU. All five models were trained on the MS COCO training dataset. The model benchmarks are shown in ascending order starting with YOLOv8n (i.e., the nano variant having the smallest model footprint to the largest model, YOLOv8x). We would be training the Nano and Small variant of YOLOv8 as it would fit well into the OAK’s computer power.

The innovation is not just limited to YOLOv8’s extensibility. Some more prominent innovations that directly relate to its performance and accuracy include

• a new backbone network
• a new anchor-free detection head
• a new loss function

YOLOv8 is also highly efficient and can run on various hardware platforms, from CPUs to GPUs to Embedded Devices like OAK. And as you already know, our goal is to run YOLOv8 on an embedded hardware platform (i.e., an OAK edge device).

Figure 2 compares YOLOv8 with previous YOLO versions: YOLOv7, YOLOv6, and Ultralytics YOLOv5. The comparison is made in two fashions: mAP vs. model parameters and mAP vs. Latency measured on A100 GPU. The figure shows that almost all the YOLOv8 variants achieve the highest mAP on the COCO validation dataset. Also, YOLOv8 has fewer model parameters and less Latency benchmarked on the NVIDIA Ampere 100 architecture.

Overall, YOLOv8 is hands down a powerful and flexible framework for object detection offered in PyTorch.

This tutorial is the first in our OAK-102 series, and we hope you have followed the series of tutorials in our OAK-101 series. If not, we highly recommend you check out the OAK-101 series, which will build a strong foundation for the OpenCV AI Kit. You will learn the OAK hardware and the software stack from the ground level, and not just that. For example, you would learn to train and deploy an image classification TensorFlow model on an OAK edge device.

This tutorial will cover more advanced Computer Vision applications and how to deploy these advanced applications onto the OAK edge device.

Now, let’s start with today’s tutorial and learn to train the hand gesture recognition model for OAK!

Configuring Your Development Environment

To follow this guide, you need to clone the Ultralytics repository and pip install all the necessary packages via the setup and requirements files.

Luckily, to run the YOLOv8 training, you can do a pip install on the ultralytics cloned folder, meaning all the libraries are pip-installable!

One good news is that YOLOv8 has a command line interface, so you do not need to run Python training and testing scripts. With just the yolo command, you get most functionalities like modes, tasks, etc. Do not worry; today’s tutorial will cover the important command line arguments!

$ git clone https://github.com/ultralytics/ultralytics
$ pip install ultralytics

Need Help Configuring Your Development Environment?

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 immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides 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!

About the Dataset

For today’s experiment, we will train the YOLOv8 model on the Hand Gesture Recognition Computer Vision Project dataset hosted on Roboflow.

These datasets are public, but we download them from Roboflow, which provides a great platform to train your models with various datasets in the Computer Vision domain. Even more interesting is that you can download the datasets in multiple formats like COCO JSON, YOLO Darknet TXT, and YOLOv8 PyTorch. This process saves time for writing helper functions to convert the ground-truth annotations to the format required by these object detection models.

YOLOv8 Label Format

Since we will train the YOLOv8 PyTorch model, we will download the dataset in YOLOv8 format. The ground-truth annotation format of YOLOv8 is the same as other YOLO formats (see Figure 4), so you could write a script on your own that does this for you. There is one text file with a single line for each bounding box for each image. For example, if four objects exist in one image, the text file would have four rows containing the class label and bounding box coordinates. The format of each row is

class_id center_x center_y width height

where fields are space-delimited, and the coordinates are normalized from 0 to 1. To convert to normalized xywh from pixel values:

1. divide x and the box width by the image’s width
2. divide y and the box height by the image’s height

Hand Gesture Recognition Dataset

This dataset contains 839 images of 5 hand gesture classes for object detection: one, two, three, four, and five. With the help of five fingers, one- to five-digit combinations are formed, and the object detection model is trained on these hand gestures with respective labels, as shown in Figure 5. The dataset is split into training, validation, and testing sets. The dataset comprises 587 training, 167 validation, and 85 testing images. Each image has a 416×416 resolution with only one object (or instance).

Figure 5 shows sample images from the dataset with ground-truth bounding boxes annotated in red, belonging to classes four, five, two, and three.

Since only one object (gesture or class) is present in each image, there are 587 regions of interest (objects) in 587 training images, meaning there is precisely one object per image. Based on the heuristic shown in Figure 6, class five contributes to more than 45% of the objects. In contrast, the remaining classes: one, two, three, and four, are under-represented relative to gesture class five.

The Python code for data visualization (Figure 5) and class distribution graph (Figure 6) computation is provided inside the Google Colab Notebook of this tutorial!

YOLOv8 Training

This section is the heart of today’s tutorial, where we will cover most of the tasks, including

• Selecting the model
• Downloading the dataset
• Creating the data configuration
• Understanding the YOLOv8 command line interface
• Training the YOLOv8 nano model
• Visualizing the YOLOv8 nano model artifacts
• Qualitative and quantitative evaluation of testing data
• Training the YOLOv8 small model
• Evaluating the YOLOv8 small variant on testing data

Selecting the Model

Figure 7 shows 5 YOLOv8 variants starting with the most miniature YOLOv8 nano model built for running on mobile and embedded devices to the YOLOv8 XLarge on the other end of the spectrum. For today’s experiment, we will work with mainly two variants: Nano and Small. We chose these two variants because our final goal is to run the YOLOv8 model on an OAK-D device that can recognize hand gestures. The figure shows that the Nano and Small model variants have smaller memory footprints than higher-end variants.

OAK-D, an embedded device, has computation constraints, which doesn’t mean that other higher-end variants like Medium and Large won’t work on OAK-D, but the performance (FPS) would be lesser. Hence, we choose Nano and Small as they balance accuracy and performance well.

One more observation from Figure 7 is that the mAP improvements from Medium to XLarge are minute. However, the algorithm processing time increases significantly, which would pose a problem for deploying these models on OAK devices.

Downloading the Hand Gesture Recognition Dataset

# Download the vehicles-open image dataset
!mkdir hand_gesture_dataset
%cd hand_gesture_dataset
!curl -L -s "https://universe.roboflow.com/ds/zafYqbWHn8?key=n1igBaphSm" > hand_gesture.zip
!unzip -q hand_gesture.zip
!rm hand_gesture.zip

On Lines 2 and 3, we create the hand_gesture_dataset directory and cd into the directory where we download the dataset. Then, on Line 4, we use the curl command and pass the dataset URL we obtained from the Hand Gesture Recognition Computer Vision Project. Finally, we unzip the dataset and remove the zip file on Lines 5 and 6.

Let’s look at the contents of the hand_gesture_dataset folder:

$tree /content/hand_gesture_dataset -L 2
/content/hand_gesture_dataset
├── data.yaml
├── README.dataset.txt
├── README.roboflow.txt
├── test
│   ├── images
│   └── labels
├── train
│   ├── images
│   └── labels
└── valid
    ├── images
    └── labels

9 directories, 3 files

The parent directory has 3 files, out of which only data.yaml is essential, and 3 subdirectories:

• data.yaml: Has the data-related configurations, such as
  • the train and valid data directory path
  • the total number of classes in the dataset
  • the name of each class
• train: Training images along with training labels
• valid: Validation images with annotations
• test: Test images and labels

Configuration Setup

Next, we will edit the data.yaml file to have the path and absolute path for the train and valid images.

# Create configuration
config = {
   "path": "/content/hand_gesture_dataset",
   "train": "train",
   "val": "valid",
   "test": "test",
   "nc": 5,
   "names": ['five', 'four', 'one', 'three', 'two']
}


with open("hand_gesture_dataset/data.yaml", "w") as file:
   yaml.dump(config, file, default_flow_style=False)

From Lines 3-7, we define the data path, train, validation, test, number of classes, and class names in a config dictionary.

Finally, on Lines 12 and 13, we:

1. open the existing data.yaml file that was downloaded along with the dataset
2. overwrite it with the contents in config
3. store it on the disk

Understanding the YOLOv8 Command Line Interface

The good news is that YOLOv8 also comes with a command line interface (CLI) and Python scripts, making training, testing, and exporting the models much more straightforward. In addition, the YOLOv8 CLI allows for simple single-line commands without needing a Python environment. For example, as shown in the shell blocks below, all tasks related to the YOLO model can be run from the terminal using the yolo command.

!yolo TASK MODE ARGS

Please note in the above command line that TASK, MODE, and ARGS are just placeholders you will need to replace with actual values, which we discuss next.

TASK is an optional parameter; if not passed, YOLOv8 will determine the task from the model type, which means it’s intelligently designed. The TASK can be detect, segment, or classify.

MODE is a required parameter that can be either train, val, predict, export, track, or benchmark. This parameter helps tell YOLOv8 whether you want to use it for

• training the model on a custom dataset
• validating a trained model
• making predictions with the trained weights on images/videos
• converting or exporting the trained model to a format that can be deployed
• training a YOLOv8 detection or segmentation model for use in conjunction with tracking algorithms like BoT-SORT or ByteTrack to perform object tracking on video streams
• benchmarking the YOLOv8 exports such as TensorRT for speed and accuracy (for example, see Table 1)

Finally, ARGS is an optional parameter with various custom configuration settings used during training, validation/testing, prediction, exporting, and all the YOLOv8 hyperparameters. Examples of ARGS can be image size, batch size, learning rate, etc. To learn more about all the available configurations, check out the default.yaml file in the Ultralytics repository.

In short, the YOLOv8 CLI is a powerful tool that allows you to operate YOLOv8 at the tip of your fingers by providing features such as

• model training
• model validation and testing
• exporting a trained model to various formats
• 10-15 types of data augmentations
• training logs
• model checkpoints
• mAP and loss plots
• file management

Let’s look at a few examples of how YOLOv8 CLI can be leveraged to train, predict, and export the trained model.

• Fine-tune a pretrained YOLOv8 nano detection model for 20 epochs with an initial learning_rate of 0.01.

!yolo train data=coco128.yaml model=yolov8n.pt epochs=20 lr0=0.01

• Predict a YouTube video using a pretrained YOLOv8 nano segmentation model at image size 320×320.

!yolo predict model=yolov8n-seg.pt source='https://youtu.be/Zgi9g1ksQHc' imgsz=320

• Export a YOLOv8n classification model to ONNX (Open Neural Network Exchange) format at image size 224×224.

!yolo export model=yolov8n-cls.pt format=onnx imgsz=224,224

Voila! Isn’t that surprising? How easy it was to perform training, prediction, and even model conversion in just one single command.

Training the YOLOv8n Model

Alright! We are almost ready to train the YOLOv8 nano and small object detection model. However, before we run the training, let’s understand a few parameters that we will use while training:

We define a few standard model parameters:

• imgsz: Image size or network input while training. The images will be resized to this value before being fed to the network. The preprocessing pipeline will resize them to 416 pixels.
• data: Path to the data .yaml file, which has training, validation, and testing data paths and class label information.
• batch: Number of images fed as a single batch into the network for a forward pass. You can modify it according to the GPU memory available. We have set it to 32.
• epochs: Number of times we want to train the model on the entire hand gesture training dataset. We will train the model for 20 epochs.
• model: Path to the base model we want to use for training. We use the nano model yolov8n from the YOLOv8 family.
• project: This will create a project directory inside the working directory (gesture_train_logs).
• name: Each time you run this model, it will create a subdirectory yolov8n under the project directory, which would have a lot of information on the model (e.g., weights, sample input images, a few validation predictions outputs, metrics plot, etc.).

!yolo train model=yolov8n.pt data=hand_gesture_dataset/data.yaml epochs=20 imgsz=416 \
batch=32 project=gesture_train_logs name=yolov8n device=0

The training will start if there are no errors, as shown below. The logs indicate that the YOLOv8 model would train with Torch version 1.13.1 on a Tesla T4 GPU, showing initialized hyperparameters.

The yolov8n.pt weights are downloaded, which means the YOLOv8n model is initialized with the parameters trained with the MS COCO dataset. Finally, we can see that two epochs have been completed with a mAP@0.5=0.238.

Downloading https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt to yolov8n.pt...
100% 6.23M/6.23M [00:00<00:00, 80.7MB/s]
Ultralytics YOLOv8.0.55 🚀 Python-3.9.16 torch-1.13.1+cu116 CUDA:0 (Tesla T4, 15102MiB)
yolo/engine/trainer: task=detect, mode=train, model=yolov8n.pt, data=hand_gesture_dataset/data.yaml, epochs=20, patience=50, batch=32, imgsz=416, save=True, save_period=-1, cache=False, device=0, workers=8, project=gesture_train_logs, name=yolov8n, exist_ok=False, pretrained=False, optimizer=SGD, verbose=True, seed=0, deterministic=True, single_cls=False, image_weights=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, show=False, save_txt=False, save_conf=False, save_crop=False, hide_labels=False, hide_conf=False, vid_stride=1, line_thickness=3, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, boxes=True, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=None, workspace=4, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, fl_gamma=0.0, label_smoothing=0.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0, cfg=None, v5loader=False, tracker=botsort.yaml, save_dir=gesture_train_logs/yolov8n
Downloading https://ultralytics.com/assets/Arial.ttf to /root/.config/Ultralytics/Arial.ttf...
100% 755k/755k [00:00<00:00, 17.2MB/s]
Overriding model.yaml nc=80 with nc=5

                   from  n    params  module                                       arguments                     
  0                  -1  1       464  ultralytics.nn.modules.Conv                  [3, 16, 3, 2]                 
  1                  -1  1      4672  ultralytics.nn.modules.Conv                  [16, 32, 3, 2]                
  2                  -1  1      7360  ultralytics.nn.modules.C2f                   [32, 32, 1, True]             
. . . . .
. . . . .
 21                  -1  1    493056  ultralytics.nn.modules.C2f                   [384, 256, 1]                 
 22        [15, 18, 21]  1    752287  ultralytics.nn.modules.Detect                [5, [64, 128, 256]]           
Model summary: 225 layers, 3011823 parameters, 3011807 gradients, 8.2 GFLOPs

Transferred 319/355 items from pretrained weights
TensorBoard: Start with 'tensorboard --logdir gesture_train_logs/yolov8n', view at http://localhost:6006/
AMP: running Automatic Mixed Precision (AMP) checks with YOLOv8n...
AMP: checks passed ✅
optimizer: SGD(lr=0.01) with parameter groups 57 weight(decay=0.0), 64 weight(decay=0.0005), 63 bias
train: Scanning /content/hand_gesture_dataset/train/labels... 587 images, 0 backgrounds, 0 corrupt: 100% 587/587 [00:00<00:00, 2371.06it/s]
train: New cache created: /content/hand_gesture_dataset/train/labels.cache
albumentations: Blur(p=0.01, blur_limit=(3, 7)), MedianBlur(p=0.01, blur_limit=(3, 7)), ToGray(p=0.01), CLAHE(p=0.01, clip_limit=(1, 4.0), tile_grid_size=(8, 8))
val: Scanning /content/hand_gesture_dataset/valid/labels... 167 images, 0 backgrounds, 0 corrupt: 100% 167/167 [00:00<00:00, 2200.35it/s]
val: New cache created: /content/hand_gesture_dataset/valid/labels.cache
Plotting labels to gesture_train_logs/yolov8n/labels.jpg... 
Image sizes 416 train, 416 val
Using 2 dataloader workers
Logging results to gesture_train_logs/yolov8n
Starting training for 20 epochs...

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       1/20      2.05G      1.315      3.383      1.556         23        416: 100% 19/19 [00:12<00:00,  1.57it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% 3/3 [00:02<00:00,  1.07it/s]
                   all        167        167    0.00357      0.974      0.119      0.064

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       2/20      2.41G      1.132       2.83      1.326         24        416: 100% 19/19 [00:10<00:00,  1.83it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% 3/3 [00:01<00:00,  2.07it/s]
                   all        167        167     0.0209      0.989      0.238      0.164

Voila! With this, you have learned to train a YOLOv8 nano object detector on a hand gesture recognition dataset you downloaded from Roboflow. Isn’t that amazing?

As discussed in the Understanding the YOLOv8 CLI section, YOLOv8 logs the model artifacts inside the runs directory, which we will look at in the next section.

Once the training is complete, you will see the output similar to the one shown below:

​​   Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
      19/20      2.41G     0.7323      1.017       1.07         11        416: 100% 19/19 [00:05<00:00,  3.46it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% 3/3 [00:01<00:00,  1.57it/s]
                   all        167        167      0.786      0.824      0.878      0.681

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
      20/20      2.41G     0.7141     0.9552      1.061         11        416: 100% 19/19 [00:05<00:00,  3.32it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% 3/3 [00:03<00:00,  1.15s/it]
                   all        167        167      0.805      0.772       0.86       0.67

20 epochs completed in 0.061 hours.
Optimizer stripped from gesture_train_logs/yolov8n/weights/last.pt, 6.2MB
Optimizer stripped from gesture_train_logs/yolov8n/weights/best.pt, 6.2MB

Validating gesture_train_logs/yolov8n/weights/best.pt...
Ultralytics YOLOv8.0.55 🚀 Python-3.9.16 torch-1.13.1+cu116 CUDA:0 (Tesla T4, 15102MiB)
Model summary (fused): 168 layers, 3006623 parameters, 0 gradients, 8.1 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% 3/3 [00:03<00:00,  1.13s/it]
                   all        167        167      0.786      0.824      0.877      0.681
                  five        167         77      0.801      0.857       0.92      0.696
                  four        167         21      0.814      0.832      0.937      0.726
                   one        167         19       0.76      0.789      0.813      0.646
                 three        167         27      0.829      0.815      0.845        0.7
                   two        167         23      0.726      0.826      0.873      0.637
Speed: 0.1ms preprocess, 2.1ms inference, 0.0ms loss, 2.0ms postprocess per image
Results saved to gesture_train_logs/yolov8n

The above results show that the YOLOv8n model achieved an mAP of 0.877@0.5 IoU and 0.681@0.5:0.95 IoU in all classes on the validation set. It also indicates class-wise mAP, and the model achieved the best score for gesture class four (i.e., 0.937 mAP@0.5 IoU). Moreover, since the training dataset is not huge, the model took hardly 3.66 minutes to complete the training for 20 epochs on a Tesla T4 GPU.

Visualizing Model Artifacts

Now that we have trained our model, let’s look at the results generated inside the gesture_train_logs directory.

All training results are logged by default to yolov8/runs/train with a new incrementing directory created for each run as runs/train/exp, runs/train/exp1, etc. However, while training the model, we passed the PROJECT and the RUN_NAME, so in this case, it does not create the default directory to log the training results. Hence, in this experiment, runs is yolov8n.

Next, let’s look at the files created in the experiment.

$ tree gesture_train_logs/yolov8n
gesture_train_logs/yolov8n
├── args.yaml
├── confusion_matrix.png
├── events.out.tfevents.1679594913.1b3064e8db41.10831.0
├── F1_curve.png
├── labels_correlogram.jpg
├── labels.jpg
├── P_curve.png
├── PR_curve.png
├── R_curve.png
├── results.csv
├── results.png
├── train_batch0.jpg
├── train_batch190.jpg
├── train_batch191.jpg
├── train_batch192.jpg
├── train_batch1.jpg
├── train_batch2.jpg
├── val_batch0_labels.jpg
├── val_batch0_pred.jpg
├── val_batch1_labels.jpg
├── val_batch1_pred.jpg
├── val_batch2_labels.jpg
├── val_batch2_pred.jpg
└── weights
    ├── best.pt
    └── last.pt

1 directory, 25 files

On Line 1, we use the tree command followed by the PROJECT and RUN_NAME, displaying various evaluation metrics and weights files for the trained object detector. As we can observe, it has a precision curve, recall curve, precision-recall curve, confusion matrix, prediction on validation images, and finally, the best and last epoch weights file in PyTorch format.

Now, look at a few images from the runs directory.

Figure 8 shows the training images batch with Mosaic data augmentation. There are 16 images clubbed together; if we pick one image from the 4th row × 1st column, we can see that the image combines four different images. We explain the concept of Mosaic data augmentation in the YOLOv4 post, so do check that out if you haven’t already.

Next, we look at the results.png, which comprises training and validation loss for bounding box, objectness, and classification. It also has the metrics: precision, recall, mAP@0.5, and mAP@0.5:0.95 for training (Figure 9).

Figure 10 shows the ground-truth images and the YOLOv8n model prediction on the Hand Gesture Recognition validation dataset. From the two images below, it is clear that the model did a great job detecting the objects. The model has no False Negative predictions; however, the model did have a few False Positive detections. For example, in the 1st row × 4th column, the model detected a class four hand gesture as class five, and a rather difficult one in the 2nd row × 4th column, a class five gesture was detected as class one. But overall, it did great on these images.

Figure 10: Ground-truth images (top) and YOLOv8n model prediction (bottom) on a sample validation dataset fine-tuned with all layers (source: image by the author).

Evaluating YOLOv8n on the Test Dataset

Now that the training is complete, we have also looked at the few artifacts generated during the training, like loss and mAP plots and YOLOv8n model prediction on the validation dataset. Next, let’s put our model to evaluation on the test dataset. To achieve this, we would write a HandGesturePredictor class.

from ultralytics import YOLO

class HandGesturePredictor:
   def __init__(self, model_path, test_folder_path):
       self.model = YOLO(model_path)
       self.test_folder = glob.glob(test_folder_path)

   def classify_random_images(self, num_images=10):
       # Generate num_images random numbers between
       # 0 and length of test folder
       random_list = random.sample(range(0, len(self.test_folder)), num_images)
       plt.figure(figsize=(20, 20))

       for i, idx in enumerate(random_list):
           plt.subplot(5,5,i+1)
           plt.xticks([])
           plt.yticks([])
           plt.grid(True)
           img = cv2.imread(self.test_folder[idx])
           results = self.model.predict(source=img)
           res_plotted = results[0].plot()
           # cv2_imshow(res_plotted)
           # convert the image frame BGR to RGB and display it
           image = cv2.cvtColor(res_plotted, cv2.COLOR_BGR2RGB)
           plt.imshow(image)
       plt.show()

On Line 1, we import the YOLO module from the ultralytics Python package. This would help us to load the trained YOLOv8n model weights directly as a parameter.

Then, on Line 3, we define the HandGesturePredictor class. On Lines 4-6, the class constructor is defined that takes two parameters: model_path and test_folder_path. We then use the model_path to initialize the YOLO model instance and store all the .jpg image paths using the glob module in the test_folder attribute.

On Lines 8-26, we define the classify_random_images method that takes num_images as an optional parameter (default value is 10). This parameter tells the number of images we would infer with trained hand gesture recognition YOLOv8 model and plot the results.

Further in classify_random_images:

• We generate a list of random numbers between 0 and the length of the test folder. This would ensure every run generates predictions for different sets of images.
• Next, we create a figure of 20×20 inches using the matplotlib Python package.
• Then, we start a for loop over each of the 10 test images, create a subplot in the current 20×20 figure with a grid of five rows and five columns, selecting the (i+1)th subplot.
• Inside the for loop, we read the image using OpenCV, perform object detection on the img using the YOLOv8n hand gesture recognition model, and store the results in the results variable.
• Continuing the loop, we call the ultralytics method .plot(), which creates a new image with overlaid object detection results. We convert the result into RGB color space and display all subplots with predictions using plt.show().

classifier = HandGesturePredictor("gesture_train_logs/yolov8n/weights/best.pt", "hand_gesture_dataset/test/images/*.jpg")
classifier.classify_random_images(num_images=10)

Now that we have the HandGesturePredictor class defined, we create a classifier instance of the class by passing in the best weights of the YOLOv8n hand gesture model and the test images path. The class instance then invokes the classify_random_images method with num_images set to 10.

Figure 11 shows the object detection predictions on the 10 test images we obtain by running the above code. The results show that the YOLOv8n hand gesture recognition model did a brilliant job, given that it’s the most lightweight model in the YOLOv8 family.

Figure 11: Ground-truth images (top) and YOLOv8n model prediction (bottom) fine-tuned with all layers (source: image by the author).

The best part is that the model did not miss any detections, and it did have a few False Positive detections, like detecting a class three hand gesture twice as a class five gesture and a class four gesture again as a class five. Well, if we look at the 1st row × 2nd image, we can clearly see that the confidence for both detections is less than 0.5, so we can ignore the detections with confidence scores less than 0.5.

Now that we have observed the qualitative results of the YOLOv8n hand gesture model, we run the quantitative evaluation of the model on the 85 test set images using the YOLO CLI in val mode.

# Validate YOLOv8n on hand gesture test data
!yolo val model=gesture_train_logs/yolov8n/weights/best.pt \
project=gesture_train_logs/yolov8n data=hand_gesture_dataset/data.yaml split=test

Figure 12 shows that the YOLOv8n hand gesture recognition model achieved an mAP of 0.824@0.5 IoU and 0.656@0.5:0.95 IoU in all classes on the test set. It also indicates class-wise mAP, and the model achieved the best score for gesture class two (i.e., 0.927 mAP@0.5 IoU).

Training the YOLOv8s Model

Alright! So now that we have trained the YOLOv8 nano model on the Hand Gesture Recognition dataset, let’s take one step further into the YOLOv8 family and train the YOLOv8 small variant on the same dataset, and find out which one trumps the other!

!yolo train model=yolov8s.pt data=hand_gesture_dataset/data.yaml epochs=20 imgsz=416 \
batch=32 project=gesture_train_logs name=yolov8s device=0

To train the YOLOv8 small variant, we need to change the model parameter to yolov8s.pt, the pretrained weights of the YOLOv8 small variant. Next, we also need to change the name (run name) parameter to yolov8s, which would create a directory inside the gesture_train_logs project directory.

20 epochs completed in 0.062 hours.
Optimizer stripped from gesture_train_logs/yolov8s/weights/last.pt, 22.5MB
Optimizer stripped from gesture_train_logs/yolov8s/weights/best.pt, 22.5MB

Validating gesture_train_logs/yolov8s/weights/best.pt...
Ultralytics YOLOv8.0.55 🚀 Python-3.9.16 torch-1.13.1+cu116 CUDA:0 (Tesla T4, 15102MiB)
Model summary (fused): 168 layers, 11127519 parameters, 0 gradients, 28.4 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% 3/3 [00:03<00:00,  1.23s/it]
                   all        167        167      0.803      0.803      0.871      0.688
                  five        167         77      0.793      0.766      0.885      0.683
                  four        167         21      0.732       0.81      0.832      0.624
                   one        167         19      0.786      0.842      0.884      0.733
                 three        167         27       0.84      0.778      0.849      0.717
                   two        167         23      0.862      0.818      0.904      0.683
Speed: 0.6ms preprocess, 2.5ms inference, 0.0ms loss, 2.3ms postprocess per image
Results saved to gesture_train_logs/yolov8s

The above results show that the YOLOv8n model achieved an mAP of 0.871@0.5 IoU and 0.688@0.5:0.95 IoU in all classes on the validation set. It also indicates class-wise mAP, and the model achieved the best score for gesture class two (i.e., 0.904 mAP@0.5 IoU).

Since the training dataset is not huge, the model took hardly 3.72 minutes to complete the training for 20 epochs on a Tesla T4 GPU.

A few surprising findings after training YOLOv8s on the Hand Gesture dataset are:

• The mAP@0.5 IoU is slightly less than the YOLOv8n model, and the mAP@0.5:0.95 IoU is marginally better than YOLOv8n.
• The time taken to train both variants is also quite similar; there’s hardly a difference of a few seconds.

It would be interesting to see how the YOLOv8s model performs qualitatively and quantitatively on the test dataset. So let’s find out in the next section!

Evaluating YOLOv8s on the Test Dataset

Similar to the YOLOv8n evaluation, we put the YOLOv8s hand gesture variant to qualitative and quantitative assessments on the test dataset.

For the qualitative analysis, we create a classifier instance of the HandGesturePredictor class by passing in the best weights of the YOLOv8s hand gesture model and test images path. The class instance then invokes the classify_random_images method with num_images set to 10.

classifier = HandGesturePredictor("gesture_train_logs/yolov8s/weights/best.pt", "hand_gesture_dataset/test/images/*.jpg")
classifier.classify_random_images(num_images=10)

Figure 13 shows the object detection predictions on the 10 test images we obtain by running the above code. From the results, we can see that the YOLOv8s hand gesture recognition model does a better job than the YOLOv8n model. In fact, there are no False Positive predictions made by the model. Of course, the images are sampled randomly, and the best comparison can be made only if the same set of images is used with the YOLOv8s hand gesture model as with YOLOv8n.

Figure 13: Ground-truth images (top) and YOLOv8s model prediction (bottom) fine-tuned with all layers (source: image by the author).

However, we would better understand the quantitative (mAP scores) analysis improvements.

Next, we run the quantitative evaluation of the YOLOv8s hand gesture model on the 85 test set images using the YOLO CLI in val mode.

# Validate YOLOv8n on hand gesture test data
!yolo val model=gesture_train_logs/yolov8s/weights/best.pt \
project=gesture_train_logs/yolov8n data=hand_gesture_dataset/data.yaml split=test

Figure 14 shows that the YOLOv8n hand gesture recognition model achieved an mAP of 0.887@0.5 IoU and 0.706@0.5:0.95 IoU in all classes on the test set. It also indicates class-wise mAP, and the model achieved the best score for gesture class five (i.e., 0.93 mAP@0.5 IoU).

Comparing the results with the YOLOv8n hand gesture model, we can observe a significant improvement in the mAP scores across all five classes.

What's next? I recommend PyImageSearch University.

Course information:
75 total classes • 86 hours of on-demand code walkthrough videos • Last updated: April 2023
★★★★★ 4.84 (128 Ratings) • 15,900+ 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:

• ✓ 75 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 75 Certificates of Completion
• ✓ 86 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 500+ 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

Summary

We have now reached the end of this tutorial, and we hope you have gained valuable insights into training the YOLOv8 object detector for OAK-D. In this tutorial, we provided a comprehensive guide on training the YOLOv8 object detector for the OAK-D device.

We started by giving an introduction to YOLOv8 and discussed its quantitative benchmarks with previous YOLO versions. The tutorial then discussed the dataset used for training, specifically focusing on the hand gesture recognition dataset and YOLOv8 label format.

The training process is explained in detail, including

1. selecting the appropriate model
2. downloading the dataset
3. setting up the configuration
4. using the YOLOv8 Command Line Interface (CLI)

We then covered the training and evaluation of two different YOLOv8 models (i.e., YOLOv8n and YOLOv8s) with visualization of model artifacts and evaluation on the test dataset.

This tutorial serves as a foundation for an upcoming tutorial, where we will deploy the gesture recognition model on the OAK device and perform inference using the DepthAI API on images and camera streams. Stay tuned for the next tutorial in this series to dive deeper into the deployment and practical applications of the trained model.

Citation Information

Sharma, A. “Training the YOLOv8 Object Detector for OAK-D,” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/9qcei

@incollection{Sharma_2023_YOLOv8-OAK-D,
  author = {Aditya Sharma},
  title = {Training the {YOLOv8} Object Detector 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/9qcei},
}

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!

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!

Download the code!

Website

The post Training the YOLOv8 Object Detector for OAK-D appeared first on PyImageSearch.

Table of Contents

What’s Behind PyTorch 2.0? TorchDynamo and TorchInductor (primarily for developers)

In our previous lesson of this series, we saw how PyTorch 2.0’s new feature, torch.compile, can provide significant speedups of 30-200% in eager mode for most models we run daily. PyTorch 2.0 introduces the following new technologies:

• TorchDynamo
• TorchInductor
• AOT Autograd
• PrimTorch

These technologies make the PyTorch 2.0 code run faster (with less memory) by JIT-compiling the PyTorch 2.0 code into optimized kernels, all while requiring minimal code changes.

This lesson covers these technologies in more detail and digs deeper into the usage and behavior of both TorchDynamo and TorchInductor. We will see how we can accelerate our models just like torch.compile but using TorchDynamo/TorchInductor and demonstrate some of the advantages of both TorchDynamo and TorchInductor over previous compiler solutions like TorchScript and FX (Function eXtraction) Tracing. This tutorial is primarily for developers who want to accelerate their deep learning models with PyTorch 2.0.

In this series, you will learn about Accelerating Deep Learning Models with PyTorch 2.0.

This lesson is the 2nd of a 2-part series on Accelerating Deep Learning Models with PyTorch 2.0:

1. What’s New in PyTorch 2.0? torch.compile
2. What’s Behind PyTorch 2.0? TorchDynamo and TorchInductor (primarily for developers) (this tutorial)

To learn what’s behind PyTorch 2.0, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

What’s Behind PyTorch 2.0? TorchDynamo and TorchInductor (primarily for developers)

We will start with understanding the new technologies behind PyTorch 2.0.

Technology Behind PyTorch 2.0

TorchDynamo

TorchDynamo (shown in Figure 1) is PyTorch’s latest compiler solution that leverages JIT (Just In Time) compilation to transform a general Python program into an FX Graph. The FX Graph is an intermediate representation of your code that can be further compiled and optimized. TorchDynamo extracts FX Graphs by inspecting Python bytecode at runtime and detecting calls to PyTorch operations.

Additionally, TorchDynamo is designed to mix Python execution with compiled backends to get the best of both worlds: usability and performance. Finally, it can support the arbitrary program by breaking the FX Graph whenever it encounters any unsupported feature (e.g., data-dependent flows or non-PyTorch libraries).

The approach is based on a CPython feature known as Frame Evaluation API that can safely and correctly capture graphs 99% of the time without incurring additional overhead. Other previous compiler solutions like TorchScript and FX Tracing need help to capture graphs even 50% of the time, and that too with additional overhead.

Here’s how TorchDynamo’s behavior differs from the default Python behavior.

TorchInductor

TorchDynamo should be paired with a backend that can efficiently convert captured FX Graphs into fast machine code to make PyTorch 2.0 programs faster. However, a lot is lost while exporting the FX Graphs in different existing backends. Some have fundamentally different execution models than PyTorch, and others are only optimized for inference and not training.

TorchInductor is the new compiler backend that compiles the FX Graphs generated by TorchDynamo into optimized C++/Triton kernels. TorchInductor uses a Pythonic define-by-run loop level intermediate representation to automatically map PyTorch models into generated Triton code on GPUs and C++/OpenMP on CPUs. TorchInductor’s core loop level intermediate representation contains only ~50 operators, and it is implemented in Python, making it easily hackable and extensible.

AOT Autograd

AOT Autograd is PyTorch’s new automatic differentiation engine that leverages PyTorch’s torch_dispatch mechanism to trace through the autograd engine and create ahead-of-time backward passes. This accelerates both forward and backward passes.

PrimTorch

With more than 2000+ operators, writing a backend for PyTorch is challenging and draining. Hence, with PrimTorch (Figure 2), PyTorch 2.0 canonicalizes these operations into smaller primitive operations. As a result, there are mainly two operation sets in PrimTorch:

• Prim ops: contain ~250 low-level operations are well suited for compilers and can be fused for good performance.
• ATen ops: contain ~750 canonical operators suited for backends that won’t have a compilation to recover performance from a lower-level operator set like Prim ops.

Configuring Your Development Environment

Installation

As in previous versions, PyTorch 2.0 is available as a Python pip package. However, to successfully install PyTorch 2.0, your system should have installed the latest CUDA (Compute Unified Device Architecture) versions (11.6 and 11.7). Here’s how you can install PyTorch 2.0 nightly version via pip:

For CUDA version 11.7:

$ pip3 install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

For CUDA version 11.6:

$ pip3 install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu116

However, if you don’t have CUDA 11.6 or 11.7 installed on your system, you can download all the required dependencies in the PyTorch nightly binaries with docker.

$ sudo apt install -y nvidia-docker2
$ sudo systemctl restart docker
$ docker pull ghcr.io/pytorch/pytorch-nightly
$ docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

Be sure to specify --gpus all so your container can access all your GPUs.

Verification

Optionally, you can verify your installation via:

$ git clone https://github.com/pytorch/pytorch
$ cd tools/dynamo
$ python verify_dynamo.py

Also, ensure you have a C++ compiler installed. You can install the g++ compiler via:

$ sudo apt-get install g++

Need Help Configuring Your Development Environment?

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 immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides 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

We first need to review our project directory structure.

Start by accessing the “Downloads” section of this tutorial to retrieve the source code.

From there, take a look at the directory structure:

├── dynamo_speedup.py
├── data_dependent_cf.py
├── non_pytorch_function.py
├── dynamo_graph.py

The project directory contains four files.

• First, the dynamo_speedup.py file demonstrates how to achieve speedups on real models using both TorchDynamo and TorchInductor.
• Next, the file data_dependent_cf.py and non_pytorch_function.py demonstrates the advantages of TorchDynamo over existing compiler solutions like TorchScript and FX Tracing.
• Finally, the file dynamo_graph.py shows how we can analyze and view graphs and graph breaks created by TorchDynamo. We will discuss these files in detail in subsequent sections.

Accelerating DNNs Using TorchDynamo

We will start our hands-on exercise by looking at speeding up deep neural networks (DNNs) like torch.compile but using TorchDynamo. First, look into the dynamo_speedup.py file in the project directory.

import torch
import torch._dynamo as dynamo
import time
import numpy as np

# running a model
def run_model(model, inputs, steps=20):
    # load model on GPU
    model = model.cuda() 
    # define an optimizer
    optimizer = torch.optim.Adam(model.parameters()) 
    times = []
    for step in range(steps):
        begin = time.time()
        # zero gradients
        optimizer.zero_grad() 
        # forward pass
        output = model(inputs.cuda()) 
        # back propagate 
        output.sum().backward() 
        # optimize weights
        optimizer.step() 
        end = time.time()
        # calculate step time
        times.append(float(end - begin)) 
        print(f"Time for {step}-th forward pass is {end - begin}")
    
    # calculate median step time
    median = np.median(times) 
    print("Median step time is {:.3f} seconds".format(median))
    return median

On Lines 1-4, we import the torch, numpy, and time libraries. Next, we import the TorchDynamo engine from torch._dynamo. Then on Lines 7-31, we define the run_model() function that takes model, inputs, and steps as arguments and runs the model training on inputs for a given number of steps. First, on Line 9, we load the model on the GPU. Then on Line 11, we define optimizer over model parameters.

Finally, on Lines 13-26, we run the model training for given steps wherein we pass the given inputs, backpropagate gradients, and update network weights in each step. We print and store the time taken by each step in a list of times. Finally, on Lines 29 and 30, we calculate and print the median step time taken by our compiled model.

# loading pretrained resnet50
model = torch.hub.load('pytorch/vision:v0.10.0', "resnet50", pretrained=True) 
# random input image
inputs = torch.randn(256, 3, 224, 224) 

# run original model
print("Running Original ResNet50...")
time1 = run_model(model, inputs, 10) 

# optimize model using dynamo
optimized_model = dynamo.optimize("inductor")(model)
# run optimized model 
print("Running Optimized ResNet50...")
time2 = run_model(optimized_model, inputs, 10)

#calculate training speedup
print("Speedup : {:.3f}%".format(100*(time1-time2)/time1))

Next, on Lines 34-36, we load a pre-trained ResNet50 model using torch.hub and define a dummy input image inputs with batch size 256. On Line 40, we run the original model and calculate its median step time. Then on Line 43, we optimize the model with dynamo.optimize() using the inductor backend. This function is analogous to torch.compile and is expected to provide similar results. Next, on Line 46, we run the optimized_model and calculate its median step time. Finally, on Line 49, we calculate the speedup given by optimized_model over the original model.

After running the above script using python dynamo_speedup.py, your output should look similar to Figure 3.

By using TorchDynamo, we can achieve a 13% speedup for ResNet50. Note that your numbers differ slightly depending on your hardware accelerator. This experiment is done on NVIDIA A6000s. Newer GPUs like A100s are likely to provide more speedups.

Also, notice that TorchDynamo and TorchInductor take much longer to complete in eager mode at the first step. This is because TorchDynamo and TorchInductor compile the model into optimized kernels as it executes.

TorchDynamo vs. TorchScript vs. FX Tracing

Now that we have seen that TorchDynamo can provide speedup for real models, it is important to understand in what other aspects it is better than previous compiler solutions like TorchScript and FX Tracing. The primary advantage of TorchDynamo lies in its ability to handle arbitrary Python code without significant changes. In this section, we will discuss these advantages in detail.

Data-Dependent Control Flow

Data-dependent control flows are one case that TorchDynamo can handle that other compilers fail to capture. Data-dependent control flow means that the flow of data depends on the value of the data itself. The file data_dependent_cf.py in the project directory illustrates this concept.

import torch
import torch._dynamo as dynamo
import torch.jit as jit
import torch.fx as fx
import traceback as tb
import warnings
warnings.simplefilter("ignore", UserWarning)

# define a function with data-dependent control flow
def f1(x):
    if x.sum() < 0:
        return torch.tensor(0)
    return torch.tensor(1)

# Check if output of two functions is same for given input
# fn1 is the original function
# fn2 is the compiled function
def test_fns(f1, f2, x):
    out1 = f1(x)
    out2 = f2(x)
    return torch.allclose(out1, out2)

On Lines 1-6, we import standard libraries and packages like torch, TorchDynamo torch._dynamo, TorchScript torch.jit, and FX Tracing torch.fx. Next, on Lines 10-13, we define a function f1() that takes a PyTorch tensor x and returns 0 or 1 depending on the polarity of its sum. Here, Line 11 (i.e., if x.sum() < 0:) is an example of data-dependent control flow.

Then on Lines 18-21, we define a function test_fns that takes two functions, f1 and f2, as arguments along with an input x to check whether both functions return the same output for the same set of input x. This function will tell us whether our original or compiled function has the same behavior.

# define random input
x = torch.randn(5, 5)

# compile using TorchScipt
traced_f1 = jit.trace(f1, (x,))
# compare f1, and traced_f1 on -x
print("TorchScript:", test_fns(f1, traced_f1, (-x)))

# compile using FX
fx_f1 = fx.symbolic_trace(f1, concrete_args={"x": x})
# compare f1, and fx_f1 on -x
print("FX:", test_fns(f1, fx_f1, (-x)))

# compile using Dynamo
dynamo_f1 = dynamo.optimize("inductor")(f1)
# compare f1, and dynamo_f1 on -x
print("TorchDynamo:", test_fns(f1, dynamo_f1, (-x)))

Next, on Line 24, we define a dummy input x. On Line 27, we compile the function f1 using TorchScript jit.trace(f1, (x,)) that takes the original function f1 and a dummy input x. On Line 29, we compare the compiled model traced_f1 and the original function f1 on -x instead of x.

On Line 32, we compile the function f1 using FX Tracing fx.symbolic_trace(f1, concrete_args={"x": x}) that also takes the original function f1 and a dummy input x. On Line 34, we compare the compiled model fx_f1 and the original function f1 on -x instead of x.

Finally, on Line 37, we compile the function f1 using TorchDynamo dynamo.optimize("inductor")(f1) that takes the original function f1 as input. Then, on Line 39, we compare the compiled model dynamo_f1 and the original function f1 on -x.

After running the script python data_dependent_cf.py, your output should look like Figure 4.

As you can see from the above output, TorchScript and FX Tracing fails to correctly capture the behavior of our function when a data-dependent control flow is present. This is because these solutions only capture or trace the path taken by the provided dummy input, which is x in our case. However, when we flip the input to -x, they follow the same old traced path and yield the wrong result.

TorchDynamo, however, correctly captures the function and yields the correct result. It does so by breaking the computation graph whenever it encounters a data-dependent control flow and letting the Python interpreter handle the unsupported code and resume graph capturing afterward.

Non-PyTorch Libraries

Another case where TorchDynamo has an advantage over TorchScript and FX Tracing is when the program contains non-PyTorch functions. The file non_pytorch_function.py in the project directory illustrates this use case. For this example, you need to install scipy, which you can install via pip install scipy.

import torch
import torch._dynamo as dynamo
import torch.jit as jit
import torch.fx as fx
import traceback as tb
import scipy
import warnings
warnings.simplefilter("ignore", UserWarning)

# define a function with non-PyTorch library
def f1(x):
    x = x * 2
    x = scipy.fft.dct(x.numpy())
    x = torch.from_numpy(x)
    x = x * 2
    return x

# Check if output of two functions is same for given input
# fn1 is the original function
# fn2 is the compiled function
def test_fns(f1, f2, x):
    out1 = f1(x)
    out2 = f2(x)
    return torch.allclose(out1, out2)

# define random input
x = torch.randn(5, 5)

# compile using TorchScipt
traced_f1 = jit.trace(f1, (x,))
# compare f1, and traced_f1 on -x
print("TorchScript:", test_fns(f1, traced_f1, (-x)))

# compile using FX
fx_f1 = fx.symbolic_trace(f1, concrete_args={"x": x})
# compare f1, and fx_f1 on -x
print("FX:", test_fns(f1, fx_f1, (-x)))

# compile using Dynamo
dynamo_f1 = dynamo.optimize("inductor")(f1)
# compare f1, and dynamo_f1 on -x
print("TorchDynamo:", test_fns(f1, dynamo_f1, (-x)))

On Lines 1-7, we import standard libraries and packages like torch, scipy, TorchDynamo torch._dynamo, TorchScript torch.jit, and FX Tracing torch.fx. Next, on Lines 11-16, we define a function f1() that takes a PyTorch tensor x and operates on it using functions from the scipy package. Note that scipy.fft.dct() is a non-PyTorch function. The rest of the code (Lines 21-42) is the same as in the previous example.

After running the script python non_pytorch_function.py, your output should look like Figure 5.

As with the previous example, we can see that TorchScript and FX Tracing fail whenever they encounter any non-PyTorch code. This is because these solutions treat output from non-PyTorch functions as constants in the computation graph which can silently yield wrong results.

TorchDynamo successfully captures the graph even when non-Pytorch functions are used.

It should be noted that it’s different from TorchScript or FX Tracing and cannot handle data-dependent control flows or non-PyTorch functions. They can handle these cases; however, the solution can require significant code changes and will raise errors when unsupported Python is used.

TorchDynamo FX Graphs

We will now understand how we can analyze the FX Graphs captured by TorchDynamo, along with graph breaks. Generally, TorchDynamo captures an arbitrary Python code and converts it into an FX Graph, which TorchInductor then consumes to yield optimized kernels. Hence, TorchDynamo can work even with other backends. First, however, we will implement our custom backend to analyze the FX Graphs. We will refer to the dynamo_graph.py file in our project directory for this demo.

import torch
import torch._dynamo as dynamo
from typing import List

# define a custom backend to print FX graph
def custom_backend(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print("custom backend called with FX graph:")
    gm.graph.print_tabular()
    return gm.forward

# Reset since we are using a different backend (a custom one).
dynamo.reset()

# define a function with data-dependent control flow
def f1(x, y):
    z = torch.nn.functional.sigmoid(y)
    if x.sum() < 0:
        x = x*-1
    return z*x

x, y = torch.randn(10), torch.randn(10)

explanation, out_guards, graphs, ops_per_graph, break_reasons, explanation_verbose = dynamo.explain(f1, x, y)
print(explanation_verbose)
for graph in graphs:
    print(graph.graph.print_tabular())

On Lines 1-3, we import the torch and torch._dynamo packages. Then on Lines 6 and 7, we implement our custom_backend() function that takes an FX Graph instance gm and a list of input tensors example_inputs and prints the computation graph. On Line 12, we subsequently reset the dynamo engine to use our custom backend. Next, on Lines 15-19, we implement a function f1 ( with data-dependent control flow) that takes two inputs, x and y, and returns an output based on the polarity of x.sum().

On Line 21, we define dummy inputs. Then, on Lines 23-26, we use dynamo.explain to print the graphs and graph breaks in the code.

After running the script python dynamo_graph.py, you will see the following output on your console (Figure 6).

From Figure 6, we can see that TorchDynamo extracts two graphs from the code corresponding to graph break at the statement if x.sum() < 0: x = -x.

• Graph 1: z = torch.nn.functional.sigmoid(y); x = x*-1; z*x
• Graph 2: z = torch.nn.functional.sigmoid(y); z*x

So whenever x.sum() < 0, dynamo executes Graph 1, i.e., evaluates the first statement y = torch.nn.functional.sigmoid(y) and then lets the Python interpreter evaluate the conditional statement and then runs statement return y*x. However, when x.sum() >= 0, dynamo executes Graph 2 where it first evaluates y = torch.nn.functional.sigmoid(y) and then lets the Python interpreter evaluate the conditional statement and then run statement return y*x.

This highlights a significant difference between TorchDynamo and previous PyTorch compiler solutions. When encountering unsupported Python features, previous solutions either raise an error or silently fail. TorchDynamo, on the other hand, will break the computation graph.

What's next? I recommend PyImageSearch University.

Course information:
75 total classes • 86 hours of on-demand code walkthrough videos • Last updated: April 2023
★★★★★ 4.84 (128 Ratings) • 15,900+ 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:

• ✓ 75 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 75 Certificates of Completion
• ✓ 86 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 500+ 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

Summary

TorchDynamo and TorchInductor are the primary technologies that power PyTorch 2.0. TorchInductor is capable of safely and correctly capturing/generating computation graphs from arbitrary Python code without significant changes in the code. TorchInductor, on the other hand, takes the computation graph generated by TorchDynamo and converts it into optimized low-level kernels.

By using dynamo.optimize(), one can speed up real deep learning models up to 13% on NVIDIA A6000s. Furthermore, TorchDynamo outperforms existing solutions like TorchScript and FX Tracing by handling data-dependent control flows and non-PyTorch code without any significant changes to the code. It does so by breaking the computation graph whenever encountering an unsupported feature. It also allows the Python interpreter to handle the unsupported code and resume graph capturing afterward.

Citation Information

Mangla, P. “What’s Behind PyTorch 2.0? TorchDynamo and TorchInductor (primarily for developers),” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/u6xmk

@incollection{Mangla_2023_PT2TDTI4dev,
  author = {Puneet Mangla},
  title = {What's Behind {PyTorch} 2.0? {TorchDynamo} and {TorchInductor} (primarily for developers)},
  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/u6xmk},
}

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!

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!

Download the code!

Website

The post What’s Behind PyTorch 2.0? TorchDynamo and TorchInductor (primarily for developers) appeared first on PyImageSearch.

Table of Contents

Learning JAX in 2023: Part 3 — A Step-by-Step Guide to Training Your First Machine Learning Model with JAX

In this tutorial, you will learn how to train your first machine learning model with JAX.

This lesson is the last of a 3-part series on Learning JAX in 2023:

1. Learning JAX in 2023 | Part 1 — The Ultimate Guide to Accelerating Numerical Computation and Machine Learning
2. Learning JAX in 2023 | Part 2 — JAX’s Power Tools grad, jit, vmap, and pmap
3. Learning JAX in 2023: Part 3 — A Step-by-Step Guide to Training Your First Machine Learning Model with JAX (this tutorial)

To learn how to train your first machine learning model with JAX and PyTrees, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Learning JAX in 2023: Part 3 — A Step-by-Step Guide to Training Your First Machine Learning Model with JAX

We conclude our “Learning JAX in 2023” series with a hands-on tutorial. Throughout the series, we have covered the theoretical concepts of JAX, and in this post, we will apply those concepts to train a machine learning model. By the end of this tutorial, you will have a solid understanding of how to train a machine learning model using JAX and will be able to apply this knowledge to other ML problems.

The PyImageSearch team is dedicated to providing readers with the resources to better understand machine learning concepts through code. We will put everything we learned so far into gradually building a multilayer perceptron (MLP) with PyTrees. We have slowly broken down the tutorial into bits and pieces to present new concepts.

We hope this post will be a valuable resource as you continue learning and exploring the world of JAX.

Configuring Your Development Environment

To follow this guide, you need to have the JAX library installed on your system.

Luckily, JAX is pip-installable:

$ pip install jax

Having Problems Configuring Your Development Environment?

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 on your Windows, macOS, or Linux system now?

Then join 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!

🚝 Train a Simple Model with JAX

In this section, we will construct two basic models. The first model comprises a single weight and bias, whereas the second model has two weights and two biases.

But will they be able to learn? Of course, they will. We will first create and model a linear dataset using our linear model. Next, we will introduce nonlinearity using various activation functions.

The gradual increase in complexity will not only assist us in understanding the concepts better but also make it easier to implement code using JAX.

Before we build models using JAX, we first need to acquire a dataset. A simple one will suffice, as our goal is to demonstrate how to build models using JAX.

Build a Linear Dataset

Let us build a linear dataset. We will take the help of the well-known equation of a line to produce our data:

where is the dependent variable, is the independent variable, is the slope of the line, and is the -intercept of the line.

# Build the PRNG key
key = jax.random.PRNGKey(42)

xs = random.normal(key, shape=(128, 1))

# Hyperparameters for the linear function
m = 2.0
c = 4.0

# Build the linear function
ys = (m*xs) + c

We create a pseudorandom number generator (PRNG) key with the value of 42, which is used to seed the random number generation.

The normal function from the random module of JAX generates an array of 128 random values with a shape of (128, 1), where each value is sampled from a normal distribution.

Then we define the hyperparameters for the linear function.

The hyperparameters define the linear function, where xs is the array of random values, and ys is the resulting array of values computed by the linear function.

Figure 1 visualizes the linear function using matplotlib.

# Visualizing the function
plt.scatter(xs, ys)
plt.xlabel("xs")
plt.ylabel("ys")
plt.title("Linear Function")
plt.show()

Linear Model

Now that our linear function is ready to be modeled, let’s start creating our first machine learning model.

To combat a linear dataset, we build a linear regression model. The model will consist of a single weight and a single bias parameter that will be learned.

def linear_model(theta, x):
    """
    The linear model that will fit on the linear function.

    Args:
        - theta: Collection of weights and biases.
        - x: The input to the model.
    
    Returns:
        The prediction from the model.
    """
    weight, bias = theta
    pred = x * weight + bias
    return pred

The linear_model function is a simple linear model that takes in an input x and a collection of parameters called theta (which is a tuple of weight and bias) and returns a prediction based on the linear equation y = x * weight + bias.

def get_loss(theta, x, y):
    """
    The loss function. We use the L2 distance as our loss function.

    Args:
        - theta: Parameters of the model.
        - x: Input to the model.
        - y: The ground truth of the dataset.
    
    Returns:
        A scalar loss between the ground truth and the predicted value.
    """
    pred = linear_model(theta, x)
    loss = jnp.mean((y - pred)**2)
    return loss

The loss function is used to compute the loss of a linear model based on the following:

• input x
• ground truth y
• model’s parameters theta

It first calls the linear_model function, passing in theta and x as arguments to get the predicted output of the linear model. Then it calculates the loss between the predicted output and the ground truth y using the mean squared error (MSE) as a loss function.

@jit
def update_step(theta, x, y, lr):
    """
    A single update step for our model. We use the simple stochastic
    gradient update as our optimizer.

    Args:
        - theta: Parameters of the model.
        - x: Input to the model.
        - y: The ground truth of the dataset.
        - lr: Learning rate of the optimizer.
    
    Returns:
        - The current updated state of the model parameters.
        - The current loss
    """
    loss, gradient = value_and_grad(get_loss)(theta, x, y)
    updated_theta = theta - lr * gradient
    return updated_theta, loss

The update_step function performs a single update step (as the name suggests) on the linear model’s parameters theta, using the input x and the ground truth y, and the learning rate lr.

It uses the jax.value_and_grad function from the JAX library, which takes a function as an argument and returns the value of the function and its gradient with respect to its inputs.

Then we calculate the loss and its gradient by calling the get_loss function with theta, x, and y.

We update the parameters theta by subtracting the product of the learning rate lr and the gradient from the current value of theta.

# Initialize the parameters of the model
weight = 0.0
bias = 0.0
theta = jnp.array([weight, bias])

Here, we initialize the model. The parameters of the model are initialized with zeros. We encourage you to try other initialization techniques as well.

# Iterate and update the parameters
epochs = 10000
for iter in range(epochs):
    theta, loss = update_step(theta, xs, ys, 1e-4)
    if iter % 1000 == 0 and iter != 0:
        print(f"ITER {iter} | LOSS {loss:.4f}")

Finally, we train our linear model using stochastic gradient descent.

Each loop iteration calls the update_step function with the current parameters theta, the input xs, the ground truth ys, and a learning rate of 1e-4. The update_step function returns the updated parameters and the current loss of the model.

The updated parameters and the current loss are then assigned to the variables theta and loss, respectively. The final value of theta obtained from the loop will be the trained parameters of the model that can be used to make predictions on new data.

>>> ITER 1000 | LOSS 12.5518
>>> ITER 2000 | LOSS 8.7530
>>> ITER 3000 | LOSS 6.1060
>>> ITER 4000 | LOSS 4.2608
>>> ITER 5000 | LOSS 2.9740
>>> ITER 6000 | LOSS 2.0763
>>> ITER 7000 | LOSS 1.4499
>>> ITER 8000 | LOSS 1.0126
>>> ITER 9000 | LOSS 0.7074

Once we have our model trained, let’s see how it did on the training dataset. From Figure 2, we can see that the model fits the dataset quite well.

plt.scatter(xs, ys, label="true")
plt.scatter(xs, linear_model(theta, xs), label="pred")
plt.legend()
plt.show()

Build a Nonlinear Dataset

Let’s take this implementation up a notch. What do we do if our data comes from a nonlinear distribution? The linear regression model does not work anymore. We have to include some nonlinearity into the model now!

# Build the dataset
xs = random.normal(key, shape=(128, 1))

mu = 0.0
sigma = 0.1
noise = mu + sigma * random.normal(key, shape=(128, 1))

ys = jnp.sin(xs) + noise

We generate xs the same way as shown above.

We define the parameters for a Gaussian noise that we will add to the xs values to create the ground truth ys for the model.

We set the following:

• mu: the mean of the noise, to 0.0
• sigma: the standard deviation of the noise, to 0.1

The sine of each element in the array xs is added to the random noise. The result is stored in the variable ys. We will now use this dataset to train the model. We visualize the nonlinear function using matplotlib, as shown in Figure 3.

plt.scatter(xs, ys)
plt.xlabel("xs")
plt.ylabel("ys")
plt.title("Non Linear Function")
plt.show()

Nonlinear Model

To model a nonlinear function, we need to add a pinch of nonlinearity to our model. We will use the relu function as the intermediate nonlinearity.

def nonlinear_model(theta, x):
    """
    The non-linear model that will fit on the non-linear function.

    Args:
        - theta: Collection of weights and biases.
        - x: The input to the model.
    
    Returns:
        The prediction from the model.
    """
    weight1, weight2, bias1, bias2 = theta
    inter_output = (x * weight1) + bias1
    pred = (jax.nn.relu(inter_output) * weight2) + bias2
    return pred

The nonlinear_model function takes in as input x and a collection of parameters called theta (which are a tuple of weight1, weight2, bias1, and bias2). It returns a prediction based on a combination of a linear equation and a nonlinear activation function (ReLU in our case).

def get_loss(theta, x, y):
    """
    The loss function. We use the L1 distance as our loss function.

    Args:
        - theta: Parameters of the model.
        - x: Input to the model.
        - y: The ground truth of the dataset.
    
    Returns:
        A scalar loss between the ground truth and the predicted value.
    """
    pred = nonlinear_model(theta, x)
    loss = jnp.mean((y - pred)**2)
    return loss

The loss function is the same as defined earlier for our linear regression model.

@jit
def update_step(theta, x, y, lr):
    """
    A single update step for our model. We use the simple stochastic
    gradient update as our optimizer.

    Args:
        - theta: Parameters of the model.
        - x: Input to the model.
        - y: The ground truth of the dataset.
        - lr: Learning rate of the optimizer.
    
    Returns:
        - The current updated state of the model parameters.
        - The current loss
    """
    loss, gradient = value_and_grad(get_loss)(theta, x, y)
    updated_theta = theta - lr * gradient
    return updated_theta, loss

In the update step, we maintain the same approach as in the linear regression section by calculating the updated state of the model parameters and the current loss value.

# Initialize the parameters of the model
theta = jnp.array([1.0, 1.0, 1.0, 1.0])

# Iterate and update the parameters
epochs = 10000
for iter in range(epochs):
    theta, loss = update_step(theta, xs, ys, 1e-4)

    if iter % 1000 == 0 and iter != 0:
        print(f"ITER {iter} | LOSS {loss:.4f}")

Finally, we train our linear model utilizing the stochastic gradient descent algorithm. This step is similar to the one previously used with a linear model.

>>> ITER 1000 | LOSS 1.8689
>>> ITER 2000 | LOSS 1.0591
>>> ITER 3000 | LOSS 0.7026
>>> ITER 4000 | LOSS 0.5205
>>> ITER 5000 | LOSS 0.4171
>>> ITER 6000 | LOSS 0.3528
>>> ITER 7000 | LOSS 0.3093
>>> ITER 8000 | LOSS 0.2774
>>> ITER 9000 | LOSS 0.2523

In this step, we plot the ground truth and predicted values. As shown in Figure 4, the model converges but could be more effective. It cannot fully capture the nonlinearity present in the data.

plt.scatter(xs, ys, label="true")
plt.scatter(xs, nonlinear_model(theta, xs), label="pred")
plt.legend()
plt.show()

🌲 Training with PyTrees

But why are we suddenly interested in linear and nonlinear models? Aren’t they a little too basic? It’s always best to revisit fundamentals and remember why we are learning a new technique. Otherwise, we will be left with the knowledge of many techniques but without wisdom on when to use them.

A linear model is a model that is based on a linear equation, meaning that the output is a linear combination of the inputs. While linear models are simple and easy to understand, they have some limitations when modeling nonlinear relationships. Nonlinear models, on the other hand, can overcome these limitations by allowing for more complex relationships between inputs and outputs.

However, these (nonlinear) models require more parameters to model nonlinear functions, which makes training them more complicated. One of the main issues with having more parameters is that we need to pass each through weights and biases and apply gradients on each individually. This is a lengthy process, and where PyTrees from JAX comes in.

A PyTree is a container that can hold leaf elements, which are non-container objects such as arrays and other PyTrees. This allows for nested structures, where the container types do not need to match. This means a PyTree can be a list, tuple, dictionary of leaf elements, or other PyTrees.

In the context of a neural network, a PyTree can be used to represent the weights and biases of the network. Instead of initializing a separate list for each weight and bias, a PyTree can be used as a container for all the weights and biases. This makes it easier to handle many parameters, as the PyTree can be easily iterated over and modified.

JAX provides built-in support for PyTrees through the jax.tree_utils module, which includes functions such as jax.tree_map and jax.tree_leaves to manipulate and traverse the PyTree easily. The most common functions are also available as jax.tree_*.

def init_mlp_params(layer_widths):
    """
    Function to initialize the parameters of a Multilayer Perceptron.

    Args:
        - layer_widths: The list of layer widths.
    
    Returns:
        Randomly initialized parameters of the model.
    """
    key = random.PRNGKey(42)
    
    # params of the MLP will be a pytree
    params = []
    for number_in, number_out in zip(layer_widths[:-1], layer_widths[1:]):
        key, subkey = random.split(key)
        params.append(
            dict(
                weights=random.normal(subkey, shape=(number_in, number_out)) * jnp.sqrt(2/number_in),
                biases=jnp.ones(shape=(number_out,))
            )
        )
    return params

The init_mlp_params function takes in a single argument, layer_widths, which is a list of each layer’s widths in an MLP model.

We then initialize an empty list called params, which will be used to store the parameters of the MLP as a PyTree.

The function then enters a for loop, iterating over the layer_widths list, using the zip function to pair the input and output number of neurons for each layer.

On each iteration of the loop, the function splits the key into two parts, the key and subkey, using the random.split function (read how JAX handles randomness in Part 2 of this series). Then it appends a dictionary to the params list, which contains two keys, weights and biases, each with its corresponding value.

The value of the weights key is a random normal distribution generated using the random.normal function, with shape (number_in, number_out), and it is scaled by jnp.sqrt(2/number_in), which is a common initialization technique to ensure that the variance of the input is preserved.

The value of the biases key is an array of ones with shape (number_out,).

Finally, the function returns the params list, which contains the initialized parameters of the MLP model as a PyTree.

# Initialize the parameters of the MLP
params = init_mlp_params([1, 128, 128, 1])

# Custom function to get shape
get_shape = lambda x:x.shape
shape_pytree = jax.tree_map(get_shape, params)

for idx, layer in enumerate(shape_pytree):
    print(f"Layer {idx}")
    print(layer)

We use the init_mlp_params defined previously to initialize the parameters of a multilayer perceptron (MLP) model with layer widths of [1, 128, 128, 1]. The returned parameters are stored in the variable parameters.

We define a lambda function get_shape, which takes a single argument x and returns its shape using the shape attribute.

The function jax.tree_map(get_shape, params) is called. This function applies the get_shape function to every leaf element in the params PyTree. The result is a new PyTree (shape_pytree), where the shape of the corresponding leaf element in params PyTree replaces each leaf element.

Finally, the code enters a for loop, iterating over the shape_pytree list. On each iteration, the loop prints the index of the current layer using the enumerate function and the shape of the current layer by accessing the current value in the shape_pytree list.

This code snippet can be used to check the shapes of the parameters of each layer of the MLP model. The output will show the number of neurons in the input and output layers of each layer of the MLP.

>>> Layer 0
>>> {'biases': (128,), 'weights': (1, 128)}
>>> Layer 1
>>> {'biases': (128,), 'weights': (128, 128)}
>>> Layer 2
>>> {'biases': (1,), 'weights': (128, 1)}

Now that we have our model initialized, we define the forward step.

def forward(params, x):
    """
    The forward propagation step.

    Args:
        - params: The parameters of the model.
        - x: The input to the model.
    
    Returns:
        Prediction from the model.
    """
    # Get the hidden layers and the last layer separately.
    *hidden, last = params

    # Iterate over the hidden layers and forward propagate the
    # input through the layers.
    for layer in hidden:
        x = jax.nn.relu(x @ layer["weights"] + layer["biases"])

    # Get the prediction
    pred = x @ last["weights"] + last["biases"]
    return pred

The forward function performs the forward propagation step of a neural network. The function uses two arguments, params and x.

The params argument is the model’s parameters, which is expected to be a PyTree containing the weights and biases for each network layer. The x argument is the input to the model.

The function starts by using an unpacking operator * to extract the hidden layers and the last layer from the params PyTree. It then enters a for loop, iterating over the hidden layers.

On each iteration, the function performs a matrix multiplication of the input x with the weights of the current layer, and adds the biases of the current layer. Then it applies the ReLU (Rectified Linear Unit) activation function to the result using the jax.nn.relu function.

This process is repeated for all hidden layers.

After the for loop, the function performs a matrix multiplication of the final output from the hidden layers with the weights of the last layer, and adds the biases of the last layer. This produces the final prediction from the model. Finally, the function returns the prediction.

def get_loss(params, x, y):
    """
    The loss function. We use the L2 distance as our loss function.

    Args:
        - params: Parameters of the model.
        - x: Input to the model.
        - y: The ground truth of the dataset.
    
    Returns:
        A scalar loss between the ground truth and the predicted value.
    """
    pred = forward(params, x)
    loss = jnp.mean((pred - y) ** 2)
    return loss

You might notice that the loss function get_loss is very similar to what we defined earlier. The only difference is with the params. Previously, params were tuples of parameters. Now the params are PyTrees consisting of the parameters of each layer of the MLP.

@jax.jit
def update_step(params, x, y, lr):
    """
    A single update step for our model. We use the simple stochastic
    gradient update as our optimizer.

    Args:
        - params: Parameters of the model.
        - x: Input to the model.
        - y: The ground truth of the dataset.
        - lr: Learning rate of the optimizer.
    
    Returns:
        - The current updated state of the model parameters.
        - The current loss
    """
    loss, gradients = jax.value_and_grad(get_loss)(params, x, y)
    sgd = lambda param, gradient: param - lr * gradient
    updated_params = jax.tree_map(
        sgd, params, gradients
    )
    return updated_params, loss

To this point, we have defined our model, the forward propagation step, and the loss function. The missing piece of training a neural network is the update_step.

The update_step function takes in four arguments: params, x, y, and lr. The params argument represents the model’s parameters, x represents the input to the model, y represents the ground truth of the dataset, and lr represents the learning rate of the optimizer.

The main motivation of the update_step function is to grab the gradient of the loss wrt the parameters and update the current parameters. The update should, in turn, optimize the loss function.

Here, the jax.value_and_grad function lets us compute the loss and the gradient. Now we define a lambda function sgd, which takes params and gradients and returns the updated params according to the stochastic gradient descent algorithm.

We use the jax.tree_map method to apply the sgd function to each leaf node in the params PyTree.

The reason behind using a PyTree should become clear now. We did not have to change our code much from that of linear and nonlinear model training. The only necessary change was to use the jax.tree_utils functions to include PyTree in the training procedure.

Train the MLP

We have all of our tools in hand. Let’s slay this monster now.

# Build the dataset
xs = random.normal(key, shape=(128, 1))

mu = 0.0
sigma = 0.1
noise = mu + sigma * random.normal(key, shape=(128, 1))

ys = jnp.sin(xs) + noise

We reuse the same code to build our nonlinear dataset with sinusoids and a pinch of noise.

epochs = 1000
for iter in range(epochs):
    params, loss = update_step(params, xs, ys, 1e-5)

    if iter % 100 == 0 and iter != 0:
        print(f"ITER {iter} | LOSS {loss:.4f}")

Using the same code to iterate through the dataset and update the model’s parameters.

>>> ITER 100 | LOSS 0.9694
>>> ITER 200 | LOSS 0.4728
>>> ITER 300 | LOSS 0.2778
>>> ITER 400 | LOSS 0.1969
>>> ITER 500 | LOSS 0.1600
>>> ITER 600 | LOSS 0.1404
>>> ITER 700 | LOSS 0.1279
>>> ITER 800 | LOSS 0.1185
>>> ITER 900 | LOSS 0.1107

And finally, plot our results to show the nonlinear model’s convergence rate. Figure 5 shows that our multilayer perceptron model has modeled on the dataset quite well.

plt.scatter(xs, ys, label="true")
plt.scatter(xs, forward(params, xs), label="pred")
plt.legend()
plt.show()

As shown in the plot, these powerful tools can expertly fit even the most nonlinear data with just a few parameters. Embrace the challenge and see what amazing results you can achieve with your own model and dataset. And don’t forget to share your stunning convergence plots on Twitter, tagging us for a chance to win a surprise!

What's next? I recommend PyImageSearch University.

Course information:
75 total classes • 86 hours of on-demand code walkthrough videos • Last updated: April 2023
★★★★★ 4.84 (128 Ratings) • 15,900+ 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:

• ✓ 75 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 75 Certificates of Completion
• ✓ 86 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 500+ 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

Summary

So, we’ve covered a lot of ground in this blog post! We looked at how JAX can train linear and nonlinear regression models and even showed you how to use the PyTrees library to train a multilayer perceptron (MLP) model.

But that’s just the tip of the iceberg! We’re excited to inform you that this is just the first in a series of posts that will dive deeper into the world of JAX. We’ll cover more advanced topics and show you even more cool things you can do with this powerful library.

It’s no secret that JAX is becoming increasingly popular in machine learning, and the foundations we’ve covered here will be super useful in your future projects. So, watch out for our next posts, and let’s keep exploring the amazing world of JAX together!

Citation Information

A. R. Gosthipaty and R. Raha. “Learning JAX in 2023: Part 3 — A Step-by-Step Guide to Training Your First Machine Learning Model with JAX,” PyImageSearch, P. Chugh, S. Huot, K. Kidriavsteva, and A. Thanki, eds., 2023, https://pyimg.co/pidru

@incollection{ARG-RR_2023_Jax2023Pt3,
  author = {Aritra Roy Gosthipaty and Ritwik Raha},
  title = {Learning {JAX} in 2023: Part 3 — A Step-by-Step Guide to Training Your First Machine Learning Model with {JAX}},
  booktitle = {PyImageSearch},
  editor = {Puneet Chugh and Susan Huot and Kseniia Kidriavsteva and Abhishek Thanki},
  year = {2023},
  url = {https://pyimg.co/pidru},
}

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!

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!

Download the code!

Website

The post Learning JAX in 2023: Part 3 — A Step-by-Step Guide to Training Your First Machine Learning Model with JAX appeared first on PyImageSearch.

Table of Contents

Deploying a Custom Image Classifier on an OAK-D

In this tutorial, you will learn to deploy the image classification model on OAK that you trained in the TensorFlow framework in the previous tutorial in the OAK-101 series. With the help of the OpenVINO toolkit, you would convert and optimize the TensorFlow FP32 (32-bit floating point) model to the MyriadX blob file format expected by the Visual Processing Unit of the OAK device. The converted blob file would then run image classification inference on the OAK-D using the DepthAI API.

This lesson is the last in our 4-part series on OAK-101:

1. Introduction to OpenCV AI Kit (OAK)
2. OAK-D: Understanding and Running Neural Network Inference with DepthAI API
3. Training a Custom Image Classification Network for OAK-D
4. Deploying a Custom Image Classifier on an OAK-D (this tutorial)

To learn how to deploy and run an image classification network inference on OAK-D, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Deploying a Custom Image Classifier on an OAK-D

Introduction

As a deep learning engineer or practitioner, you may be working in a team building a product that requires you to train deep learning models on a specific data modality (e.g., computer vision) on a daily basis.

As an engineer, your work might include more than just running the deep learning models on a cluster equipped with high-end GPUs and achieving state-of-the-art results on the test data. However, optimizing and deploying those best models onto some edge device allows you to put your deep learning models to actual use in an industry where deployment on edge devices is mandatory and can be a cost-effective solution.

For example, for security, traffic management, manufacturing, healthcare, and agriculture applications, a coin-size edge device like OAK-D can be a great hardware to deploy your deep learning models.

In the previous tutorial of this series, we learned to train a custom image classification network for OAK-D using the TensorFlow framework. The image classification model we trained can classify one of the 15 vegetables (e.g., tomato, brinjal, and bottle gourd).

In today’s tutorial, we will take one step further and deploy the image classification model on OAK-D. First, we would learn the process of converting and optimizing the TensorFlow image classification model and then test the converted model on OAK-D with both images and the OAK device camera stream.

This is the last tutorial in our OAK-101 4-part series, and if you followed the series of tutorials from the beginning, we hope you have built strong foundations of the OpenCV AI Kit. We learned the OAK hardware and software stack from the ground level. With this tutorial, we would also learn to deploy an image classification application on the device.

Now, let’s start with today’s tutorial and learn about the deployment on OAK!

Configuring Your Development Environment

To follow this guide, you need to have depthai, opencv, and imutils installed on your system.

Luckily, all these libraries are pip-installable:

$ pip install depthai
$ pip install opencv-python
$ pip install imutils

Having Problems Configuring Your Development Environment?

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 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

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:

$ tree .
.
├── classify_camera.py
├── classify_image.py
├── openvino_model
│   ├── vegetable_classifier.bin
│   ├── vegetable_classifier.blob
│   ├── vegetable_classifier.mapping
│   └── vegetable_classifier.xml
├── pyimagesearch
│   ├── __init__.py
│   ├── config.py
│   └── utils.py
├── results
│   ├── pred_camera.mov
│   └── pred_images
│       ├── bean.jpg
│       ├── bitter_gourd.jpg
│       ├── brinjal.jpg
│       ├── cauliflower.jpg
│       ├── cucumber.jpg
│       ├── papaya.jpg
│       └── radish.jpg
└── test_data
    ├── bean.jpg
    ├── bitter_gourd.jpg
    ├── brinjal.jpg
    ├── cauliflower.jpg
    ├── cucumber.jpg
    ├── papaya.jpg
    └── radish.jpg

5 directories, 24 files

In the pyimagesearch directory, we have the following:

• config.py: The configuration file for the task
• utils.py: The utilities for running the image classification on OAK (e.g., creating pipelines and a few other helper functions)

In the core directory, we have the following:

• openvino_model: Houses the image classification trained model files converted to OpenVINO format (.blob) as required by OAK hardware
• test_data: It contains a few vegetable images from the test set, which the classify_image.py script will use
• classify_image.py: The inference script to leverage OAK’s neural accelerator for classifying images
• classify_camera.py: The inference script to run image classification with OAK’s camera

In the results directory, we have:

• pred_camera.mov: The prediction output file when inference is performed with OAK’s 4K color camera
• pred_images: Hosts the prediction results performed on the test_data images

Deploying the Model on OAK

In this section, we will broadly discuss the steps required to deploy your custom deep learning model to the OAK device.

The deep learning model could be in any format like PyTorch, TensorFlow, or Caffe, depending on the framework where the model was trained. Now the goal is to deploy the model on the OAK device and perform inference. However, the limitation is that the OAK device does not directly support any of these frameworks, so we need to convert the model to the MyriadX blob format supported by the OAK device.

The main reason why only a specific model format is required and the prominent deep learning frameworks don’t work directly on an OAK device is that the hardware has a visual processing unit based on Intel’s MyriadX processor, which requires the model in blob file format.

Figure 2 shows the steps required to convert the deep learning model from frameworks like PyTorch or TensorFlow to MyriadX blob file format and deploy them on the OAK device.

The process of converting the image classification model (in the TensorFlow framework) can be achieved in mainly 4-5 steps:

1. Installation and validation of the OpenVINO toolkit 2021.4.752 version
2. Converting the image classification TensorFlow model to IR (intermediate representation) using OpenVINO Model Optimizer. The IR consists of the model configuration in .xml format and the weights of the trained model in .bin format.
3. An optional step is to validate the intermediate representation by running inference on sample test images. It’s a sanity check to ensure the conversion from TensorFlow to IR works as expected.
4. Next, we convert the intermediate representation to MyriadX blob file format using the Model Compiler.
5. Finally, we can deploy the .blob file format on the OAK device and perform inference using the file and DepthAI API software stack.

Launch Jupyter Notebook on Google Colab by clicking on the link above

Most of the heavy lifting in the entire conversion process is done with the OpenVINO toolkit, which takes care of the most important intermediate step. The OpenVINO toolkit consists of a Model Optimizer and a Myriad Compiler.

The Model Optimizer can input any of the mentioned formats (first block Fig. 2) to produce an intermediate representation which can then be compiled to a .blob file format using the MyriadX compiler and finally deployed directly into the OAK device.

Configuring the Prerequisites

Before we start our implementation, let’s review our project’s configuration pipeline. Then, we will move on to the config.py script located in the pyimagesearch directory.

# import the necessary packages
import os
import glob

# define path to the model, test data directory and results
VEGETABLE_CLASSIFIER = os.path.join(
     "openvino_model","vegetable_classifier.blob"
)
TEST_DATA = glob.glob("test_data/*.jpg")
OUTPUT_IMAGES = os.path.join("results", "pred_images")
OUTPUT_VIDEO = os.path.join("results", "pred_camera.mov")

# define image height and width, and camera preview dimensions
IMG_DIM = (224, 224)
CAMERA_PREV_DIM = (480, 480)

# define the class label names list
LABELS = [
     "Bean", "Bitter_Gourd", "Bottle_Gourd", "Brinjal", "Broccoli",
     "Cabbage", "Capsicum", "Carrot", "Cauliflower", "Cucumber",
     "Papaya", "Potato", "Pumpkin", "Radish", "Tomato",
]

The config.py script sets up the necessary variables and paths for running the image classification model on images and camera streams to classify vegetables.

On Lines 2 and 3, we import the os and glob modules. Then, from Lines 6-11, we define the following:

• path to the vegetable classifier model
• test data directory
• output locations for images and videos

From Lines 14-22, we also define the dimensions for images and camera previews and a list of class label names to help decode class predictions to human-readable class names.

Defining the Utilities

Now that the configuration has been defined, we can determine the utilities for creating OAK pipelines and a few helper functions for resizing the input and normalizing the predictions. The utils.py script defines several functions:

• Helps create the pipeline for inference on OAK with images
• Pipeline for inference on OAK with color camera stream
• Define a softmax function to convert predictions into probabilities and a function to resize input and swap channel dimensions

Creating the Images Pipeline

# import the necessary packages
from pyimagesearch import config
from pathlib import Path
import numpy as np
import cv2
import depthai as dai

def create_pipeline_images():
    print("[INFO] initializing pipeline...")
    # initialize a depthai pipeline
    pipeline = dai.Pipeline()

    # configure inputs for depthai pipeline
    classifierIN = pipeline.createXLinkIn()
    classifierIN.setStreamName("classifier_in")

On Lines 2-6, we import the necessary packages:

• config from the pyimagesearch module
• Path from pathlib, which would help in loading the model from disk
• numpy for softmax computation
• cv2 for resizing and other image-related utilities
• depthai library

We define the function create_pipeline_images() on Line 8. Then, a depthai pipeline is initialized on the host, which helps define the nodes, the flow of data, and communication between the nodes (Line 11).

On Lines 14 and 15, we configure the pipeline’s inputs by creating an XLinkIn object and setting the stream name to classifier_in. This stream name is used to specify the input source for the pipeline. In this case, the image classifier model will classify objects in the images.

In short, the XLinkIn, if you recall from the 2nd tutorial of this series, will help send image data from the host to the OAK device, which then would be fed to the classifier for prediction.

    # configure vegetable classifier model and set its input
    print("[INFO] initializing vegetable classifier network...")
    classifierNN = pipeline.create(dai.node.NeuralNetwork)
    classifierNN.setBlobPath(
        str(Path(config.VEGETABLE_CLASSIFIER).resolve().absolute())
    )
    classifierIN.out.link(classifierNN.input)

    # configure outputs for depthai pipeline
    classifierNNOut = pipeline.createXLinkOut()
    classifierNNOut.setStreamName("classifier_nn")
    classifierNN.out.link(classifierNNOut.input)

    # return the pipeline
    return pipeline

On Lines 19-22, we create a NeuralNetwork node and set the blob path to the path of the classifier model. We then configured the vegetable classifier model, and this step is similar to how we load weights to a neural network model in TensorFlow or PyTorch.

On Line 23, the classifierNN object is linked to the classifierIN object, which was created earlier to define the input stream name.

Then from Lines 26-31,

• We configure the outputs for the depthai pipeline by creating an XLinkOut object.
• The output stream name is set to classifier_nn. This stream name is used to specify the output source for the pipeline, in this case, the output of the neural network classifier.
• The classifierNN object is linked to the classifierNNOut object, which was created to define the output stream name.
• Finally, the function returns the pipeline object configured with the classifier model and input/output streams to the calling function.

Creating the Camera Pipeline

def create_pipeline_camera():
    print("[INFO] initializing pipeline...")
    # initialize a depthai pipeline
    pipeline = dai.Pipeline()

    # configure vegetable classifier model and set its input
    print("[INFO] initializing vegetable classifier network...")
    classifierNN = pipeline.create(dai.node.NeuralNetwork)
    classifierNN.setBlobPath(
        str(Path(config.VEGETABLE_CLASSIFIER).resolve().absolute())
    )

Next, we define the create_pipeline_camera() that initializes a depthai pipeline on Line 36. We then define the NeuralNetwork node and set the blob path to the path of the classifier model (Lines 40-43).

This function is similar to the create_pipeline_images() function, but here we do not define the input stream or the XLinkIn node since we would leverage the OAK module’s in-built camera as an input to the image classifier model.

    # create and configure the color camera properties
    print("[INFO] Creating Color Camera...")
    cam_rgb = pipeline.create(dai.node.ColorCamera)
    cam_rgb.setPreviewSize(config.CAMERA_PREV_DIM)
    cam_rgb.setInterleaved(False)
    cam_rgb.setResolution(
        dai.ColorCameraProperties.SensorResolution.THE_1080_P
    )
    cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)
    cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.RGB)

Now we create and configure the color camera properties by creating a ColorCamera node and setting the preview size, interleaved status, resolution, board socket, and color order.

From Lines 47-54,

• The cam_rgb object is created from the pipeline object, allowing it to be connected to the pipeline. This would act as an input to the classifier model.
• Then the setPreviewSize method sets the resolution of the preview window to the CAMERA_PREV_DIM variable defined earlier in config.
• Next, the setInterleaved(False) method disables the interleaved mode. This means the depth and color frames will not be interleaved together.
• The setResolution method sets the camera resolution to 1080p. It can also be set to different resolutions.
• The setBoardSocket method connects the camera to the RGB port of the camera board.
• Finally, the setColorOrder method sets the order of the color channels to RGB.

    # create XLinkOut node for displaying frames
    cam_xout = pipeline.create(dai.node.XLinkOut)
    # set stream name as rgb
    cam_xout.setStreamName("rgb")
    # link the camera preview to XLinkOut node input
    cam_rgb.preview.link(cam_xout.input)

From Lines 56-61,

• An XLinkOut node is created (i.e., cam_xout variable), which would help display frames on the host side.
• The stream name is set to rgb. This name would be passed as a parameter to the OutputQueue for fetching the color frames.
• Then the camera preview is linked to the XLinkOut node (cam_xout) input by linking the output of the cam_rgb object’s preview property to the input of the cam_xout object. This allows the frames captured by the color camera to be output to the specified stream, named rgb. This stream can then be used for displaying or saving video frames.

    # resize the camera frames to dimensions expected by neural network
    print("[INFO] Creating ImageManip node...")
    manip = pipeline.create(dai.node.ImageManip)
    manip.initialConfig.setResize(config.IMG_DIM)
    # link the camera preview to the manipulation node for resize
    cam_rgb.preview.link(manip.inputImage)
    # link the output of resized frame to input of neural network
    manip.out.link(classifierNN.input)

Then, on Lines 65 and 66, we create the ImageManip node, which is used for image manipulation and the initial configuration of the manipulation node is set with the dimensions defined in the config.IMG_DIM variable.

It then links the camera preview to the manipulation node input by linking the output of the cam_rgb object’s preview property to the input of the manip object on Line 68.

On Line 70, the output of the image manipulation node is then linked to the input of the classifierNN object, which is the neural network classifier model. This allows the camera frames to be passed through the image manipulation node to be resized to the required dimensions before being passed to the classifier model.

    # configure outputs for depthai pipeline
    xout_nn = pipeline.create(dai.node.XLinkOut)
    xout_nn.setStreamName("nn")
    classifierNN.out.link(xout_nn.input)

    # return the pipeline
    return pipeline

On Lines 73-75, we link the classifierNN (image classifier) output to an XLinkOut node, allowing us to display or save the image classification predictions. We would extract the predictions using the nn stream name.

Finally, on Line 78, the function returns the pipeline object, which has been configured with the classifier model, color camera, image manipulation node, and input/output streams.

def softmax(x):
    # compute softmax values for each set of scores in x.
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum(axis=0)


def to_planar(arr: np.ndarray, shape: tuple) -> np.ndarray:
    # resize the image array and modify the channel dimensions
    resized = cv2.resize(arr, shape)
    return resized.transpose(2, 0, 1)

From Lines 80-83, we define the softmax() function, which calculates the softmax values for a given set of scores in x. The softmax function is a commonly used activation function in neural networks, particularly in the output layer, to return the probability of each class.

The to_planar() function takes in two arguments: an array and a tuple representing the shape (Lines 86-89).

• It resizes the array to the given shape using the cv2 library
• Then modifies the channel dimensions by transposing the array so that
  • the first dimension represents the channels
  • the second dimension represents the rows
  • the third dimension represents the columns

In short, the to_planar() function helps reshape image data before passing it to the neural network.

Classify Images

With the configurations and utilities implemented, we can finally get into the code walkthrough of classifying images on OAK-D.

# import the necessary packages
from pyimagesearch import config
from pyimagesearch import utils
import os
import numpy as np
import cv2
import depthai as dai

# initialize a depthai images pipeline
print("[INFO] initializing a depthai image pipeline...")
pipeline = utils.create_pipeline_images()

We start by importing the necessary packages, including the config and utils modules from pyimagesearch, and the os, numpy, cv2, and depthai modules on Lines 2-7.

The magic happens on Line 11, where we initialize the depthai images pipeline by calling the create_pipeline_images() function from the utils module.

The create_pipeline_images() function

• Creates and configures a pipeline for running the vegetable classifier model on images.
• The pipeline object returned by the function is assigned to the variable pipeline.
• It would create a pipeline that is ready to process images and perform inference using the depthai library and the vegetable classifier model.

# pipeline defined, now the device is assigned and pipeline is started
with dai.Device(pipeline) as device:
    # define the queues that will be used in order to communicate with
    # depthai and then send our input image for predictions
    classifierIN = device.getInputQueue("classifier_in")
    classifierNN = device.getOutputQueue("classifier_nn")

Now that the pipeline is created, a context manager is created using the with statement and the Device class from depthai on Line 14.

It assigns the pipeline object created earlier to the Device class. This will establish a connection between the pipeline and the device, which is necessary for the pipeline to run.

Then on Lines 17 and 18, two queues are defined, classifierIN and classifierNN, that will be used to communicate with the device and send input images for predictions.

The classifierIN variable is assigned the input queue for the classifier_in stream, and the classifierNN variable is assigned the output queue for the classifier_nn stream, defined in the create_pipeline_images() function. These queues will send images to the pipeline for image classification and receive the predictions from the pipeline.

    print("[INFO] loading image from disk...")
    for img_path in config.TEST_DATA:
        # load the input image and then resize it
        image = cv2.imread(img_path)
        image_copy = image.copy()
        nn_data = dai.NNData()
        nn_data.setLayer(
            "input",
            utils.to_planar(image_copy, config.IMG_DIM)
        )
        classifierIN.send(nn_data)

On Line 21, we start to iterate over the list of image paths stored in the config.TEST_DATA. On Lines 23 and 24, we read the image using cv2 from the disk and create a copy of the original image.

Then on Lines 25-29,

• An nn_data object is created and set to the neural network’s input layer using the setLayer() method.
• The utils.to_planar() function is called to reshape the image to the dimensions specified in the IMG_DIM variable of the config module.

Finally, on Line 30, the send() method of the classifierIN queue is called to send the image to the pipeline for image classification. The pipeline uses the classifier model to classify objects in the image and returns the predictions through the classifierNN queue, which can be used for further processing or display purposes.

        print("[INFO] fetching neural network output for {}".
            format(img_path.split('/')[1]))
        in_nn = classifierNN.get()

        # apply softmax on predictions and
        # fetch class label and confidence score
        if in_nn is not None:
            data = utils.softmax(in_nn.getFirstLayerFp16())
            result_conf = np.max(data)
            if result_conf > 0.5:
                result = {
                    "name": config.LABELS[np.argmax(data)],
                    "conf": round(100 * result_conf, 2)
                }
            else:
                result = None

On Line 34, the neural network (stored in the variable classifierNN) is used to classify the image, and the output is stored in in_nn. In short, we get the neural network output from the queue.

From Lines 38-47,

• If the output is not None (i.e., the classification was successful), it applies the softmax function on the output to get the probability scores for each class.
• Next, the function extracts the class label by getting the index of the maximum probability and then using it to look up the corresponding label in the config.LABELS list.
• It also extracts the confidence score by getting the maximum probability value itself.
• If the maximum probability is greater than 0.5, it stores the class label and confidence score in a dictionary and assigns it to the result variable. Otherwise, it assigns result to None.

        # if the prediction is available,
        # annotate frame with prediction and show the frame
        if result is not None:
            cv2.putText(
                image,
                "{}".format(result["name"]),
                (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.9,
                (255, 255, 255),
                thickness=2,
                lineType=cv2.LINE_AA
            )
            cv2.putText(
                image,
                "Conf: {}%".format(result["conf"]),
                (10, 60),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7,
                (255, 255, 255),
                thickness=2,
                lineType=cv2.LINE_AA
            )

            cv2.imwrite(
                config.OUTPUT_IMAGES +"/"+img_path.split('/')[1],
                image
            )

Now that we have the image classification prediction for an image, we are all set to display the class label (name) and confidence score (conf) of the image classification on the image itself using the cv2 module.

From Lines 51-76,

• If the classification result is not None (i.e., the classification was successful), we use the cv2.putText() function to write the class label and confidence score on the image.
• The first cv2.putText() call writes the class label at coordinates (10, 30) on the image, and the second one writes the confidence score at (10, 60).
• The cv2.imwrite() function saves the image with the annotation to the specified location.

Results

Figure 3 shows the image classification results on a few sample test images when fed to the OAK device. From the below figure, we can see that the converted (.blob) and optimized image classification model does pretty well in classifying the test images. However, it misclassified papaya as pumpkin with a confidence score of 52.49%.

Classify Camera Stream

Great! So we just learned how to run neural network inference on images with an OAK device. Let’s now dive one step further and use the OAK’s color camera to classify the frames, which in our opinion, is where you put your OAK module to real use, catering to a wide variety of applications discussed in the 1st blog post of this series.

# import the necessary packages
from pyimagesearch import config
from pyimagesearch import utils
from imutils.video import FPS
import numpy as np
import cv2
import depthai as dai


# initialize a depthai camera pipeline
print("[INFO] initializing a depthai camera pipeline...")
pipeline = utils.create_pipeline_camera()

frame = None
fps = FPS().start()

fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(
    config.OUTPUT_VIDEO,
    fourcc,
    20.0,
    config.CAMERA_PREV_DIM
)

On Lines 2-7, we import the necessary packages like config, utils, OpenCV, NumPy, and the depthai library. We also import the FPS module, which would help us compute how long the image classification takes to perform inference on each frame.

On Line 12, we call utils.create_pipeline_camera(), which initializes the depthai pipeline for capturing video frames from the OAK camera and performing image classification.

Then, on Line 15, a new instance of the FPS class is defined, which is used to keep track of the frames per second (FPS) at which the pipeline is running or, let’s say, the time inference takes on an OAK device.

Since we are also interested in saving the prediction of the camera stream as a video, on Line 17, we specify a codec (fourcc) to be used for writing the video and creating a VideoWriter object. The codec being used is XVID.

From Lines 18-23, we define the video writer object, which takes several of the following parameters:

• the output video path along with the filename
• next, we pass the codec defined on Line 17
• the frame rate at which the video will be saved is set to 20.0
• the dimensions of the video frames via the config.CAMERA_PREV_DIM parameter

# Pipeline defined, now the device is assigned and pipeline is started
with dai.Device(pipeline) as device:

    # Output queues will be used to get the rgb frames and
    # nn data from the outputs defined above
    q_rgb = device.getOutputQueue(name="rgb", maxSize=1, blocking=False)
    q_nn = device.getOutputQueue(name="nn", maxSize=1, blocking=False)

Similar to the classifying images section, a context manager is created using the with statement and the Device class from depthai on Line 26.

Inside the block, on Lines 30 and 31,

• We create two output queues, one for the RGB frames and one for the neural network data.
• The first queue, q_rgb, is created by calling the device.getOutputQueue() function, passing in the name of the output queue (rgb), the maximum size of the queue (1), and the blocking parameter (False).
• The second queue, q_nn, is also created in the same way as the first queue, but with the name nn instead of rgb.

    # `get_frame()` fetches frame from OAK,
    # resizes and returns the frame to the calling function
    def get_frame():
        in_rgb = q_rgb.get()
        new_frame = np.array(in_rgb.getData())\
            .reshape((3, in_rgb.getHeight(), in_rgb.getWidth()))\
            .transpose(1, 2, 0).astype(np.uint8)
        new_frame = cv2.cvtColor(new_frame, cv2.COLOR_BGR2RGB)
        return True, np.ascontiguousarray(new_frame)

    result = None

Next, we define a function named get_frame() which

• fetches a frame from the OAK device
• resizes it
• returns the frame to the calling function

On Line 36, we fetch a frame from the q_rgb queue by calling q_rgb.get(). Then, on Lines 37-39,

• Reshape the frame data into a 3-dimensional array, with dimensions (height, width, 3) to match the expected dimensions of an image.
• Transpose the array’s dimensions to (height, width, 3).
• Change the data type of the array to uint8, which is used to represent 8-bit unsigned integers in NumPy.

On Line 40, the color space of the frame is converted from BGR to RGB using the cv2.cvtColor() function.

Finally, the function returns a tuple containing a Boolean value (True) and the processed frame as a contiguous array on Line 41.

    while True:
        read_correctly, frame = get_frame()

        if not read_correctly:
            break

        # fetch neural network prediction
        in_nn = q_nn.tryGet()

With the frame and neural network data queues defined and the frame postprocessing helper function in place, we start the while loop on Line 45.

In general, the above code runs a loop that captures video frames from the OAK device, processes them, and fetches neural network predictions from the q_nn queue.

On Line 46, we call the get_frame() function to fetch a new frame. The function returns a tuple containing a Boolean value indicating whether the frame was read correctly and the frame itself.

On Lines 48 and 49, we check if the Boolean value is false, which would indicate that the frame was not read correctly. If that’s the case, the loop is broken.

If the Boolean value is true, the code fetches a neural network prediction from the q_nn queue by calling the q_nn.tryGet() function (Line 52).

        # apply softmax on predictions and
        # fetch class label and confidence score
        if in_nn is not None:
            data = utils.softmax(in_nn.getFirstLayerFp16())
            result_conf = np.max(data)
            if result_conf > 0.5:
                result = {
                    "name": config.LABELS[np.argmax(data)],
                    "conf": round(100 * result_conf, 2)
                }
            else:
                result = None

Now that we have the neural network prediction, we apply a softmax function on the output of the neural network in_nn and then extract the class label and confidence score from the resulting data.

On Lines 56-65,

• We check if the neural network output is not None.
• If not None, we apply the softmax function on the output to get the probability scores for each class by calling utils.softmax(in_nn.getFirstLayerFp16()).
• We extract the confidence score by getting the maximum probability value with np.max(data).
• The function extracts the class label by getting the index of the maximum probability and then using it to look up the corresponding label in the config.LABELS list.
• If the maximum probability is greater than 0.5, we store the class label and confidence score in a dictionary and assign it to the result variable. Otherwise, assign result to None.

        # if the prediction is available,
        # annotate frame with prediction and display it
        if result is not None:
            cv2.putText(
                frame,
                "{}".format(result["name"]),
                (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.9,
                (255, 255, 255),
                thickness=2,
                lineType=cv2.LINE_AA
            )
            cv2.putText(
                frame,
                "Conf: {}%".format(result["conf"]),
                (10, 60),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7,
                (255, 255, 255),
                thickness=2,
                lineType=cv2.LINE_AA
            )

On Lines 69-89, the OpenCV library puts text on the frame. The text added to the frame is the class label of the prediction and the confidence score of the prediction made on the frame.

        # update the FPS counter
        fps.update()

        # display the output
        cv2.imshow("rgb", frame)

        # break out of the loop if `q` key is pressed
        if cv2.waitKey(1) == ord('q'):
            break

        # write the annotated frame to the file
        out.write(frame)

# stop the timer and display FPS information
fps.stop()
print("[INFO] elapsed time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))

# do a bit of cleanup
out.release()
cv2.destroyAllWindows()

From Lines 92-99,

• The frames per second (FPS) counter is updated using the fps.update() method.
• The output image is displayed on the screen using cv2.imshow("rgb", frame).
• If the user presses the q key, the while loop breaks, and the program continues to execute the next line of code outside the loop.

The output image is also written to a file using out.write(frame) on Line 102. After the loop is broken, the fps.stop() method is called to stop the timer on Line 105.

The elapsed time and approximate FPS are printed to the console on Lines 106 and 107. The out.release() method is then used to release the file, and cv2.destroyAllWindows() is used to close all open windows, effectively ending the program (Lines 110 and 111).

Results

Below are the inference results on the video stream, and the predictions seem good. Initially, the network misclassified capsicum as brinjal. The image classification network achieved 30 FPS real-time speed on the OAK device.

What's next? I recommend PyImageSearch University.

Course information:
74 total classes • 84 hours of on-demand code walkthrough videos • Last updated: March 2023
★★★★★ 4.84 (128 Ratings) • 15,800+ 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:

• ✓ 74 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 74 Certificates of Completion
• ✓ 84 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 500+ 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

Summary

With this, we have come to the end of the OAK-101 series. We hope you enjoyed this series on OpenCV AI Kit as much as we did! But hold tight, as we will soon come back with OAK-102, where we would like to cover a few more advanced computer vision applications using OAK devices.

For now, let’s quickly summarize what we learned today.

In last week’s tutorial, we trained an image classification model on a vegetable image dataset in the TensorFlow framework. In this tutorial, we optimized that trained TensorFlow classification model for the OAK device.

We first learned the model deployment process for OAK; specifically, we discussed the process of taking the deep learning trained model in various frameworks (e.g., TensorFlow, PyTorch, Caffe) to the intermediate representation to the .blob format supported by the OAK device. To achieve this, we discussed the role of the OpenVINO toolkit.

Then, we covered the conversion and optimization process of the trained image classification TensorFlow model to the .blob format. Finally, as a sanity check, we tested the model in Google Colab with some sample vegetable test images before feeding the OAK with the optimized model.

With our image classification model converted to the .blob format, we used the DepthAI API to run inference on OAK with images and camera video streams.

Finally, we examined the classification results of both images and video streams. Then, we computed the end-to-end runtime performance for the video inference pipeline, and the OAK device achieved real-time speed (i.e., 30 FPS).

Citation Information

Sharma, A. “Deploying a Custom Image Classifier on an OAK-D,” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/tnx0u

@incollection{Sharma_2023_Deploy-CIC-OAK-D,
  author = {Aditya Sharma},
  title = {Deploying a Custom Image Classifier on an {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/tnx0u},
}

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!

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!

Download the code!

Website

The post Deploying a Custom Image Classifier on an OAK-D appeared first on PyImageSearch.

If you want a gentle introduction to machine learning for computer vision, you’re in the right spot. Here at PyImageSearch we’ve been helping people just like you master deep learning for computer vision.

Also, you might want to check out our computer vision for deep learning program before you go.

Machine Learning for Computer Vision

Spotlight on the Key Differences: Machine Learning vs. Computer Vision Simply put, computer vision is a specialized branch of machine learning that empowers computers to acquire a profound understanding by processing videos and digital images.

Unraveling the Connection Between Computer Vision and Machine Learning Curious about the relationship between computer vision and machine learning? Join us as we unveil the fascinating details of what they are, how they function, and how they intertwine with one another.

What is Machine Learning

A Closer Look at Machine Learning Machine learning (ML) revolves around creating algorithms and related systems that learn behavior strategies in specific environments using instructions and training data sets.

As a subfield of artificial intelligence, machine learning sidesteps some of the more profound and philosophical AI questions, focusing instead on learning and training methods that result in efficient machines for various applications. This field encompasses statistical models, algorithms, and learning techniques that find use in diverse industries such as manufacturing, retail, logistics, food production, and construction.

Machine learning encompasses several strategies that teach algorithms to recognize patterns in data, guiding informed actions in similar settings. These strategies include:

1. Supervised Learning: In this approach, data scientists provide ML systems with training data sets containing inputs and corresponding desired outputs. The machine learning system then grasps the intended outcomes of specific actions and formulates optimal strategies

to achieve those results.

2. Unsupervised Learning: As the name suggests, unsupervised learning involves unstructured data sets without predefined, ideal outputs. The machine learning system is responsible for identifying patterns in the data sets and devising strategies for behavior.
3. Reinforcement Learning: Primarily used for training independent machine agents within a system, reinforcement learning employs cumulative reward models to teach agents how to act in different environments. This ML application is utilized in various industries, with substantial research conducted in the realm of online multiplayer games.
4. Deep Learning and Neural Networks: Traditional machine learning and AI systems relied on linear or iterative learning methods. However, since the 1980s, researchers have developed “neural network” architectures using node-cluster structures and weighted decision-making strategies. This approach enables machine learning systems to simplify complex problems, with the outcomes of simpler problems combining to form a comprehensive solution for larger ones.

Deep learning takes this concept a step further by designing layered neural networks where levels of solution-based networks essentially create an emergent problem-solving engine. For instance, a deep learning system could feature layers where basic pattern-recognition processes collaborate to execute intricate tasks like facial recognition in images.

The primary focus of these and other machine learning strategies is to educate ML systems, simulate training environments, and utilize machine learning to drive comprehensive artificial intelligence and autonomous systems.

What is Computer Vision

Discover the World of Computer Vision: Making AI and ML More Insightful

Discover the World of Computer Vision: Making AI and ML More Insightful Creating artificial intelligence or machine learning systems can be a complex task. To streamline the process, data scientists often focus on specific applications, such as managing a supply chain or manufacturing process, rather than expecting a system to function across multiple environments.

However, many researchers aim to expand the capabilities of these systems to cover multiple domains. Since the early days of artificial intelligence, the idea of a versatile, general AI that could power various systems has captivated scientists. This notion drove research in the 1970s and early 1980s, exploring AI applications in image processing, language recognition, and robotics.

One significant challenge in fields like robotics is enabling AI systems to process, interpret, and react to visual information. Consequently, there was a push to develop advanced visual and acoustic sensors, environmental navigation systems, and mobility features.

The resurgence of artificial intelligence and machine learning in recent years, fueled by high-performance computing, big data, and cloud systems, has increased the demand for hardware capable of collecting visual data for machine learning applications.

Introducing computer vision – a specialized branch of machine learning and artificial intelligence that extracts meaningful insights from digital images and videos, allowing for informed decision-making.

Like most ML systems, computer vision relies on large amounts of data to train algorithms in interpreting this information.

Computer vision typically employs two main technologies:

1. Deep Learning: As previously mentioned, deep learning supports intricate problem-solving.

More importantly, deep learning using neural networks can essentially train machine “brains” to process visual data and retain knowledge of patterns, strategies, and environmental changes over time.

2. Convolutional Neural Networks (CNNs): These networks process visual information like images by breaking them down into pixels. They use “convolutions” (the operation of deriving a mathematical function from two other functions) to make predictions based on that data.

In essence, computer vision combines CNNs and deep learning to conduct high-speed, high-volume unsupervised learning on visual information. This process trains machine learning systems to interpret data in a way that resembles how the human eye functions, making AI and ML more insightful and versatile across various applications.

Relationship between Machine Learning and Computer Vision

Exploring the Intricate Relationship Between Machine Learning and Computer Vision

Machine learning and computer vision are two captivating domains within the broader field of artificial intelligence. While they share many similarities, they each have unique roles and applications. In this section, we’ll delve into the relationship between machine learning and computer vision, revealing how they are intertwined and complement each other.

Machine Learning: Laying the Foundation

Machine learning (ML) is a subset of artificial intelligence that focuses on developing algorithms and systems capable of learning from and making predictions or decisions based on data. It involves various techniques, including supervised learning, unsupervised learning, reinforcement learning, and deep learning.

These techniques enable machine learning algorithms to identify patterns, adapt to new data, and improve their performance over time, providing a foundation for many AI applications, such as natural language processing, recommendation systems, and fraud detection.

Computer Vision: Building on the ML Foundation

Computer vision, on the other hand, is a specialized field within machine learning that deals with teaching computers to interpret and understand visual information from the world, such as images and videos. By leveraging machine learning techniques, computer vision systems can process vast amounts of visual data, recognize patterns, and make informed decisions based on their interpretations.

Convolutional Neural Networks (CNNs) and deep learning are essential tools in computer vision, enabling these systems to break down complex visual data into simpler components and analyze them more effectively. Some common applications of computer vision include facial recognition, object detection, and autonomous vehicle navigation.

The Synergy Between Machine Learning and Computer Vision

The relationship between machine learning and computer vision is akin to that of a tree and its branches. Machine learning lays the groundwork, forming the trunk, while computer vision extends from it as a specialized branch. Both domains share the same roots in artificial intelligence and are closely intertwined in their goals and methods.

The success of computer vision depends heavily on the foundation laid by machine learning. ML techniques enable computer vision systems to extract meaningful insights from visual data, making them smarter and more adaptable. In return, computer vision expands the range of applications and capabilities of machine learning, pushing the boundaries of what AI can achieve.

This interplay between machine learning and computer vision has led to remarkable advancements in both fields. For example, deep learning-based computer vision algorithms have vastly improved object recognition, making it possible for AI-powered systems to quickly and accurately identify objects in images or videos. In turn, these innovations have fueled breakthroughs in areas such as robotics, self-driving cars, and medical imaging.

Conclusion: A Powerful Partnership for the Future of AI

The relationship between machine learning and computer vision is a testament to the power of collaboration within the AI community. By building upon each other’s strengths and pushing the limits of what is possible, these two fields have paved the way for remarkable innovations that continue to transform the world around us.

As researchers and developers continue to explore new techniques and applications, the partnership between machine learning and computer vision will undoubtedly strengthen, resulting in even more groundbreaking advancements in the realm of artificial intelligence. The future is bright for both fields, and their synergy promises to unlock new possibilities and revolutionize the way we interact with technology.

Tasks involved in Computer Vision

Unveiling the Power of Computer Vision: Revolutionizing Industries Through Visual Intelligence

Introduction: The Rise of Computer Vision Applications

In recent years, computer vision has emerged as a groundbreaking technology, transforming various industries and changing the way we interact with the world. By teaching computers to interpret and understand visual information from images and videos, computer vision applications have made it possible to automate and enhance numerous tasks that were once heavily reliant on human vision.

In this blog post, we will explore some of the most notable tasks accomplished using computer vision applications, showcasing how this technology is revolutionizing different sectors and unlocking new possibilities.

1. Image and Object Recognition

One of the primary tasks enabled by computer vision is image and object recognition. This capability allows computers to identify and classify objects within images or videos, making it possible to create AI-powered systems that can quickly and accurately recognize specific items. Some common applications of image and object recognition include:

• Security and surveillance: Facial recognition systems can identify individuals in real-time, enhancing security measures in public spaces, airports, and other sensitive areas.
• Retail: Computer vision technology can enable cashierless stores, where customers can simply pick up items and walk out, with the system automatically recognizing and billing them for the products.

2. Autonomous Vehicles

Computer vision for transportation plays a crucial role in the development of autonomous vehicles. By processing visual data from cameras and other sensors, these systems can recognize and track objects in real-time, such as pedestrians, other vehicles, and traffic signs. This information enables the vehicle to make informed decisions, navigate complex environments, and avoid potential hazards.

3. Medical Imaging and Diagnosis

In the medical field, computer vision applications have made significant strides in improving the accuracy and efficiency of medical imaging and diagnosis. By analyzing images from X-rays, MRIs, and other scans, AI-powered computer vision systems can detect abnormalities, identify potential health issues, and even predict the likelihood of certain diseases. This technology is particularly beneficial in areas such as:

• Cancer detection: Computer vision can assist radiologists in identifying tumors and cancerous cells in medical images, reducing the risk of misdiagnosis and improving patient outcomes.
• Ophthalmology: AI-driven computer vision systems can analyze retinal images to detect early signs of diabetic retinopathy and other eye conditions, helping doctors provide timely and accurate treatment.

image credits

4. Agriculture and Precision Farming

Computer vision applications are also making their way into the agriculture industry, revolutionizing traditional farming practices. By analyzing images captured by drones, satellites, or ground-based cameras, these systems can provide valuable insights into crop health, soil conditions, and pest infestations. Some key applications in agriculture include:

• Crop monitoring: Computer vision can identify and track various crop health indicators, enabling farmers to make informed decisions about irrigation, fertilization, and pest control.
• Harvesting automation: AI-powered harvesting robots can use computer vision to recognize ripe fruits and vegetables, picking them with precision and efficiency.

5. Augmented Reality (AR) and Virtual Reality (VR)

Computer vision plays a critical role in the development of AR and VR experiences. By processing and interpreting visual data in real-time, these systems can create immersive, interactive environments that seamlessly blend digital

VR include:

• Gaming: AR and VR games leverage computer vision to track the user’s movements and gestures, providing a more immersive and engaging gaming experience.
• Education and training: AR and VR applications can enhance learning experiences by overlaying digital information on physical objects, enabling users to explore and interact with complex subjects in a more intuitive way.

You might also like to learn about computer vision for logistics.

Conclusion Embracing the Future of Machine Learning for Computer Vision

As we’ve seen, computer vision applications are reshaping a wide range of industries and tasks, demonstrating the incredible potential of this technology. The key driver behind these advancements is the synergy between computer vision and machine learning. By harnessing the power of machine learning techniques, computer vision systems can process vast amounts of visual data, recognize patterns, and continuously improve their performance.

The future of machine learning for computer vision holds immense promise, with new applications and innovations emerging at a rapid pace. As researchers and developers continue to push the boundaries of what’s possible, we can expect even more groundbreaking advancements that will revolutionize the way we live, work, and interact with the world around us. By embracing these cutting-edge technologies, we are not only unlocking new possibilities but also paving the way for a more intelligent, efficient, and connected future.

The post Machine Learning Computer Vision appeared first on PyImageSearch.

Table of Contents

What’s New in PyTorch 2.0? torch.compile

Over the last few years, PyTorch has evolved as a popular and widely used framework for training deep neural networks (DNNs). The success of PyTorch is attributed to its simplicity, first-class Python integration, and imperative style of programming. Since the launch of PyTorch in 2017, it has strived for high performance and eager execution. It has provided some of the best abstractions for distributed training, data loading, and automatic differentiation.

With continuous innovation from the PyTorch team, PyTorch has moved from version 1.0 to the most recent version, 1.13. However, over all these years, hardware accelerators like GPUs have become 15x and 2x faster in compute and memory access, respectively. Thus, to leverage these resources and deliver high-performance eager execution, the team moved substantial parts of PyTorch internals to C++.

On December 2, 2022, the team announced the launch of PyTorch 2.0, a next-generation release that will make training deep neural networks much faster and support dynamic shapes. The stable release of PyTorch 2.0 is planned for March 2023. This blog series aims to understand and test the capabilities of PyTorch 2.0 via its beta release.

In this series, you will learn about Accelerating Deep Learning Models with PyTorch 2.0.

This lesson is the 1st of a 2-part series on Accelerating Deep Learning Models with PyTorch 2.0:

1. What’s New in PyTorch 2.0? torch.compile (today’s tutorial)
2. What’s Behind PyTorch 2.0? TorchDynamo and TorchInductor

To learn what’s new in PyTorch 2.0, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

What’s New in PyTorch 2.0? torch.compile

We start this lesson by learning to install PyTorch 2.0.

Configuring Your Development Environment

Installation

Like previous versions, PyTorch 2.0 is available as a Python pip package. However, to successfully install PyTorch 2.0, your system should have installed the latest CUDA (Compute Unified Device Architecture) versions (11.6 and 11.7). Here’s how you can install the PyTorch 2.0 nightly version via pip:

For CUDA version 11.7

$ pip3 install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

For CUDA version 11.6

$ pip3 install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu116

However, if you don’t have CUDA 11.6 or 11.7 installed on your system, you can download all the required dependencies in the PyTorch nightly binaries with docker.

$ sudo apt install -y nvidia-docker2
$ sudo systemctl restart docker
$ docker pull ghcr.io/pytorch/pytorch-nightly
$ docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

Be sure to specify --gpus all so that your container can access all your GPUs.

Verification

Optionally, you can verify your installation via:

$ git clone https://github.com/pytorch/pytorch
$ cd tools/dynamo
$ python verify_dynamo.py

Overview of PyTorch 2.0

Before understanding what’s new in PyTorch 2.0, let us first understand the fundamental difference between eager and graph executions (Figure 1).

Eager Execution: An eager execution evaluates the operations immediately and at run time. The programs are generally easy to write, test, and debug with a natural Python-like syntax design. However, because of its nature, it fails to fully leverage the capabilities of hardware accelerators like GPUs. PyTorch is a common example that follows eager execution.

Graph Execution: Graph execution, on the other hand, builds a graph of all operations and operands before running. Such an execution is much faster than an eager one, as the graph formed can be optimized to leverage the capabilities of hardware accelerators. However, such programs take more work to write and debug. TensorFlow is a typical example that follows graph execution.

PyTorch has always strived for high performance and eager execution while delivering some of the best abstractions for distributed learning, data loading, and automatic differentiation. To make PyTorch programs faster, the team moved its internals to C++, making the executions faster and less hackable without compromising the flexibility offered by eager mode.

The PyTorch 2.0 release aims to make the training of deep neural networks faster with low memory usage, along with supporting dynamic shapes. In addition, PyTorch 2.0 aims to leverage the capabilities of hardware accelerators and offers better speedups in eager mode.

The backbone of PyTorch 2.0 is four new technologies (TorchDynamo, AOT Autograd, PrimTorch, and TorchInductor) aiming to make PyTorch programs run faster and with less memory.

• TorchDynamo safely captures the PyTorch programs using a new CPython feature called Frame Evaluation API introduced in PEP 523. TorchDynamo can acquire graphs 99% safely, without errors, and with negligible overhead.
• AOT Autograd is the new PyTorch autograd engine that generates ahead-of-time (AOT) backward traces.
• With the PrimTorch project, the team could canonicalize 2000+ PyTorch operations (which used to make its backend challenging) to a set of 250 primitive operators that cover the complete PyTorch backend. This makes it easy to implement any new feature in the PyTorch backend.
• The new OpenAI Triton-based deep learning compiler (TorchInductor) can generate fast code for multiple accelerators and backends.

We will discuss more on these new technologies in a future lesson. This high-level overview should set the background and context on what makes PyTorch 2.0 programs faster.

What’s New in PyTorch 2.0? torch.compile

The core of the PyTorch 2.0 is a torch.compile function that wraps your standard PyTorch model, optimizes it under the hood, and returns a compiled version.

torch.compile Definition

def torch.compile(model: Callable,
  *,
  mode: Optional[str] = "default",
  dynamic: bool = False,
  fullgraph:bool = False,
  backend: Union[str, Callable] = "inductor",
  # advanced backend options go here as kwargs
  **kwargs
) -> torch._dynamo.NNOptimizedModule

Here:

• On Line 1, model is your nn.Module instance. In other words, your standard PyTorch model instance.
• On Line 3, mode specifies how much the compiler should optimize while compiling. There are three types:
  • default mode: compiles your model efficiently without taking too much time to compile.
  • reduce-overhead mode: reduces the framework overhead by a lot more but consumes a small amount of extra memory.
  • max-autotune mode: compiles for a long time, giving you the fastest code it can generate.
• On Line 4, dynamic specifies where the optimization should be done for dynamic shapes. Since specific compiler optimizations are not applicable for dynamic shapes, it is important to specify this before compiling.
• On Line 5, fullgraph compiles the entire program into a single graph. Most users don’t need it unless they are very performance specific.
• On Line 6, backend specifies which compiler backend to use. By default, TorchInductor is used, but a few others are available, like aot_cudagraphs and nvfuser.

torch.compile, in its default, is intended to provide you with most of the speedups PyTorch 2.0 has to offer. Hence you only need to use other modes if you are keen on getting the best speed. Based on our discussion, here (Figure 2) are the three execution modes you can run your program on.

Here’s a quick differentiation between the three optimization modes offered by torch.compile:

Since torch.compile is backward compatible, all other operations (e.g., reading and updating attributes, serialization, distributed learning, inference, and export) would work just as PyTorch 1.x.

Whenever you wrap your model under torch.compile, the model goes through the following steps before execution (Figure 3):

1. Graph Acquisition: The model is broken down and re-written into subgraphs. Subgraphs that can be compiled/optimized are flattened, whereas other subgraphs which can’t be compiled fall back to the eager model.
2. Graph Lowering: All PyTorch operations are decomposed into their chosen backend-specific kernels.
3. Graph Compilation: All the backend kernels call their corresponding low-level device operations.

Now, let’s start some experimentation.

Accelerating DNNs with PyTorch 2.0

Project Structure

We first need to review our project directory structure.

Start by accessing the “Downloads” section of this tutorial to retrieve the source code.

From there, take a look at the directory structure:

├── cnn.py
├── vit.py
├── bert.py
├── utils.py

The project directory contains four files. The utils.py file implements basic utility functions to parse command line arguments and run/report a model’s speed. The cnn.py, vit.py, and bert.py files load a specified CNN (convolutional neural network), ViT (vision transformer), or a BERT (bidirectional encoder representations from transformers) model, compile it with torch.compile, and report its speed on a random input. We will discuss these files in detail in subsequent sections.

Accelerating Convolutional Neural Networks

Using torch.compile is easy and is expected to provide 30%-200% speedups on most models you run daily. However, first, we will look into some utility functions in utils.py to parse command line arguments and run a model on a given input.

Parsing Command Line Arguments and Running a Model

import torch
import time
import numpy as np
import argparse

# command line arguments
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, help='model type', default='resnet50')
    parser.add_argument('--batch_size', type=int, help='Batch size', default=128)
    parser.add_argument('--steps', type=int, help='Steps', default=10)
    parser.add_argument('--mode', type=str, help='Mode', default='default')
    parser.add_argument('--backend', type=str, help='Backend', default='inductor')

    args = parser.parse_args()
    return args

# running a model
def run_model(model, inputs, steps=20):
    # load model on GPU
    model = model.cuda() 
    # define an optimizer
    optimizer = torch.optim.Adam(model.parameters()) 
    times = []
    for step in range(steps):
        begin = time.time()
        # zero gradients
        optimizer.zero_grad() 
        # forward pass
        output = model(inputs.cuda()) 
        # back propagate 
        if not isinstance(output, torch.Tensor):
            output = output.logits
        output.sum().backward() 
        # optimize weights
        optimizer.step() 
        end = time.time()
        # calcuate step time
        times.append(float(end - begin)) 
        print(f"Time for {step}-th forward pass is {end - begin}")
    
    # calcuate median step time
    median = np.median(times) 
    print("Median step time is {:.3f} seconds".format(median))

On Lines 1-4, we import the torch, time, argparse, and numpy libraries. Then, on Lines 7-16, we define the parse_args() function that parses the following command line arguments:

• --model: specifies the model to load (default set to resnet50)
• --batch_size: specifies the batch size of the inputs (default set to 128)
• --steps: specifies the number of steps to run the model (default is set to 10)
• --mode: specifies whether to use default, reduce-overhead, or original mode for compilation. We won’t experiment with the max-autotune mode as it takes very long to compile and doesn’t always work.
• --backend: specifies the compiler backend (default set to inductor)

Then on Lines 19-44, we define the run_model() function that takes model, inputs, and steps as arguments and runs the model training on inputs for a given number of steps. First, on Line 21, we load the model on the GPU. Then on Line 23, we define optimizer over model parameters.

Finally, on Lines 25-40, we run the model training for given steps wherein we pass the given inputs, backpropagate gradients, and update network weights in each step. We print and store the time taken by each step in a list of times. Finally, on Lines 43 and 44, we calculate and print the median step time taken by our compiled model.

Now let’s start experimenting with convolutional neural networks.

Evaluating Convolutional Neural Networks

import torch
from utils import parse_args, run_model

args = parse_args()

# loading pretrained resnet50
model = torch.hub.load('pytorch/vision:v0.10.0', args.model, pretrained=True) 

# compile your model
if args.mode in ['default', 'reduce-overhead']:
        model = torch.compile(model, mode=args.mode, backend=args.backend)   

# random input image
inputs = torch.randn(args.batch_size, 3, 224, 224) 
run_model(model, inputs, args.steps)

We start by loading the torch library and utilities from utils.py (Lines 1 and 2). On Line 4, we read the command line arguments. Then on Line 7, we load the given args.model from TorchHub. If you are unfamiliar with TorchHub, we highly recommend watching our tutorials.

On Lines 10 and 11, we compile the model using torch.compile with specified mode args.mode. Note that if a user specifies any other mode apart from default and reduce-overhead, we return the original model. By default, we use the inductor backend. On Line 14, we define a random input image with a given batch size args.batch_size. Finally, on Line 15, we run and report the time taken by the model using the run_model utility function.

Here is a sample command to run the above code snippet. The following command tests a ResNet-50 model with default mode and batch size 256.

$ python cnn.py --model resnet50 --batch_size 256 --mode default --steps 10

Figure 4 displays how the output should look. Note that your numbers might differ depending on your GPU specs.

When you run the above code snippet, you will notice that the first step takes an abnormally long time while the subsequent steps are faster. This is because the torch.compile is a lazy wrapper and only compiles the model during the first forward pass.

To notice the speedup, you will need to compare the speed of the compiled model with the original model (by using --mode original in the command).

In Figure 5, we compare several convolutional models like ResNets, GoogleNet, AlexNet, SqueezeNet, DenseNet, MobileNet, and Wide ResNet. On average, CNNs give a 10% training speedup on NVIDIA A6000s. Among all the models, MobileNetV2 and SqueezeNet provide close to 20% speedup, while AlexNet and Wide ResNet give <10% speedup. Please note that the speedup might differ depending on your hardware accelerator. You are likely to see more significant speedups with newer GPUs like A100s.

Accelerating Vision Transformers

Similarly, using the torch.compile wrapper, one can speed up a vision transformer for image classification tasks. We will use the PyTorch image models (timm) library that can be installed via pip:

$ pip install timm

For this example, we will refer to the vit.py file in our project directory.

Evaluating Vision Transformers

import torch
import timm
from utils import parse_args, run_model

args = parse_args()

# loading pretrained ViT model
model = timm.create_model(args.model, pretrained=True) 

# compile your model
if args.mode in ['default', 'reduce-overhead']:
        model = torch.compile(model, mode=args.mode, backend=args.backend)   

# random input image
inputs = torch.randn(args.batch_size, 3, 224, 224) 
run_model(model, inputs, args.steps)

Like our previous example, we start by loading torch, the timm library, and utilities from utils.py (Lines 1-3). Next, on Line 5, we read the command line arguments. Then on Line 8, we load the given args.model from TIMM. The remainder of the code is the same.

Here’s a sample command to run the above code snippet. The following command tests a ViT-B/16 (vision transformer base architecture and patch size 16) model with default mode and batch size 256. You can check out the list of available models using timm.list_models().

$ python vit.py --model vit_base_patch16_224 --batch_size 256 --mode default --steps 10

In Figure 6, we compare several state-of-the-art vision transformers. We notice that, on average, transformers give only 2%-3% speedup compared to >10% speedup for CNNs. This is likely because of the self-attention module, which operates on the global view of the image rather than the local view (e.g., convolutions in CNNs). Hence are difficult to optimize. Models (e.g., MLP-Mixer) give negative speedup, on the other hand.

Accelerating BERT

A similar concept works for natural language processing (NLP) models like BERT. We will use the Hugging Face transformers library that can be installed via pip:

$ pip install transformers==4.26.1

For this example, we will refer to the bert.py file in our project directory.

Evaluating BERT

import torch
from transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification
from utils import parse_args, run_model

args = parse_args()

# loading pretrained BERT model
config = AutoConfig.from_pretrained(args.model)
tokenizer = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForSequenceClassification.from_config(config)

# compile your model
if args.mode in ['default', 'reduce-overhead']:
        model = torch.compile(model, mode=args.mode, backend=args.backend)   

# random input text
text = ", ".join(["This is a very long text" for i in range(20)])
inputs = tokenizer(text, return_tensors='pt')
inputs = inputs["input_ids"].repeat(args.batch_size, 1)
run_model(model, inputs, args.steps)

On Lines 1-3, we import the torch, transformers, and utils.py libraries. Next, on Line 5, we read the command line arguments. Then, on Lines 8-10, we load the given args.model and its config and tokenizer from the Hugging Face transformers library.

On Lines 13 and 14, we compile the model using torch.compile with specified mode args.mode. Then on Lines 17-19, we define a dummy tokenized input text with a given batch size args.batch_size. Finally, on Line 20, we run and report the time taken by the model using the run_model utility function.

Here’s a sample command to run the above code snippet. The following command tests a BERT model with default mode and batch size 256.

$ python bert.py --model bert-base-uncased --batch_size 256 --mode default --steps 10

Figure 7 compares some of the state-of-the-art NLP models (e.g., BERT, DistillBERT, and XLM-RoBERTa) from the Hugging Face library. On average, PyTorch 2.0 provides a 5%-6% speedup on these models. On the other hand, DistillBERT achieves a maximum speedup of 8.5%.

Miscellaneous

Different Benchmarks: Figure 8 shows the speedup of PyTorch 2.0 on NVIDIA A100 GPUs across 163 open source models from different libraries (e.g., TIMM, TorchBench, and Hugging Face). At Float32 precision, it runs 21% faster on average, and at AMP (automatic mixed precision), it runs 51% faster on average. The figure reports the uneven weighted average speedup of 0.75 * AMP + 0.25 * float32 since we find AMP is more common in practice.

Different Backends: By default, we have used the “inductor” compiler backend in our experiments so far. However, there are plenty of backends supported by PyTorch 2.0. You can find the list of supported backends using torch._dynamo.list_backends(). Figure 9 compares a few different compiler backends with the default TorchInductor backend for the ResNet-50 model. We can see that TorchInductor, by default, gives the maximum speedup.

You can experiment with other backends. Note that these backends are hardware-dependent; some might not work on your hardware.

What's next? I recommend PyImageSearch University.

Course information:
74 total classes • 84 hours of on-demand code walkthrough videos • Last updated: March 2023
★★★★★ 4.84 (128 Ratings) • 15,800+ 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:

• ✓ 74 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 74 Certificates of Completion
• ✓ 84 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 500+ 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

Summary

PyTorch has been one of the most popular and widely used frameworks for training and deploying deep learning models. Continuous innovation in PyTorch has resulted in an elegant, high-performance, and eager execution framework. PyTorch 2.0, a next-generation release, brings significant speedup in eager execution by leveraging the best of hardware accelerators through the latest technologies (e.g., TorchDynamo, TorchInductor, PrimTorch, and AOT Autograd).

At the core, PyTorch 2.0 introduces torch.compile, a function that wraps your nn.Module instances, optimizes its graph, and provides a fast model for several backends and architectures. Besides being easy to use, torch.compile is backward compatible. All other operations (e.g., reading and updating attributes, serialization, distributed learning, inference, export, etc.) would work just as in PyTorch 1.x.

On 163 open source models from different libraries (e.g., TIMM, TorchBench, and Hugging Face), torch.compile provided 30%-200% speedups on NVIDIA A100s. Moreover, at Float32 precision, it runs 21% faster on average, and at AMP (automatic mixed precision), it runs 51% faster on average. We also experimented on NVIDIA A6000s and observed that PyTorch 2.0 could provide up to 20% speedup on vision architectures (e.g., SqueezeNet, DenseNet, etc.).

PyTorch has always strived for high performance and eager execution while delivering some of the best abstractions for distributed learning, data loading, and automatic differentiation. With this new release, training deep neural networks in eager modes will become much faster!

Citation Information

Mangla, P. “What’s New in PyTorch 2.0? torch.compile,” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/fh15d

@incollection{Mangla_2023_PT2TC,
  author = {Puneet Mangla},
  title = {What's New in PyTorch 2.0? torch.compile},
  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/fh15d},
}

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!

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!

Download the code!

Website

The post What’s New in PyTorch 2.0? torch.compile appeared first on PyImageSearch.

Table of Contents

Training and Making Predictions with Siamese Networks and Triplet Loss

In this tutorial, we will learn to train our Siamese network based face recognition application using Keras and TensorFlow. Furthermore, we will discuss how we can use our model to make predictions in real-time.

In the previous tutorial of this series, we tried to understand the formulation of triplet loss. We discussed how it could be used to learn an embedding space where “similar faces” (i.e., from the same person) reside close to each other and “dissimilar face” (i.e., from different people) reside farther apart. Additionally, we discussed a typical Siamese network pipeline and how it can be used to build our face recognition model.

Furthermore, we implemented the triplet loss and developed our Siamese network based face recognition pipeline in Keras and TensorFlow.

In this tutorial, we will take this further and learn how to train our face recognition model using Keras and TensorFlow. Once our model is trained, we will use it to predict new unseen faces in real-time.

This lesson is the 4th in a 5-part series on Siamese networks and their application in face recognition:

1. Face Recognition with Siamese Networks, Keras, and TensorFlow
2. Building a Dataset for Triplet Loss with Keras and TensorFlow
3. Triplet Loss with Keras and TensorFlow
4. Training and Making Predictions with Siamese Networks and Triplet Loss (this tutorial)
5. Evaluating Siamese Network Accuracy (ROC, Precision, and Recall) with Keras and TensorFlow

To learn how to train and make predictions with Siamese networks and triplet loss, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Training and Making Predictions with Siamese Networks and Triplet Loss

In the second part of this series, we developed the modules required to build the data pipeline for our face recognition application. Furthermore, in the previous tutorial, we developed modules to build our Siamese Model and triplet loss function. In this tutorial, we will put everything together and build our end-to-end face recognition application using the modules that we built previously. Additionally, we will learn to train our end-to-end face recognition model and discuss how we can make predictions using it in real-time.

For this tutorial, we will use Keras and TensorFlow, as we have done in the previous parts of this series. Keras and TensorFlow provide various functionalities that allow us to elegantly put all the modules together and develop our end-to-end pipeline. This provides a simple and intuitive way in which the different parts of our application can communicate and work in tandem to create an efficient and effective face recognition application. Furthermore, Keras provides a simple API that helps us implement, compile, train, and save our model with minimal code.

Figure 1 depicts the overview of our face recognition pipeline and shows how the modules we built in the previous parts of this series work together to develop our final end-to-end application.

Let us revisit our modules and understand the structure and flow of our application.

First, we develop the data pipeline, as shown in Figure 1. Next, the Data Generator (created using the TripleGenerator Class) is used to create our training and validation Dataset with the help of the tf.data.Dataset functionality provided by TensorFlow. This dataset is then used to create our DataLoaders, allowing us to apply pre-processing transformations using the MapFunction Class and generate batches of data samples. Finally, the data pipeline returns two data loaders (i.e., trainDs and valDs for training and validation, respectively).

Next, we develop our Siamese network pipeline. We create our embedding module with the help of the get_embedding_module() function, which we had defined in earlier tutorials. Then, we use the embedding module to embed the anchor, positive, and negative images to build our Siamese network using the get_siamese_network() function. Finally, we pass our Siamese network to the SiameseModel Class which implements the triplet loss and training and test step code.

In the end, we compile and train our model using Keras and finally save our trained model so we can use it in the inference phase for making real-time predictions.

Now that we have discussed the overview of our pipeline, let us dive into the code to train and make predictions with our Siamese network based face recognition application.

Configuring Your Development Environment

To follow this guide, you need to have the TensorFlow and OpenCV libraries installed on your system.

Luckily, both TensorFlow and OpenCV are pip-installable:

$ pip install tensorflow
$ pip install opencv-contrib-python

If you need help configuring your development environment for OpenCV, we highly recommend that you read our pip install OpenCV guide — it will have you up and running in a matter of minutes.

Having Problems Configuring Your Development Environment?

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 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

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.

├── crop_faces.py
├── face_crop_model
│   ├── deploy.prototxt.txt
│   └── res10_300x300_ssd_iter_140000.caffemodel
├── inference.py
├── pyimagesearch
│   ├── config.py
│   ├── dataset.py
│   └── model.py
└── train.py

In the previous tutorial, we presented a step-by-step walkthrough of our model.py from the pyimagesearch folder, which allows us to implement the triplet loss function and build our Siamese network model.

In this tutorial, we will discuss in detail the train.py file, which implements the code to train our face recognition pipeline, and the inference.py file, which will help us make predictions using our Siamese network based face recognition application.

Training Our Siamese Network Model with Triplet Loss

Now that we have discussed the overview of our face recognition pipeline and the function performed by the modules we have built, let us put everything together and train our Siamese network based face recognition pipeline using Keras and TensorFlow.

We open our train.py file and get started.

# USAGE
# python train.py

# import the necessary packages
from pyimagesearch.dataset import TripletGenerator
from pyimagesearch.model import get_embedding_module
from pyimagesearch.model import get_siamese_network
from pyimagesearch.model import SiameseModel
from pyimagesearch.dataset import MapFunction
from pyimagesearch import config
from tensorflow import keras
import tensorflow as tf
import os

# create the data input pipeline for train and val dataset
print("[INFO] building the train and validation generators...")
trainTripletGenerator = TripletGenerator(
	datasetPath=config.TRAIN_DATASET)
valTripletGenerator = TripletGenerator(
	datasetPath=config.TRAIN_DATASET)
print("[INFO] building the train and validation `tf.data` dataset...")
trainTfDataset = tf.data.Dataset.from_generator(
	generator=trainTripletGenerator.get_next_element,
	output_signature=(
		tf.TensorSpec(shape=(), dtype=tf.string),
		tf.TensorSpec(shape=(), dtype=tf.string),
		tf.TensorSpec(shape=(), dtype=tf.string),
	)
)
valTfDataset = tf.data.Dataset.from_generator(
	generator=valTripletGenerator.get_next_element,
	output_signature=(
		tf.TensorSpec(shape=(), dtype=tf.string),
		tf.TensorSpec(shape=(), dtype=tf.string),
		tf.TensorSpec(shape=(), dtype=tf.string),
	)
)

First, we import the important modules we built earlier to train our face recognition model on Lines 5-13. In previous tutorials, we noted that the pyimagesearch folder contains the code for the dataset module (dataset.py), the model definition (model.py), and the configuration file (config.py), which we discussed in detail. We will now use these modules to train our face recognition application.

We start by importing the TripletGenerator class from the dataset module (Line 5) and the get_embedding_module, get_siamese_network, and SiameseModel class from the model definition (Lines 6-8). We also import the MapFunction class and the config file on Lines 9 and 10, respectively. Finally, we import the Keras, TensorFlow, and os packages on Lines 11-13.

Next, we develop our data pipeline to allow us to sample batches for training and validation. We use the TripletGenerator class (that we developed earlier) to define the training data generator (i.e., trainTripletGenerator) and validation data generator (i.e., valTripletGenerator) on Lines 17-19. The TripletGenerator class takes as input the path to the respective dataset (i.e., config.TRAIN_DATASET) as shown on Line 20.

Now that we have defined our data generators, we use the tf.data.Dataset.from_generator functionality to define our training and validation datasets on Lines 22-37. First, we define our training dataset (i.e., trainTfDataset) on Lines 22-29. Note that tf.data.Dataset.from_generator takes as input a callable generator function (i.e., trainTripletGenerator.get_next_element), whose outputs must be compatible with the output format defined by the output_signature argument.

Similarly, we create the validation dataset (i.e., valTfDataset) using tf.data.Dataset.from_generator on Lines 30-37.

# preprocess the images
mapFunction = MapFunction(imageSize=config.IMAGE_SIZE)
print("[INFO] building the train and validation `tf.data` pipeline...")
trainDs = (trainTfDataset
    .map(mapFunction)
    .shuffle(config.BUFFER_SIZE)
    .batch(config.BATCH_SIZE)
    .prefetch(config.AUTO)
)
valDs = (valTfDataset
    .map(mapFunction)
    .batch(config.BATCH_SIZE)
    .prefetch(config.AUTO)
)

# build the embedding module and the siamese network
print("[INFO] build the siamese model...")
embeddingModule = get_embedding_module(imageSize=config.IMAGE_SIZE)
siameseNetwork =  get_siamese_network(
	imageSize=config.IMAGE_SIZE,
	embeddingModel=embeddingModule,
)
siameseModel = SiameseModel(
	siameseNetwork=siameseNetwork,
	margin=0.5,
	lossTracker=keras.metrics.Mean(name="loss"),
)

# compile the siamese model
siameseModel.compile(
	optimizer=keras.optimizers.Adam(config.LEARNING_RATE)
)

# train and validate the siamese model
print("[INFO] training the siamese model...")
siameseModel.fit(
	trainDs,
	steps_per_epoch=config.STEPS_PER_EPOCH,
	validation_data=valDs,
	validation_steps=config.VALIDATION_STEPS,
	epochs=config.EPOCHS,
)

# check if the output directory exists, if it doesn't, then
# create it
if not os.path.exists(config.OUTPUT_PATH):
	os.makedirs(config.OUTPUT_PATH)

# save the siamese network to disk
modelPath = config.MODEL_PATH
print(f"[INFO] saving the siamese network to {modelPath}...")
keras.models.save_model(
	model=siameseModel.siameseNetwork,
	filepath=modelPath,
	include_optimizer=False,
)

On Line 40, we define the pre-processing operations that we want to apply to our data samples using the MapFunction, which takes the config.IMAGE_SIZE parameter as an argument.

Finally, on Lines 42-52, we use the training and validation dataset (i.e., trainTfDataset and valTfDataset) to define our training and validation data loaders (i.e., trainDs and valDs). Note that TensorFlow allows us to apply different functionalities to our generated data samples.

• The map function (which takes as an argument our mapFunction class) applies pre-processing transformations to our data samples.
• The shuffle functionality (which takes as argument config.BUFFER_SIZE) to randomly sample elements from a buffer of config.BUFFER_SIZE number of elements.
• The batch functionality (which takes as argument config.BATCH_SIZE) allows us to sample batches of data samples with the number of elements per batch defined by the config.BATCH_SIZE argument.
• The prefetch functionality directs TensorFlow to prepare later elements while current elements are being processed.

Now that we have created the data pipeline, we can define our Siamese Model. We first build our embedding module using the get_embedding_module function, which takes as input the imageSize (Line 56).

Next, we use the get_siamese_network function, which takes as an argument the imageSize and embeddingModule to build and return our siameseNetwork (Lines 57-60).

Now that we have our siameseNetwork, we use the SiameseModel class to build our Siamese network based face recognition model (Lines 61-65). The SiameseModel class takes as arguments the siameseNetwork, the margin distance that we discussed earlier, and the keras.metrics.Mean(name="loss") metric.

We can now compile our model using siameseModel.compile and use the Adam optimizer with the learning rate equal to config.LEARNING_RATE, as shown on Lines 68-70.

Finally, we use the fit functionality to train our Siamese network based face recognition model (Lines 74-80). This takes as input the training data loader (i.e., trainDs), the steps_per_epoch, the validation data loader (i.e., valDs), the number of validation steps (i.e., config.VALIDATION_STEPS), and the total number of epochs (i.e., config.EPOCHS).

Next, we prepare to save our model. We check if the output directory exists, and if it does not, we create it (Lines 84 and 85). On Line 88, we define the modelPath, and on Lines 90-94, we use the keras.models.save_model function to save our trained model.

Making Predictions with Our Siamese Network Based Face Recognition Model

Now that we have discussed the code required to train our model, let us implement the code to make predictions in real-time with our trained Siamese network based face recognition application.

We start by opening the inference.py file.

# USAGE
# python inference.py

# import the necessary packages
from pyimagesearch.dataset import TripletGenerator
from pyimagesearch.dataset import MapFunction
from pyimagesearch.model import SiameseModel
from matplotlib import pyplot as plt
from pyimagesearch import config
from tensorflow import keras
import tensorflow as tf
import os

# create the data input pipeline for test dataset
print("[INFO] building the test generator...")
testTripletGenerator = TripletGenerator(
	datasetPath=config.TEST_DATASET)
print("[INFO] building the test `tf.data` dataset...")
testTfDataset = tf.data.Dataset.from_generator(
	generator=testTripletGenerator.get_next_element,
	output_signature=(
		tf.TensorSpec(shape=(), dtype=tf.string),
		tf.TensorSpec(shape=(), dtype=tf.string),
		tf.TensorSpec(shape=(), dtype=tf.string),
	)
)
mapFunction = MapFunction(imageSize=config.IMAGE_SIZE)
testDs = (testTfDataset
	.map(mapFunction)
	.batch(4)
	.prefetch(config.AUTO)
)

As always, we first import the necessary modules like TripletGenerator, MapFunction, and SiameseModel on Lines 5-7. Additionally, we import the necessary packages like pyplot (from matplotlib), config, keras, tensorflow, and os, as shown on Lines 8-12.

We start by using the TripletGenerator class to define the test data generator (i.e., testTripletGenerator), which takes as input the path to the test dataset (i.e., config.TEST_DATASET), as shown on Lines 16 and 17.

Then, we use the tf.data.Dataset.from_generator function to define our test dataset (i.e., testTfDataset). Similar to what we discussed for the training pipeline, this function takes as inputs the test data generator function testTripletGenerator.get_next_element and the output_signature, as shown on Lines 19-26.

Similar to what we did in the training phase, now we use the MapFunction class to define the pre-processing transformations for the test set (Line 27) and define our test data loader testDs (Lines 28-32). Here, we use the map functionality to apply the pre-process transformations and use a batch size of 4 as shown.

# load the siamese network from disk and build the siamese model
modelPath = config.MODEL_PATH
print(f"[INFO] loading the siamese network from {modelPath}...")
siameseNetwork = keras.models.load_model(filepath=modelPath)
siameseModel = SiameseModel(
	siameseNetwork=siameseNetwork,
	margin=0.5,
	lossTracker=keras.metrics.Mean(name="loss"),
)

# load the test data
(anchor, positive, negative) = next(iter(testDs))
(apDistance, anDistance) = siameseModel((anchor, positive, negative))
plt.figure(figsize=(10, 10))
rows = 4
for row in range(rows):
	plt.subplot(rows, 3, row * 3 + 1)
	plt.imshow(anchor[row])
	plt.axis("off")
	plt.title("Anchor image")
	plt.subplot(rows, 3, row * 3 + 2)
	plt.imshow(positive[row])
	plt.axis("off")
	plt.title(f"Positive distance: {apDistance[row]:0.2f}")
	plt.subplot(rows, 3, row * 3 + 3)
	plt.imshow(negative[row])
	plt.axis("off")
	plt.title(f"Negative distance: {anDistance[row]:0.2f}")

# check if the output directory exists, if it doesn't, then
# create it
if not os.path.exists(config.OUTPUT_PATH):
	os.makedirs(config.OUTPUT_PATH)

# save the inference image to disk
outputImagePath = config.OUTPUT_IMAGE_PATH
print(f"[INFO] saving the inference image to {outputImagePath}...")
plt.savefig(fname=outputImagePath)

Next, we define the modelPath, where our trained Siamese model is stored (Line 35), and use the keras.models.load_model function to load our model (i.e., siameseNetwork) (Line 37). Now that we have our pre-trained siameseNetwork, we use it to build our siameseModel using the SiameseModel class (Lines 38-42).

Now that we have defined our test data pipeline, we can sample from the test set and see our face recognition model in action.

We use the iter() method to convert our testDs to an iterator and then use the next() method to sample a batch of test data. We store the output as a tuple (i.e., (anchor, positive, negative)) (Line 45).

We then pass these samples as input to the siameseModel and get the distance between the anchor and positive samples and the distance between the anchor and negative samples (i.e., (apDistance, anDistance)) (Line 46).

We can now create a plot using matplotlib to visualize our samples. We start by defining our figure with figsize=(10, 10) using the plt.figure() function, as shown on Line 47. We also define the number of rows in our plot, which is equal to the batch size of our test loader (Line 48).

Next, we iterate over the different rows, and for each row, we create a subplot for the anchor image, positive image, and the negative image, as shown on Lines 49-61.

Finally, we prepare to save our inference images. We check if the output directory where we want to store our inference output exists, and if it doesn’t, then we create it (Lines 65 and 66).

We then define the outputImagePath (Line 69) and use the plt.savefig function to save our inference image (Line 71).

What's next? I recommend PyImageSearch University.

Course information:
69 total classes • 73 hours of on-demand code walkthrough videos • Last updated: March 2023
★★★★★ 4.84 (128 Ratings) • 15,800+ 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:

• ✓ 74 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 74 Certificates of Completion
• ✓ 84 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 500+ 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

Summary

In this tutorial, we discussed how to train our Siamese network based face recognition model using Keras and TensorFlow. Specifically, we tried to understand how the modules we built in the previous parts of this series come together to form our face recognition application.

Furthermore, we discussed and implemented the code to predict new unseen face images in real-time using our trained face recognition model.

In the upcoming tutorials of this series, we will evaluate the performance of our face recognition model using different metrics.

Citation Information

Chandhok, S. “Training and Making Predictions with Siamese Networks and Triplet Loss,” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/avjyi

@incollection{Chandhok_2023_training_and_making,
  author = {Shivam Chandhok},
  title = {Training and Making Predictions with Siamese Networks and Triplet Loss},
  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/avjyi},
}

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!

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!

Download the code!

Website

The post Training and Making Predictions with Siamese Networks and Triplet Loss appeared first on PyImageSearch.

Table of Contents

Train a MaskFormer Segmentation Model with Hugging Face 🤗 Transformers

In this tutorial, we will learn how to train MaskFormer on a Colab Notebook to perform panoptic segmentation.

We will use the 🤗Transformers and Datasets libraries to load and train a model on the Scene Parsing dataset and the Hub library to publish our model. To measure the performance of our trained model, we will use the 🤗Evaluate library.

You will learn how to:

• Load and preprocess the dataset
• Use the transformers Trainer class to train models
• Evaluate your trained segmentation model

Project Structure

For this tutorial, we will use a Colab Notebook. Feel free to jump over to the notebook or create a new notebook and code along!

Requirements

Before starting, please create an account on the Hugging Face Hub. This will allow us to push our trained models and preprocessing pipeline to the Hub for public or private use.

Bonus

Hugging Face has multiple Python libraries under its umbrella: datasets, transformers, evaluate, and accelerate, just to name a few! To familiarize yourself with these libraries, feel free to jump into their documentation.

An Introduction to Image Segmentation

Image segmentation is a massively popular computer vision task that deals with the pixel-level classification of images. Unlike the object detection task, where the goal is to predict object classes and corresponding bounding boxes, the image segmentation task aims to learn how to segment similar groups of pixels and predict the object classes of these groups.

As we might have guessed, segmentation provides higher granularity information about images compared to object detection and has a wide variety of use cases (e.g., self-driving, quality assurance, medical imaging, fashion, satellite imaging, etc.).

The three subtasks of image segmentation are:

1. Semantic segmentation
2. Instance segmentation
3. Panoptic segmentation

While semantic and instance segmentation are often used interchangeably, these tasks require different model architectures. Let’s see how these subtasks differ from each other with the example shown in Figure 1.

The image includes sky, road, multiple persons, cars, trees, and traffic light instances.

Semantic segmentation: The model tries to map each pixel to a class label (e.g., sky: 0, road: 1, person: 2, car: 3, tree: 4, traffic_light: 5). The pixels of all cars in the image are labeled as 3.

Instance segmentation: The model tries to map each pixel to a class label and differentiate between the instances of the same object class. Pixels of car 1 and car 2 have the same class id but not the same instance id. There can also be two sky instances due to a person standing in the foreground and splitting the sky pixels into two segments. However, background pixels such as sky and road are often disregarded as object categories in instance segmentation tasks.

Panoptic segmentation: The model provides the best semantic and instance segmentation. It allows assigning a single instance id to non-overlapping segments (e.g., the split sky segments). Background classes such as sky and road are always included in the class labels.

The panoptic segmentation models can be considered simple extensions of instance segmentation models. However, note that the panoptic segmentation task is very useful as it provides a more relevant context in many use cases (e.g., self-driving cars).

In short, semantic segmentation is our vanilla task, instance segmentation combines semantic segmentation with object detection, and panoptic segmentation combines semantic and instance segmentation and can handle multiple instances of the same object class when there shouldn’t be (e.g., there is only one sky!).

To learn more about the different segmentation subtasks, check out a previous blog post.

What Is MaskFormer?

MaskFormer is a state-of-the-art image segmentation model by Facebook Research that proposes a unified framework for semantic, instance, and panoptic segmentation tasks. See Figure 2 for an overview of the MaskFormer architecture.

The official MaskFormer includes checkpoints for models trained on ADE20K, Cityscapes, COCO, and Mapillary Vistas across all tasks and multiple model sizes.

In this tutorial, we will use the Hugging Face implementation of MaskFormer, which allows us to load, train, and evaluate the model on a custom dataset with a few lines of code. In addition, we will use an instance segmentation dataset to fine-tune a MaskFormer model trained on another dataset.

Do not worry if it is your first time with Hugging Face. We have got your back! This tutorial is the right mix to learn about the ecosystem inside out.

🤗Transformers Ecosystem

The 🤗 transformers library enables users to easily use, experiment with, and fine-tune pre-trained state-of-the-art models. The library also designs custom architectures for new research. Other Hugging Face libraries (e.g., Datasets and Evaluate) are designed to make it easier to load datasets and evaluate model performance. In addition, they can be smoothly integrated with transformers pipelines.

Before training a model using 🤗 transformers, we need to:

1. Load a dataset with the 🤗 Datasets library or write a custom dataloader to load our own data.
2. Create a custom XXXModelConfig object to define the model architecture, hyperparameters and labels (where applicable).
3. Initialize a XXXModel model with the custom XXXModelConfig config object.
4. Preprocess the dataset using the same configuration (e.g., expected image size).

To train a MaskFormer model, all we need to do is replace XXX in the above steps with MaskFormer.

How to Train an Instance Segmentation Model with MaskFormer?

This tutorial focuses on fine-tuning an image segmentation model (MaskFormer) on a custom dataset. This is meant to provide an understanding of how fine-tuning works and an overview of the Hugging Face ecosystem. We structure the tutorial in the following manner:

1. Setup
2. Loading the Dataset
3. Preprocessing the Dataset
4. Fine-tuning the MaskFormer Model
5. Inference and Performance Evaluation

Setup

Let’s start by installing the dependencies.

In this tutorial, we use the following:

• The 🤗 Datasets library to load the MIT Scene Parsing Benchmark dataset.
• The 🤗 Transformers library to preprocess the data and fine-tune a MaskFormer model.
• The 🤗 Evaluate library to benchmark the trained model’s performance.
• The Hub client library to push our trained model to the 🤗 Hub.
• The Albumentations library, in conjunction with the transformers preprocessor class, to augment our data for training.

# Install the necessary dependencies
!pip install datasets -qq
!pip install evaluate -qq
!pip install albumentations -qq
!pip install git+https://github.com/huggingface/transformers.git -qq


# We will use this to push our trained model to HF Hub
!pip install huggingface_hub -qq

We can now import MaskFormer and other dependencies:

# Import the necessary packages
import random
from tqdm.auto import tqdm
import matplotlib.pyplot as plt

import torch
from torch import nn
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

import albumentations as A

import numpy as np
import pandas as pd

from datasets import load_dataset
from transformers import (
    MaskFormerConfig,
    MaskFormerImageProcessor,
    MaskFormerModel,
    MaskFormerForInstanceSegmentation,
)
import evaluate
from huggingface_hub import notebook_login

Loading the Dataset

About the Dataset

This tutorial will use the Scene Parsing dataset to train an instance segmentation model. This dataset is an extension of the ADE20K dataset, consisting of over 20K images and their segmentation annotations. It is named after Adela Barriuso, who single-handedly annotated the dataset. ADE20K is among the most popular and widely used semantic segmentation datasets. The Scene Parsing dataset has 150 semantic labels for the semantic segmentation task and 100 for the instance segmentation task.

Loading a Dataset with 🤗 Datasets

We can use the 🤗 Datasets library to load this dataset with a single line of code and create train, validation, and test generators.

Note: Downloading the dataset takes 1.2 GB of disk space. If you don’t want to download the whole dataset, you can simply pass in the streaming=True argument to create an iterable dataset where samples are downloaded as you iterate over them.

Now, let’s download the dataset from the 🤗 Hub with the load_dataset() function. Calling this function will download the dataset and return an iterable DatasetDict object. Note that we can also pass in an optional split argument to download train, validation, and test sets separately.

>>> DatasetDict({
>>>     train: Dataset({
>>>         features: ['image', 'annotation'],
>>>         num_rows: 20210
>>>     })
>>>     test: Dataset({
>>>         features: ['image', 'annotation'],
>>>         num_rows: 3352
>>>     })
>>>     validation: Dataset({
>>>         features: ['image', 'annotation'],
>>>         num_rows: 2000
>>>     })
>>> })

Now let’s load the Scene Parsing dataset’s instance segmentation subset:

train = load_dataset("scene_parse_150", "instance_segmentation", split="train")
validation = load_dataset("scene_parse_150", "instance_segmentation", split="validation")
test = load_dataset("scene_parse_150", "instance_segmentation", split="test")

Each instance in the train, validation, and test set dictionaries of Scene Parsing to have two keys: image and annotation.

Note: Different datasets on the hub have different keys based on their target task, and we can refer to the dataset page to see what they contain (see the Scene Parsing dataset page).

print(train[0])
>>> {'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=256x256 at 0x7F48DF5D6310>,
 'annotation': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=256x256 at 0x7F48DF5D6250>}

Dataset Visualization

Let’s print out an image randomly selected from our train set.

# Grab a random index of the training dataset
print("[INFO] Displaying a random image and its annotation...")
index = random.randint(0, len(train))

# Using the random index grab the corresponding datapoint
# from the training dataset
image = train[index]["image"]
image = np.array(image.convert("RGB"))
annotation = train[index]["annotation"]
annotation = np.array(annotation)

# Plot the original image and the annotations
plt.figure(figsize=(15, 5))
for plot_index in range(3):
    if plot_index == 0:
        # If plot index is 0 display the original image
        plot_image = image
        title = "Original"
    else:
        # Else plot the annotation maps
        plot_image = annotation[..., plot_index - 1]
        title = ["Class Map (R)", "Instance Map (G)"][plot_index - 1]

    # Plot the image
    plt.subplot(1, 3, plot_index + 1)
    plt.imshow(plot_image)
    plt.title(title)
    plt.axis("off")

Instance Segmentation Annotations

In this section, we will look at instance segmentation annotations. As described in the Scene Parsing dataset page, the instance annotation masks are stored in RGB image format and structured as follows:

• The R(ed) channel encodes category ID.
• The G(reen) channel encodes instance ID.
• Instance IDs are assigned per image such that each object in an annotation image has a different instance ID (regardless of its class ID). In contrast, different annotation images can have objects with the same instance ID. Each image in the dataset has < 256 object instances.

We can refer to this file (train split) for the 100 class labels included in the instance segmentation subset of Scene Parsing.

Important Notes

Semantic and instance segmentation subsets of Scene Parsing don’t have the same id2label mapping. To find the mapping between the semantic categories for instance_segmentation and semantic_segmentation and the ADE20K dataset, we need to refer to Mapping.txt.

Preprocessing the Dataset

For any 🤗 Transformers vision model, we can simply use the corresponding FeatureExtractor or ImageProcessor class to preprocess raw images (and segmentation maps) and convert them to the format expected by the model.

As Scene Parsing is a large image dataset, we will call the preprocessor as we iterate over the training and validation loops instead of preprocessing all data at once.

FeatureExtractor and ImageProcessor Classes

Each model on the 🤗 Hub has a corresponding FeatureExtractor or ImageProcessor or Processor (for multi-modal models) class that is used to preprocess the raw image (or text, speech, etc.) such that it has the expected input shape and format used by the model. Some classes also include convenient post-processing methods to convert raw model outputs to final predictions.

Same as the Model and Config class, we can use the from_pretrained() method to retrieve the processor of any pre-trained model.

from transformers import XXXModelImageProcessor

# Use default configuration
preprocessor = XXXModelImageProcessor.from_pretrained(<MODEL_NAME>)

train = load_dataset(<DATASET>)["train"]
for data in train:
    model = XXXModelImageProcessor(data[<target_key>])

Defining a Preprocessor Configuration

An important point to note here is that the preprocessor configuration needs to be consistent with the model configuration chosen for the input data to be preprocessed correctly (e.g., expected image size, normalization, augmentation, …).

As there are no model checkpoints available on the hub that were trained on the Scene Parsing dataset, we will create our custom preprocessor configuration.

# Create the MaskFormer Image Preprocessor
processor = MaskFormerImageProcessor(
    reduce_labels=True,
    size=(512, 512),
    ignore_index=255,
    do_resize=False,
    do_rescale=False,
    do_normalize=False,
)

>>> preprocessor
>>> >>> 
>>> MaskFormerImageProcessor {
>>>   "_max_size": 1333,
>>>   "do_normalize": false,
>>>   "do_reduce_labels": true,
>>>   "do_rescale": false,
>>>   "do_resize": false,
>>>   "ignore_index": 255,
>>>   "image_mean": [
>>>     0.485,
>>>     0.456,
>>>     0.406
>>>   ],
>>>   "image_processor_type": "MaskFormerImageProcessor",
>>>   "image_std": [
>>>     0.229,
>>>     0.224,
>>>     0.225
>>>   ],
>>>   "resample": 2,
>>>   "rescale_factor": 0.00392156862745098,
>>>   "size": {
>>>     "height": 512,
>>>     "width": 512
>>>   },
>>>   "size_divisor": 32
>>> }

The preprocessor configuration is compact with mostly easy-to-understand fields. As specified in the MaskFormerImageProcessor docs, the reduce_label argument handles class labels that start from 1.

When set to True, the reduce_label argument replaces the background / null class pixels (denoted with 0) with the ignore_index specified and decrements all other labels by 1. This comes in handy as Python indexing starts from 0, and we will need the labels to start from 0 to calculate the loss terms correctly.

Note: Instance segmentation dataset labels start from 1 while 0 is reserved for the null / background class to be ignored.

How to Use MaskFormerImageProcessor

Remember, MaskFormer is a unified framework supporting multiple segmentation tasks: semantic, instance, and panoptic segmentation. Hence, its preprocessor is designed to handle both semantic and instance segmentation annotations.

Let’s jump into MaskFormerImageProcessorʼs documentation once again. After initializing our image processor, we need to pass in the following arguments to process our data:

• images
• segmentation_maps (optional)
• instance_id_to_semantic_id (optional)

Let’s break it down. images are the images we’d like to segment. segmentation_maps can either be pixel-wise class annotations (default) or pixel-wise instance id annotations. If we pass in instance-level annotations, we also need to pass in an instance_id_to_semantic_id dictionary (List[Dict[int, int]] or Dict[int, int]) that maps instance ids to class ids. During inference, we only need to pass in images as input.

Remember that our dataset’s annotation field stores annotations as images, where the Red channel and Green channel values correspond to pixel-wise class and instance labels, respectively. Since we would like to train instance segmentation models, we need to pass in pixel-wise instance segmentation maps as the segmentation_maps argument and create either a global instance_id_to_semantic_id dictionary or a list of dictionary for per-image instance id to class id mapping. Hence, we will need to do additional light preprocessing before feeding the data into our preprocessor.

Fine-Tuning the MaskFormer Model

We are now ready to fine-tune the MaskFormer model on the Scene Parsing dataset for the instance segmentation task.

Defining a Configuration

Each pre-trained model on 🤗 Hub has its own repository and a config.json file (see an example over here. The configuration file defines the model architecture (e.g., number of layers, hidden dimension), hyperparameters (e.g., dropout ratio) and other relevant metadata (e.g., labels in the dataset).

We can load the specific configuration of any available pre-trained model simply using the from_pretrained() method.

from transformers import MaskFormerConfig

config = MaskFormerConfig.from_pretrained("facebook/maskformer-swin-base-coco")
config

We can also initialize a XXXModelConfig object, which defaults to values specified in the configuration_XXXModel.py (like MaskFormerʼs configuration script) file of the respective model.

To create a custom XXXModelConfig object, we can either create a custom config.json file and pass it to the from_pretrained() method or initialize XXXModelConfig with default values and edit it as we would edit a *dictionary*.

from transformers import MaskFormerConfig

# Option 1 - Use custom configuration file
config = MaskFormerConfig.from_pretrained(<PATH TO YOUR LOCAL CUSTOM CONFIG JSON>)

# Option 2 - Edit default configuration
config = MaskFormerConfig()
config.embed_dim = 256 # Editing the configuration

We need to start by defining a model configuration in order to initialize a MaskFormerForInstanceSegmentation model, as discussed above.

For this tutorial, we will use the configuration of the pre-trained facebook/maskformer-swin-base-ade model with a few modifications and fine-tune MaskFormer on the Scene Parsing instance segmentation subset. Note that this model is trained on the ADE20K semantic segmentation dataset and can only perform the semantic segmentation task.

If we wanted to train MaskFormer from scratch, we could also define a custom configuration by initializing a MaskFormerConfig object and changing the default values or loading a custom config.json file.

# Define the name of the model
model_name = "facebook/maskformer-swin-base-ade"

# Get the MaskFormer config and print it
config = MaskFormerConfig.from_pretrained(model_name)
print("[INFO] displaying the MaskFormer configuration...")
print(config)

The config.json file shows that the label2id and id2label fields are specific to the ADE20K dataset. Therefore we’ll need to customize it.

To do this, we will download the Scene Parsing instance segmentation labels and reduce the labels such that the labels start from 0 instead of 1.

# Download SceneParsing instance segmentation labels
!wget https://raw.githubusercontent.com/CSAILVision/placeschallenge/master/instancesegmentation/instanceInfo100_train.txt

# Get a modified version of the id2label and label2id
data = pd.read_csv(
    "/content/instanceInfo100_train.txt",
    sep="\t",
    header=0,
    on_bad_lines="skip",
)
id2label = {id: label.strip() for id, label in enumerate(data["Object Names"])}
label2id = {v: k for k, v in id2label.items()}

# Edit MaskFormer config labels
config.id2label = id2label
config.label2id = label2id

Notes on Fine-Tuning HF Models

To fine-tune pre-trained models on custom datasets with different labels, we will need to make a few modifications. The classification or final prediction heads of each model are specific to the dataset it was trained on. However, we can still use the learned parameters of the base model to accelerate our training progress. This ensures that we start the training from a close-to-ideal point in terms of model parameters.

In practice, we will simply:

1. Initialize a MaskFormerForInstanceSegmentation model with random weights
2. Load a pre-trained MaskFormerModel with a compatible configuration
3. Replace the random base model (MaskFormerModel) weights of our MaskFormerForInstanceSegmentation model with that of the pre-trained MaskFormerModel model

We already know that the XXXModelConfig object defines the model architecture, hyperparameters, and the relevant metadata to create and initialize a model. This is why the first and foremost step would be to create a XXXModelConfig object.

Note: We use XXXModelConfig to make the code more generic. If you want to use your favorite model, just replace XXX with the model’s name. 🤗 Hub has a list of models from which you can choose. Head on to the 🤗 Hub to know more.

from transformers import XXXModelConfig, XXXModel

# Use default configuration
config = XXXModelConfig()

# Initialize model with config
model = XXXModel(config)

Initializing a MaskFormer Model

Now we can define our custom configuration and initialize MaskFormer for training!

# Use the config object to initialize a MaskFormer model with randomized weights
model = MaskFormerForInstanceSegmentation(config)

# Replace the randomly initialized model with the pre-trained model weights
base_model = MaskFormerModel.from_pretrained(model_name)
model.model = base_model

MaskFormer Model Inputs

The MaskFormerForInstanceSegmentation model expects the following inputs during training:

• pixel_values: Tensor of shape (batch_size, num_channels, height, width). Preprocessed image to be passed into the model as input.
• pixel_mask (optional): Mask to avoid performing attention on padding pixel values.
• mask_labels (optional): List of mask label tensors of shape (num_labels, height, width) to be fed to a model.
• class_labels (optional): List of target class label tensors of shape (num_labels, height, width) to be fed to the model. They identify the labels of mask_labels, e.g., the label of mask_labels[i][j] if class_labels[i][j].

Note: While segmentation datasets more or less have the same structure and data fields, different models such as MaskFormer and Segformer have different approaches to frame the segmentation problem. For example, MaskFormer creates binary instance segmentation maps (mask_labels) for each object instance across the training set, which requires additional preprocessing. Luckily, our preprocessor (an instance of the MaskFormerImageProcessor class) takes care of all these custom preprocessing steps.

However, the ImageProcessor classes don’t support data augmentation (e.g., image flipping, random crops, hue changes, etc.). To illustrate how you can combine custom data augmentation steps with the existing image processor classes of transformers, we will create a DataLoader class that uses both the MaskFormerImageProcessor and the Albumentations library. So, let’s go ahead and define our data loader class.

# Define the configurations of the transforms specific
# to the dataset used
ADE_MEAN = np.array([123.675, 116.280, 103.530]) / 255
ADE_STD = np.array([58.395, 57.120, 57.375]) / 255

# Build the augmentation transforms
train_val_transform = A.Compose([
    A.Resize(width=512, height=512),
    A.HorizontalFlip(p=0.3),
    A.Normalize(mean=ADE_MEAN, std=ADE_STD),
])

class ImageSegmentationDataset(Dataset):
    def __init__(self, dataset, processor, transform=None):
        # Initialize the dataset, processor, and transform variables
        self.dataset = dataset
        self.processor = processor
        self.transform = transform
        
    def __len__(self):
        # Return the number of datapoints
        return len(self.dataset)
    
    def __getitem__(self, idx):
        # Convert the PIL Image to a NumPy array
        image = np.array(self.dataset[idx]["image"].convert("RGB"))
        
        # Get the pixel wise instance id and category id maps
        # of shape (height, width)
        instance_seg = np.array(self.dataset[idx]["annotation"])[..., 1]
        class_id_map = np.array(self.dataset[idx]["annotation"])[..., 0]
        class_labels = np.unique(class_id_map)

        # Build the instance to class dictionary
        inst2class = {}
        for label in class_labels:
            instance_ids = np.unique(instance_seg[class_id_map == label])
            inst2class.update({i: label for i in instance_ids})

        # Apply transforms
        if self.transform is not None:
            transformed = self.transform(image=image, mask=instance_seg)
            (image, instance_seg) = (transformed["image"], transformed["mask"])
            
            # Convert from channels last to channels first
            image = image.transpose(2,0,1)

        if class_labels.shape[0] == 1 and class_labels[0] == 0:
            # If the image has no objects then it is skipped
            inputs = self.processor([image], return_tensors="pt")
            inputs = {k:v.squeeze() for k,v in inputs.items()}
            inputs["class_labels"] = torch.tensor([0])
            inputs["mask_labels"] = torch.zeros(
                (0, inputs["pixel_values"].shape[-2], inputs["pixel_values"].shape[-1])
            )
        else:
            # Else use process the image with the segmentation maps
            inputs = self.processor(
                [image],
                [instance_seg],
                instance_id_to_semantic_id=inst2class,
                return_tensors="pt"
            )
            inputs = {
                k:v.squeeze() if isinstance(v, torch.Tensor) else v[0] for k,v in inputs.items()
            }

        # Return the inputs
        return inputs

# Build the train and validation instance segmentation dataset
train_dataset = ImageSegmentationDataset(
    train,
    processor=processor,
    transform=train_val_transform
)
val_dataset = ImageSegmentationDataset(
    validation,
    processor=processor,
    transform=train_val_transform
)

Before moving on, let’s check if the preprocessed data is in the expected format:

# Check if everything is preprocessed correctly
inputs = val_dataset[0]
for k,v in inputs.items():
  print(k, v.shape)

Everything seems to be in the correct order! We can now create a collate function to load the data in batches and move on to model training.

def collate_fn(examples):
    # Get the pixel values, pixel mask, mask labels, and class labels
    pixel_values = torch.stack([example["pixel_values"] for example in examples])
    pixel_mask = torch.stack([example["pixel_mask"] for example in examples])
    mask_labels = [example["mask_labels"] for example in examples]
    class_labels = [example["class_labels"] for example in examples]

    # Return a dictionary of all the collated features
    return {
        "pixel_values": pixel_values,
        "pixel_mask": pixel_mask,
        "mask_labels": mask_labels,
        "class_labels": class_labels
    }

# Building the training and validation dataloader
train_dataloader = DataLoader(
    train_dataset,
    batch_size=4,
    shuffle=True,
    collate_fn=collate_fn
)
val_dataloader = DataLoader(
    val_dataset,
    batch_size=4,
    shuffle=False,
    collate_fn=collate_fn
)

Training Step

We will now fine-tune our model for 2 epochs for demonstration purposes. If we need the fine-tuned model, we will download it later for inference and evaluation, which is trained for 25 epochs. Ideally, we would fine-tune our model for >100 epochs for optimal performance.

# Use GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Initialize Adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)

# Set number of epochs and batch size
num_epochs = 2

for epoch in range(num_epochs):
    print(f"Epoch {epoch} | Training")

    # Set model in training mode 
    model.train()
    train_loss, val_loss = [], []

    # Training loop
    for idx, batch in enumerate(tqdm(train_dataloader)):
        # Reset the parameter gradients
        optimizer.zero_grad()
 
        # Forward pass
        outputs = model(
            pixel_values=batch["pixel_values"].to(device),
            mask_labels=[labels.to(device) for labels in batch["mask_labels"]],
            class_labels=[labels.to(device) for labels in batch["class_labels"]],
        )

        # Backward propagation
        loss = outputs.loss
        train_loss.append(loss.item())
        loss.backward()

        if idx % 50 == 0:
            print("  Training loss: ", round(sum(train_loss)/len(train_loss), 6))

        # Optimization
        optimizer.step()

    # Average train epoch loss
    train_loss = sum(train_loss)/len(train_loss)

    # Set model in evaluation mode
    model.eval()

    start_idx = 0
    print(f"Epoch {epoch} | Validation")
    for idx, batch in enumerate(tqdm(val_dataloader)):
        with torch.no_grad():
            # Forward pass
            outputs = model(
                pixel_values=batch["pixel_values"].to(device),
                mask_labels=[labels.to(device) for labels in batch["mask_labels"]],
                class_labels=[labels.to(device) for labels in batch["class_labels"]],
            )

            # Get validation loss
            loss = outputs.loss
            val_loss.append(loss.item())

            if idx % 50 == 0:
                print("  Validation loss: ", round(sum(val_loss)/len(val_loss), 6))

    # Average validation epoch loss
    val_loss = sum(val_loss)/len(val_loss)

    # Print epoch losses
    print(f"Epoch {epoch} | train_loss: {train_loss} | validation_loss: {val_loss}")

Saving the Trained Model

Luckily for us, 🤗 Transformers make it really easy to save trained models and custom preprocessing pipelines. We can simply call the save_pretrained() method as follows:

model.save_pretrained(<Path to target save folder>)
preprocessor.save_pretrained(<Path to target save folder>)

Alternatively, we can push our trained model to the hub for private use or share it publicly. To do this, you will need to create an account on the Hub and get an authentication token. For more information, check out the tutorial on sharing your models.

from huggingface_hub import notebook_login

# Login to your account
notebook_login()

# Push your model and preprocessor to the Hub
model.push_to_hub("my-awesome-model")
processor.push_to_hub("my-awesome-model")

Using the push_to_hub() method will create a repository named <YOUR_USERNAME>/my-awesome-model on the Hub. Whether you save your model locally or push it to the hub, saving the model will create a folder identical to any MaskFormer repo on the hub, see an example over here.

Let’s go ahead and push our fine-tuned model to the hub:

# Login to your account
notebook_login()

# We won't be using albumentations to preprocess images for inference
processor.do_normalize = True
processor.do_resize = True
processor.do_rescale = True

# Push your model and preprocessor to the Hub
model.push_to_hub("maskformer-swin-base-sceneparse-instance")
processor.push_to_hub("maskformer-swin-base-sceneparse-instance")

Inference and Evaluation

We can now load our trained model with a single line of code and perform inference. Remember, we can load a pre-trained model using the from_pretrained(<MODEL NAME OR PATH>) method. Since we pushed our trained model to the hub, we can load it using the repository name we defined earlier. Then, to perform inference, we just need to pass in the pixel_values (preprocessed input image) to the model.

# Use GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Grab the trained model and processor from the hub
model = MaskFormerForInstanceSegmentation.from_pretrained(
    "adirik/maskformer-swin-base-sceneparse-instance"
).to(device)
processor = MaskFormerImageProcessor.from_pretrained(
    "adirik/maskformer-swin-base-sceneparse-instance"
)

Let’s test our model on a test image:

# Use random test image
index = random.randint(0, len(test))
image = test[index]["image"].convert("RGB")
target_size = image.size[::-1]

# Preprocess image
inputs = processor(images=image, return_tensors="pt").to(device)

# Inference
model.eval()
with torch.no_grad():
    outputs = model(**inputs)

Next, let’s see what information is returned by our model:

# Let's print the items returned by our model and their shapes
print("Outputs...")
for key, value in outputs.items():
    print(f"  {key}: {value.shape}")

Post-Processing Outputs

As we can see, MaskFormer returns class query logits of shape (batch_size, num_queries, num_labels + 1) and mask query logits of shape (batch_size, num_queries, height, width). Therefore, to retrieve the most probable instance segmentation masks and their classes and to rescale them back to the original image/s shape, we can use the convenient post_process_instance_segmentation() method of the MaskFormerImageProcessor class.

Note that this class also has post_process_semantic_segmentation() and post_process_panoptic_segmentation() methods to post-process the outputs of MaskFormer models trained on semantic segmentation and panoptic segmentation datasets. If in doubt, refer to the model card of the checkpoint you would like to use to find out which dataset your selected model was trained on.

Let’s post-process our segmentation predictions and see what this method outputs:

# Post-process results to retrieve instance segmentation maps
result = processor.post_process_instance_segmentation(
    outputs,
    threshold=0.5,
    target_sizes=[target_size]
)[0] # we pass a single output therefore we take the first result (single)

instance_seg_mask = result["segmentation"].cpu().detach().numpy()

print(f"Final mask shape: {instance_seg_mask.shape}")
print("Segments Information...")
for info in result["segments_info"]:
    print(f"  {info}")

The MaskFormerImageProcessor.post_process_instance_segmentation() returns a list of length batch_size where each entry is a dictionary and contains two keys:

• segmentation: A tensor of shape (height, width) where each pixel represents an instance id or List[List] run-length encoding (RLE) of the segmentation map if return_coco_annotation is set to True or a concatenated tensor of binary segmentation maps (one per detected instance) of shape (num_instances, height, width) if return_binary_maps is set to True.
• segments_info: A dictionary that contains additional information on each segment.
  • id: An integer representing the instance id.
  • label_id: An integer representing the semantic class id corresponding to id.
  • score: Prediction score of a segment with id.

Let’s go ahead and visualize the post-processed per-pixel instance segmentation map.

def visualize_instance_seg_mask(mask):
    # Initialize image with zeros with the image resolution
    # of the segmentation mask and 3 channels
    image = np.zeros((mask.shape[0], mask.shape[1], 3))

    # Create labels
    labels = np.unique(mask)
    label2color = {
        label: (
            random.randint(0, 255),
            random.randint(0, 255),
            random.randint(0, 255),
        )
        for label in labels
    }

    for height in range(image.shape[0]):
        for width in range(image.shape[1]):
            image[height, width, :] = label2color[mask[height, width]]

    image = image / 255
    return image

instance_seg_mask_disp = visualize_instance_seg_mask(instance_seg_mask)

plt.figure(figsize=(10, 10))
for plot_index in range(2):
    if plot_index == 0:
        plot_image = image
        title = "Original"
    else:
        plot_image = instance_seg_mask_disp
        title = "Segmentation"
    
    plt.subplot(1, 2, plot_index+1)
    plt.imshow(plot_image)
    plt.title(title)
    plt.axis("off")

Performance Evaluation

The result looks pretty good, but we still need to quantitatively benchmark our model to see how well it performs. Since our test set does not have any segmentation maps we can use, we will use our validation set, and the Evaluate library to benchmark predicted segmentation maps against ground truth maps.

The panoptic and instance segmentation models are typically evaluated using the Panoptic Quality (PQ) and Average Precision (AP) metrics, respectively. To simplify our workflow, we can post-process our predictions to get the predicted semantic segmentation maps instead of the instance segmentation maps and use the simpler Mean Intersection over Union (IoU) metric to benchmark our fine-tuned model. Load the mean_iou metric from the evaluate library to compute the mean intersection over the union between the prediction and ground truth sets.

# Load Mean IoU metric
metrics = evaluate.load("mean_iou")

# Set model in evaluation mode
model.eval()

# Test set doesn't have annotations so we will use the validation set
ground_truths, preds = [], []

for idx in tqdm(range(200)):
    image = validation[idx]["image"].convert("RGB")
    target_size = image.size[::-1]

    # Get ground truth semantic segmentation map
    annotation = np.array(validation[idx]["annotation"])[:,:,0]
    # Replace null class (0) with the ignore_index (255) and reduce labels
    annotation -= 1
    annotation[annotation==-1] = 255
    ground_truths.append(annotation)

    # Preprocess image
    inputs = processor(images=image, return_tensors="pt").to(device)

    # Inference
    model.eval()
    with torch.no_grad():
        outputs = model(**inputs)
 
    # Post-process results to retrieve semantic segmentation maps
    result = processor.post_process_semantic_segmentation(outputs, target_sizes=[target_size])[0]
    semantic_seg_mask = result.cpu().detach().numpy()
    preds.append(semantic_seg_mask)

results = metrics.compute(
    predictions=preds,
    references=ground_truths,
    num_labels=100,
    ignore_index=255
)
print(f"Mean IoU: {results['mean_iou']} | Mean Accuracy: {results['mean_accuracy']} | Overall Accuracy: {results['overall_accuracy']}")

What's next? I recommend PyImageSearch University.

Course information:
69 total classes • 73 hours of on-demand code walkthrough videos • Last updated: February 2023
★★★★★ 4.84 (128 Ratings) • 15,800+ 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:

• ✓ 69 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 69 Certificates of Completion
• ✓ 73 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 500+ 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

Summary

Hohoo, that’s it. We fine-tuned a pre-trained MaskFormer model on the Scene Parsing dataset for instance segmentation, and got pretty good results after training for only a few epochs. We uploaded our model to the Hugging Face hub so that we could download and use it with a few lines of code. We shared our model publicly, but we could also set it as private for personal use or commercialization purposes.

On top of this, we learned how to load datasets with the Datasets library, preprocess datasets with transformers ImageProcessor classes, and combine them with other image processing libraries such as Albumentations for data augmentation during training. We also learned the transformers library’s different components and how they are organized to make it easier to create custom configurations and preprocessing pipelines to fine-tune pre-trained models. Finally, we used the Evaluate library to benchmark the performance of our fine-tuned model.

Tweet @pyimagesearch OR email ask.me@pyimagesearch.com

Citation Information

Dirik, A. “Train a MaskFormer Segmentation Model with Hugging Face Transformers,” PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/6b7z8

@incollection{Dirik_2023_HFTMF,
  author = {Alara Dirik},
  title = {Train a MaskFormer Segmentation Model with Hugging Face Transformers},
  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/6b7z8},
}

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!

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!

Download the code!

Website

The post Train a MaskFormer Segmentation Model with Hugging Face Transformers appeared first on PyImageSearch.

Table of Contents

Triplet Loss with Keras and TensorFlow

In today’s tutorial, we will try to understand the formulation of the triplet loss and build our Siamese Network Model in Keras and TensorFlow, which will be used to develop our Face Recognition application.

In the previous tutorial of this series, we built the dataset and data pipeline for our Siamese Network based Face Recognition application. Specifically, we looked at an overview of triplet loss and discussed what kind of data samples are required to train our model with the triplet loss. We also used Keras and TensorFlow to develop modules that allow us to process our dataset and generate triplet data samples that can be used to train our model.

Furthermore, we discussed and implemented face detection and cropping, which form an important part of the data pipeline and allow our face recognition model to effectively make predictions based on facial features.

In this tutorial, we will dive deeper into the definition of the triplet and discuss its mathematical formulation in detail. Furthermore, we will build our Siamese Network model and write our own triplet loss function, which will form the basis for our face recognition application and later be used to train our face recognition application.

This lesson is the 3rd of a 5-part series on Siamese Networks and their application in face recognition:

1. Face Recognition with Siamese Networks, Keras, and TensorFlow
2. Building a Dataset for Triplet Loss with Keras and TensorFlow
3. Triplet Loss with Keras and TensorFlow (this tutorial)
4. Training and Making Predictions with Siamese Networks and Triplet Loss
5. Evaluating Siamese Network Accuracy (ROC, Precision, and Recall) with Keras and TensorFlow

To learn how to write your own triplet loss with Keras and TensorFlow, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Triplet Loss with Keras and TensorFlow

In the first part of this series, we discussed the basic formulation of a contrastive loss and how it can be used to learn a distance measure based on similarity. Furthermore, in the previous tutorial, we looked at the type of data samples required to train our model with triplet loss and discussed the anchor, positive, and negative data samples.

The triplet loss uses these samples to learn an embedding space where samples from the same person or class lie close to each other and samples from different classes or persons lie farther apart. It is based on the idea that the distance between the representation of the anchor and negative samples (which belong to different people) should be at least a “margin” distance more than the distance between the anchor and positive samples. Here, “margin” is a hyperparameter similar to the one we discussed in the definition of pairwise contrastive loss in the first part of this series.

Let us try to delve deeper and understand this with an example. Figure 1 shows a typical pipeline for training a Siamese Network based Face Recognition model with Triplet Loss.

Suppose we have a database of faces, and we draw a triplet sample consisting of an anchor image , positive image , and negative image . Remember, as discussed in the previous tutorial, that the anchor and positive images are different instances of faces of the same person, and the anchor and negative images come from different people. The three images are passed through our deep network (denoted as Deep Architecture in the figure) to get the corresponding representations in the embedding space, that is, , , and . Usually, each of these representations is then normalized to have a unit norm. In other words, this means that they will lie on a hypersphere with a radius of 1 in the embedding space since each has a norm equal to 1 (as shown in Figure 1).

Let us now discuss how we use the triplet loss to train this network and understand its mathematical formulation.

The equation above shows the formulation of the triplet loss we will use to train our Siamese Network. Notice that here refers to the square of the L2 norm of , which is simply the square of the Euclidean distance between our representations and . Let us refer to this as .

Similarly, is the square of the Euclidean distance between our representations and , which we refer to as . Furthermore, the parameter refers to the margin distance.

After simplification, our final equation for triplet loss is as follows.

Notice that the minimum value of this expression is 0, and it occurs when the term . Rearranging the terms, we notice that this implies that this loss is minimized when

In other words, optimizing our model to minimize the triplet loss ensures that the distance between our anchor and negative representations is at least margin = higher than the distance between our anchor and positive representations. This allows us to learn an embedding space where our anchor and positive representations are close and the anchor and negative representations are farther apart.

Now that we have discussed the mathematical formulation of the triplet loss, let us look at the internal structure of the deep network (referred to as Deep Architecture in Figure 1) that we will be using for our face recognition application.

Figure 2 shows the internal structure of the Deep Architecture block, which we refer to as the Embedding Module in the code for this tutorial. The function of the Embedding Module is to take our anchor, positive, and negative samples in the image space and map them to d-dimensional representations in the embedding space, as shown.

The embedding module is composed of a backbone network and a trainable module. This tutorial uses the ResNet-50 (without the final head) as the backbone feature extractor. Note that the backbone is initialized with ImageNet weights, and its weights are frozen. Furthermore, we additionally use a lightweight trainable module composed of fully connected layers to get the final d-dimensional representation for our inputs. Notice that this setup allows us to use the feature extraction capabilities of the ResNet-50 backbone without having to train such a big network. Furthermore, the trainable module gives our embedding module the flexibility to learn the appropriate function for the task without training numerous parameters.

Now that we have discussed an overview of our Siamese Model pipeline, the formulation of the triplet loss, and the internal structure of our modules, let us code our Siamese Model with Keras and TensorFlow.

Configuring Your Development Environment

To follow this guide, you need to have the TensorFlow and OpenCV libraries installed on your system.

Luckily, both TensorFlow and OpenCV are pip-installable:

$ pip install tensorflow
$ pip install opencv-contrib-python

If you need help configuring your development environment for OpenCV, we highly recommend that you read our pip install OpenCV guide — it will have you up and running in a matter of minutes.

Having Problems Configuring Your Development Environment?

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 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

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.

├── crop_faces.py
├── face_crop_model
│   ├── deploy.prototxt.txt
│   └── res10_300x300_ssd_iter_140000.caffemodel
├── inference.py
├── pyimagesearch
│   ├── config.py
│   ├── dataset.py
│   └── model.py
└── train.py

In the previous tutorial, we discussed the directory structure for our Siamese Network based Face Recognition application in detail. Furthermore, we presented a step-by-step walkthrough of our config.py, dataset.py, and crop_faces.py files, which allow us to process input data and build the data pipeline.

In this tutorial, we will discuss the model.py file from the pyimagesearch folder, which implements the code for our Siamese Network Model and triplet loss function.

Implementing Siamese Model and Triplet Loss

Now that we have discussed the concepts required to build our Siamese Model, let’s dive into the code and implement the triplet loss and Siamese network in Keras and TensorFlow.

We open our model.py file from the pyimagesearch folder and get started.

# import the necessary packages
from tensorflow.keras.applications import resnet
from tensorflow.keras import layers
from tensorflow import keras
import tensorflow as tf

def get_embedding_module(imageSize):
    # construct the input layer and pass the inputs through a
    # pre-processing layer
    inputs = keras.Input(imageSize + (3,))
    x = resnet.preprocess_input(inputs)
    
    # fetch the pre-trained resnet 50 model and freeze the weights
    baseCnn = resnet.ResNet50(weights="imagenet", include_top=False)
    baseCnn.trainable=False
    
    # pass the pre-processed inputs through the base cnn and get the
    # extracted features from the inputs
    extractedFeatures = baseCnn(x)

    # pass the extracted features through a number of trainable layers
    x = layers.GlobalAveragePooling2D()(extractedFeatures)
    x = layers.Dense(units=1024, activation="relu")(x)
    x = layers.Dropout(0.2)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dense(units=512, activation="relu")(x)
    x = layers.Dropout(0.2)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dense(units=256, activation="relu")(x)
    x = layers.Dropout(0.2)(x)
    outputs = layers.Dense(units=128)(x)

    # build the embedding model and return it
    embedding = keras.Model(inputs, outputs, name="embedding")
    return embedding

We start by importing the necessary packages and modules to build our Siamese Model on Lines 2-5.

On Lines 7-35, we develop the embedding module of our Siamese model that we discussed above. This module takes as input the imageSize (i.e., height and width of the input image) as shown on Line 7.

We start by defining the input layer using the keras.Input function, as shown on Line 10. This layer takes as input the imageSize (i.e., its height and width) along the number of channels (i.e., (3,)) of our input image. Then on Line 11, we pass our inputs through the preprocessing layer of our backbone reset network.

On Line 14, we initialize our ResNet50 backbone model and load the pre-trained ImageNet weights using the weights="imagenet" argument. Furthermore, we use include_top=False since we only need the backbone network without the final fully connected layers. On Line 15, we set baseCnn.trainable=False to freeze the weights of our resnet backbone.

On Line 19, we pass our input x through our baseCnn and store the output features in extractedFeatures as shown. Next, we pass the extracted features through our trainable layers on Lines 22-30. We first consolidate the extractedFeatures using the layers.GlobalAveragePooling2D() layer (Line 22) and then pass it through a 1024-dimensional fully connected layer with relu activation using layers.Dense(units=1024, activation="relu") (Line 23). Next, we use the Dropout regularization layer on Line 24. Then, we use a series of BatchNormalization → Dense → Dropout layers (Lines 25-30) and finally pass our processed output x through the final 128-dimensional fully connected layer on Line 31.

Finally, we build the embedding model using the keras.Model function with our computed inputs and outputs as the input and output to our model (Line 34). We then return our embedding model on Line 35.

def get_siamese_network(imageSize, embeddingModel):
    # build the anchor, positive and negative input layer
    anchorInput = keras.Input(name="anchor", shape=imageSize + (3,))
    positiveInput = keras.Input(name="positive", shape=imageSize + (3,))
    negativeInput = keras.Input(name="negative", shape=imageSize + (3,))

    # embed the anchor, positive and negative images
    anchorEmbedding = embeddingModel(anchorInput)
    positiveEmbedding = embeddingModel(positiveInput)
    negativeEmbedding = embeddingModel(negativeInput)

    # build the siamese network and return it
    siamese_network = keras.Model(
        inputs=[anchorInput, positiveInput, negativeInput],
        outputs=[anchorEmbedding, positiveEmbedding, negativeEmbedding]
    )
    return siamese_network

Now we build our get_siamese_network function, which takes as input imageSize and embeddingModel and returns our siamese_network (Lines 37-53).

On Lines 39-41, we build an input layer for each anchor, positive, and negative samples using the keras.Input layer. Next, on Lines 44-46, we use the embeddingModel to embed our anchorInput, positiveInput, and negativeInput.

Finally, on Lines 49-52, we build our Siamese network model with [anchorInput, positiveInput, negativeInput] as the input and [anchorEmbedding, positiveEmbedding, negativeEmbedding] as the output. We return our siamese_network on Line 53.

class SiameseModel(keras.Model):
    def __init__(self, siameseNetwork, margin, lossTracker):
        super().__init__()
        self.siameseNetwork = siameseNetwork
        self.margin = margin
        self.lossTracker = lossTracker

    def _compute_distance(self, inputs):
        (anchor, positive, negative) = inputs
        # embed the images using the siamese network
        embeddings = self.siameseNetwork((anchor, positive, negative))
        anchorEmbedding = embeddings[0]
        positiveEmbedding = embeddings[1]
        negativeEmbedding = embeddings[2]

        # calculate the anchor to positive and negative distance
        apDistance = tf.reduce_sum(
            tf.square(anchorEmbedding - positiveEmbedding), axis=-1
        )
        anDistance = tf.reduce_sum(
            tf.square(anchorEmbedding - negativeEmbedding), axis=-1
        )
        
        # return the distances
        return (apDistance, anDistance)

    def _compute_loss(self, apDistance, anDistance):
        loss = apDistance - anDistance
        loss = tf.maximum(loss + self.margin, 0.0)
        return loss

    def call(self, inputs):
        # compute the distance between the anchor and positive,
        # negative images
        (apDistance, anDistance) = self._compute_distance(inputs)
        return (apDistance, anDistance)

    def train_step(self, inputs):
        with tf.GradientTape() as tape:
            # compute the distance between the anchor and positive,
            # negative images
            (apDistance, anDistance) = self._compute_distance(inputs)

            # calculate the loss of the siamese network
            loss = self._compute_loss(apDistance, anDistance)

        # compute the gradients and optimize the model
        gradients = tape.gradient(
            loss,
            self.siameseNetwork.trainable_variables)
        self.optimizer.apply_gradients(
            zip(gradients, self.siameseNetwork.trainable_variables)
        )

        # update the metrics and return the loss
        self.lossTracker.update_state(loss)
        return {"loss": self.lossTracker.result()}

    def test_step(self, inputs):
        # compute the distance between the anchor and positive,
        # negative images
        (apDistance, anDistance) = self._compute_distance(inputs)

        # calculate the loss of the siamese network
        loss = self._compute_loss(apDistance, anDistance)
        
        # update the metrics and return the loss
        self.lossTracker.update_state(loss)
        return {"loss": self.lossTracker.result()}

    @property
    def metrics(self):
        return [self.lossTracker]

Now, we define our SiameseModel class which will allow us to implement and apply the triplet loss function during the training and test steps of our Face Recognition Application (Lines 55-127).

We start by defining the _init_ function, which takes as arguments the siameseNetwork, the margin, and the lossTracker (which allows us to log and track the losses during training), as shown on Line 56. Then, on Lines 58-60, we assign these arguments to create the self.siameseNetwork, self.margin, and self.lossTracker attributes of the class.

Now we will discuss the _compute_distance function (Lines 62-79), which takes the inputs (i.e., (anchor, positive, negative) samples) and computes the distance between the anchor and positive sample (i.e., apDistance) and the anchor and negative sample (i.e., anDistance).

We start by unpacking our inputs (i.e., anchor, positive, and negative samples) on Line 63. Next, we embed them by passing them through our self.siameseNetwork and storing the outputs in embeddings, as shown on Line 65.

On Lines 66-69, we unpack the embeddings of our anchor sample (i.e., anchorEmbedding), positive sample (i.e., positiveEmbedding), and negative sample (i.e., negativeEmbedding).

Next, we compute the square of the Euclidean distance between the anchor and positive embedding (i.e., apDistance) on Lines 71 and 72. Similarly, we compute the distance between the anchor and negative embedding (i.e., anDistance) on Lines 74-76). Finally, we return the distances (i.e., (apDistance, anDistance)) on Line 79.

Now that we have computed the distances between our embeddings, let us write a function to compute our triplet loss. We define the _compute_loss function (Lines 81-84), which takes as input the distance between the anchor and positive (i.e., apDistance) and anchor and negative embeddings (i.e., anDistance).

On Line 81, we compute the difference between apDistance and anDistance, and on Line 83, we implement our triplet loss equation which we discussed before.

Finally, we return our computed loss on Line 84.

On Lines 86-90, we define our call function that takes the inputs and uses the self._compute_distance function to compute apDistance and anDistance (Line 89) and finally returns these distances on Line 90.

Finally, let us discuss how the triplet loss is applied during our model’s training and test phase.

We discuss the train_step function defined on Lines 92-111, which takes the inputs as arguments, as shown on Line 92. During the training phase, we want to keep track of the gradients during loss computation to use them during backpropagation. Thus, we work inside with tf.GradientTape() as tape: block, as shown on Line 93.

We first compute the distance between the anchor and positive (i.e., apDistance) and anchor and negative embeddings (i.e., anDistance) using the self._compute_distance function on Line 96. Then we use the self._compute_loss function defined above to compute our triplet loss and store it as loss as shown on Line 99.

Next, on Lines 102-104, we compute the gradients of our loss w.r.t. the trainable parameters of our Siamese Network (i.e., self.siameseNetwork.trainable_variables) as shown. We then update the weights to optimize the model using self.optimizer.apply_gradients, which takes the gradients and the corresponding trainable variables as shown on Lines 105-107. Finally, we use the computed loss to update the loss metric and return the loss on Lines 110 and 111.

Now that we have discussed the training step, let us further discuss the test_step function (Lines 113-123). Note that this function will be very similar to the train_step function, with the only difference being that the gradients will not be tracked or computed, and there will be loss optimization since we are in the test or inference phase.

Similar to the train_step function, the test_step function also takes the inputs as the argument, as shown on Line 113. Next, we compute the apDistance and anDistance on Line 116 and compute the triplet loss on Line 119 as we did previously in the training step. Finally, we use the computed loss to update the loss metric and return the loss on Lines 122 and 123.

Ultimately, we define our metrics function that simply returns our self.lossTracker, which keeps track of our metrics.

This completes the code walkthrough to build our Siamese model and triplet loss function.

What's next? I recommend PyImageSearch University.

Course information:
69 total classes • 73 hours of on-demand code walkthrough videos • Last updated: February 2023
★★★★★ 4.84 (128 Ratings) • 15,800+ 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:

• ✓ 69 courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 69 Certificates of Completion
• ✓ 73 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 500+ 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

Summary

In this tutorial, we discussed building a Siamese Network based Face Recognition Model using Keras and TensorFlow. Specifically, we tried to understand the details of the triplet loss and closely looked at a typical pipeline for training our Siamese Network Model with Triplet Loss.

Furthermore, we discussed the mathematical formulation of the triplet loss. We saw how it enforces the anchor and positive representations to be closer to each other and the anchor and negative representations to be farther apart. Additionally, we discussed the internal structure of the embedding module, which allows us to efficiently map our faces from the image space to the embedding space.

In the upcoming tutorials of this series, we will see how we can train our Siamese Network and make predictions using it in real-time.

Credits

This post is inspired by the amazing National Programme on Technology Enhanced Learning (NPTEL) Deep Learning for Computer Vision Course, which the author contributed to while working at IIT Hyderabad.

Citation Information

Chandhok, S. “Triplet Loss with Keras and TensorFlow, PyImageSearch, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2023, https://pyimg.co/2try0

@incollection{Chandhok_2023_TLwK+TF,
  author = {Shivam Chandhok},
  title = {Triplet Loss with {Keras and TensorFlow}},
  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/2try0},
}

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!

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!

Download the code!

Website

The post Triplet Loss with Keras and TensorFlow appeared first on PyImageSearch.

Table of Contents

Learning JAX in 2023: Part 2 — JAX’s Power Tools grad, jit, vmap, and pmap

In this tutorial, you will learn the power tools of JAX, grad, jit, vmap, and pmap.

This lesson is the 2nd in a 3-part series on Learning JAX in 2023:

1. Learning JAX in 2023: Part 1 — The Ultimate Guide to Accelerating Numerical Computation and Machine Learning
2. Learning JAX in 2023: Part 2 — JAX’s Power Tools grad, jit, vmap, and pmap (today’s tutorial)
3. Learning JAX in 2023: Part 3 — A Step-by-Step Guide to Training Your First Machine Learning Model with JAX

To learn how to use JAX’s power tools, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Learning JAX in 2023: Part 2 — JAX’s Power Tools grad, jit, vmap, and pmap

🙌🏻 Introduction

Welcome to our comprehensive guide on advanced JAX techniques! In the previous tutorial, we were introduced to JAX, and its predecessors autograd and xla. We also briefly looked into numerical computing with JAX.

In this post, we’ll be diving into some of the most powerful and useful features of the JAX library, including grad, jit, vmap, and pmap. These functions allow you to easily and efficiently compute gradients of functions, optimize your code for faster execution, and apply functions to arrays of data in parallel. By the end of this post, you’ll have a solid understanding of how to use these tools to improve the performance and functionality of your numerical computation and machine learning tasks.

We’ll also cover the topic of randomness in JAX, including how to generate and control random numbers for use in your computations. Randomness is an important aspect of many machine learning algorithms, and JAX provides a range of functions and techniques for working with randomness in a controlled and reproducible manner.

Whether you’re a seasoned JAX user or just getting started with the library, we hope you’ll find this post a valuable resource for improving your skills and taking your projects to the next level. So let’s get started!

Configuring Your Development Environment

To follow this guide, you need to have the JAX library installed on your system.

Luckily, JAX is pip-installable:

$ pip install numpy
$ pip install autograd
$ pip install jax

Having Problems Configuring Your Development Environment?