Title: PyTorch image classification with pre-trained networks - PyImageSearch

URL Source: https://pyimagesearch.com/2021/07/26/pytorch-image-classification-with-pre-trained-networks/

Published Time: 2021-07-26T14:00:00+00:00

Markdown Content:
In this tutorial, you will learn how to perform image classification with pre-trained networks using PyTorch. Utilizing these networks, you can accurately classify 1,000 common object categories in only a few lines of code.

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

Today’s tutorial is part four in our five part series on PyTorch fundamentals:

1.  [_What is PyTorch?_](https://pyimagesearch.com/2021/07/05/what-is-pytorch/)
2.  [_Intro to PyTorch: Training your first neural network using PyTorch_](https://pyimagesearch.com/2021/07/12/intro-to-pytorch-training-your-first-neural-network-using-pytorch/)
3.  [_PyTorch: Training your first Convolutional Neural Network_](https://pyimagesearch.com/2021/07/19/pytorch-training-your-first-convolutional-neural-network-cnn/)
4.  _PyTorch image classification with pre-trained networks_ (today’s tutorial)
5.  _August 2nd: PyTorch object detection with pre-trained networks_ (next week’s tutorial)

Throughout the rest of this tutorial, you’ll gain experience using PyTorch to classify input images using seminal, state-of-the-art image classification networks, including VGG, Inception, DenseNet, and ResNet.

**To learn how to perform image classification with pre-trained PyTorch networks,** **_just keep reading._**

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

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

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

In the first part of this tutorial, we’ll discuss what pre-trained image classification networks are, including those that are built into the PyTorch library.

From there, we’ll configure our development environment and review our project directory structure.

I’ll then show you how to implement a Python script that can accurately classify input images using pre-trained PyTorch networks.

We’ll wrap up this tutorial with a discussion of our results.

### **What are pre-trained image classification networks?**

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

**Figure 1:** Most popular, state-of-the-art neural networks come with weights pre-trained on the ImageNet dataset. The PyTorch library includes many of these popular image classification networks.

When it comes to image classification, there is no dataset/challenge more famous than [ImageNet](https://www.image-net.org/). **The goal of ImageNet is to accurately classify input images into a set of 1,000 common object categories that computer vision systems will “see” in everyday life.**

Most popular deep learning frameworks, including PyTorch, Keras, TensorFlow, fast.ai, and others, include _pre-trained_ networks. These are highly accurate, state-of-the-art models that computer vision researchers trained on the ImageNet dataset.

After training on ImageNet was complete, researchers saved their models to disk and then published them freely for other researchers, students, and developers to learn from and use in their own projects.

This tutorial will show how to use PyTorch to classify input images using the following state-of-the-art classification networks:

*   VGG16
*   VGG19
*   Inception
*   DenseNet
*   ResNet

Let’s get started!

### **Configuring your development environment**

To follow this guide, you need to have both PyTorch and OpenCV installed on your system.

Luckily, both PyTorch and OpenCV are extremely easy to install using pip:

$ pip install torch torchvision
$ pip install opencv-contrib-python

**If you need help configuring your development environment for PyTorch, I _highly recommend_ that you** [**read the PyTorch documentation**](https://pytorch.org/get-started/locally/) — PyTorch’s documentation is comprehensive and will have you up and running quickly.

And if you need help installing OpenCV, [be sure to refer to my _pip install OpenCV_ tutorial](https://pyimagesearch.com/2018/09/19/pip-install-opencv/).

### **Having problems configuring your development environment?**

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

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

All that said, are you:

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

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

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

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

### **Project structure**

Before we implement image classification with PyTorch, let’s first review our project directory structure.

Start by accessing the **_“Downloads”_** section of this guide to retrieve the source code and example images. You’ll then be presented with the following directory structure.

$ tree . --dirsfirst
.
├── images
│   ├── bmw.png
│   ├── boat.png
│   ├── clint\_eastwood.jpg
│   ├── jemma.png
│   ├── office.png
│   ├── scotch.png
│   ├── soccer\_ball.jpg
│   └── tv.png
├── pyimagesearch
│   └── config.py
├── classify\_image.py
└── ilsvrc2012\_wordnet\_lemmas.txt

Inside the `pyimagesearch` module we have a single file, `config.py`. This file stores important configurations, such as:

*   Our input image dimensions
*   Mean and standard deviation for mean subtraction and scaling
*   Whether or not we are using a GPU for training
*   Path to the human-readable ImageNet class labels (i.e., `ilsvrc2012_wordnet_lemmas.txt`)

Our `classify_image.py` script will load our `config` and then classify an input image using either VGG16, VGG19, Inception, DenseNet, or ResNet (depending on which model architecture we supply as our command line argument).

The `images` directory contains a number of sample images where we’ll apply these image classification networks.

### **Creating our configuration file**

Before we implement our image classification driver script, let’s first create a configuration file to store important configurations.

Open the `config.py` file in the `pyimagesearch` module and insert the following code:

\# import the necessary packages
import torch

# specify image dimension
IMAGE\_SIZE = 224

# specify ImageNet mean and standard deviation
MEAN = \[0.485, 0.456, 0.406\]
STD = \[0.229, 0.224, 0.225\]

# determine the device we will be using for inference
DEVICE = "cuda" if torch.cuda.is\_available() else "cpu"

# specify path to the ImageNet labels
IN\_LABELS = "ilsvrc2012\_wordnet\_lemmas.txt"

**Line 5** defines our input image spatial dimensions, meaning that each image will be resized to _224×224_ pixels before being passed through our pre-trained PyTorch network for classification.

**_Note:_** _Most networks trained on the ImageNet dataset accept images that are 224×224 or 227×227. Some networks, particularly fully convolutional networks, may accept larger image dimensions._

From there, we define the mean and standard deviation of RGB pixel intensities across our training set (**Lines 8 and 9**). Prior to passing an input image through our network for classification, we first scale the image pixel intensities by subtracting the mean and then dividing by the standard deviation — this preprocessing is typical for CNNs trained on large, diverse image datasets such as ImageNet.

From there, **Line 12** specifies whether we are using our CPU or GPU for training, while **Line 15** defines the path to our input text file of ImageNet class labels.

If you were to open this file in your favorite text editor of choice, you would see the following contents:

tench, Tinca\_tinca
goldfish, Carassius\_auratus
...
bolete
ear, spike, capitulum
toilet\_tissue, toilet\_paper, bathroom\_tissue

Each row in this text file maps to the name of a class label our pre-trained PyTorch networks were trained to recognize and classify.

### **Implementing our image classification script**

With our configuration file taken care of, let’s move on to implementing our main driver script used to classify input images using our pre-trained PyTorch networks.

Open the `classify_image.py` file in your project directory structure, and let’s get to work:

\# import the necessary packages
from pyimagesearch import config
from torchvision import models
import numpy as np
import argparse
import torch
import cv2

We start on **Lines 2-7** importing our Python packages, including:

*   `config`: The configuration file we implemented from the previous section
*   `models`: Contains PyTorch’s pre-trained neural networks
*   `numpy`: Numerical array processing
*   `torch`: Accesses the PyTorch API
*   `cv2`: Our OpenCV bindings

With our imports taken care of, let’s define a function to accept an input image and preprocess it:

def preprocess\_image(image):
	# swap the color channels from BGR to RGB, resize it, and scale
	# the pixel values to \[0, 1\] range
	image = cv2.cvtColor(image, cv2.COLOR\_BGR2RGB)
	image = cv2.resize(image, (config.IMAGE\_SIZE, config.IMAGE\_SIZE))
	image = image.astype("float32") / 255.0

	# subtract ImageNet mean, divide by ImageNet standard deviation,
	# set "channels first" ordering, and add a batch dimension
	image -= config.MEAN
	image /= config.STD
	image = np.transpose(image, (2, 0, 1))
	image = np.expand\_dims(image, 0)

	# return the preprocessed image
	return image

Our `preprocess_image` function takes a single argument, `image`, which is the image we’ll be preprocessing for classification.

We start the preprocessing operations by:

1.  Swapping from BGR to RGB channel ordering (the pre-trained networks we’re using here utilized RGB channel ordering whereas OpenCV uses BGR ordering by default)
2.  Resizing our image to fixed dimensions (i.e., _224×224_), ignoring aspect ratio
3.  Converting our image to a floating point data type and then scaling the pixel intensities to the range _\[0, 1\]_

From there, we perform a second set of preprocessing operations:

1.  Subtracting the mean (**Line 18**) and dividing by the standard deviation (**Line 19**)
2.  Moving the channels dimension to the _front_ of the array (**Line 20**), which is called **channels-first ordering** and is the default channel ordering method that PyTorch expects
3.  Adding a batch dimension to the array (**Line 21**)

The preprocessed `image` is then returned to the calling function.

Next, let’s parse our command line arguments:

\# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add\_argument("-i", "--image", required=True,
	help="path to the input image")
ap.add\_argument("-m", "--model", type=str, default="vgg16",
	choices=\["vgg16", "vgg19", "inception", "densenet", "resnet"\],
	help="name of pre-trained network to use")
args = vars(ap.parse\_args())

We have two command line arguments to parse:

1.  `--image`: The path to the input image that we wish to classify
2.  `--model`: The pre-trained CNN model we’ll be using to classify the image

Let’s now define a `MODELS` dictionary which maps the name of the `--model` command line argument to its corresponding PyTorch function:

\# define a dictionary that maps model names to their classes
# inside torchvision
MODELS = {
	"vgg16": models.vgg16(pretrained=True),
	"vgg19": models.vgg19(pretrained=True),
	"inception": models.inception\_v3(pretrained=True),
	"densenet": models.densenet121(pretrained=True),
	"resnet": models.resnet50(pretrained=True)
}

# load our the network weights from disk, flash it to the current
# device, and set it to evaluation mode
print("\[INFO\] loading {}...".format(args\["model"\]))
model = MODELS\[args\["model"\]\].to(config.DEVICE)
model.eval()

**Lines 37-43** create our `MODELS` dictionary:

*   The _key_ to the dictionary is the human-readable name of the model, passed in via the `--model` command line argument.
*   The _value_ to the dictionary is the corresponding PyTorch function used to load the model with the weights pre-trained on ImageNet

You’ll be able to use the following pre-trained models to classify an input image with PyTorch:

1.  VGG16
2.  VGG19
3.  Inception
4.  DenseNet
5.  ResNet

Specifying the `pretrained=True` flag instructs PyTorch to not only load the model architecture definition, but _also_ download the pre-trained ImageNet weights for the model.

**Line 48** then loads the model and pre-trained weights (if you’ve never downloaded the model weights before they will be automatically downloaded and cached for you) and then sets the model to run either on your CPU or GPU, depending on your `DEVICE` from the configuration file.

**Line 49** puts our `model` into evaluation mode, instructing PyTorch to handle special layers, such as dropout and batch normalization, different from how it would otherwise handle them during training. **Putting your model into evaluation mode _before making predictions is critical_, so don’t forget to do it!**

Now that our model is loaded, we need an input image — let’s take care of that now:

\# load the image from disk, clone it (so we can draw on it later),
# and preprocess it
print("\[INFO\] loading image...")
image = cv2.imread(args\["image"\])
orig = image.copy()
image = preprocess\_image(image)

# convert the preprocessed image to a torch tensor and flash it to
# the current device
image = torch.from\_numpy(image)
image = image.to(config.DEVICE)

# load the preprocessed the ImageNet labels
print("\[INFO\] loading ImageNet labels...")
imagenetLabels = dict(enumerate(open(config.IN\_LABELS)))

**Line 54** loads our input `image` from disk. We make a copy of it on **Line 55** so that we can draw on it and visualize the top prediction of our network. We also make use of our `preprocess_image` function on **Line 56** to perform resizing and scaling.

**Line 60** converts our `image` from a NumPy array to a PyTorch tensor, while **Line 61** moves the `image` to our device (either CPU or GPU).

FInally, **Line 65** loads our input ImageNet class labels from disk.

We are now ready to make predictions on input `image` using our `model`:

\# classify the image and extract the predictions
print("\[INFO\] classifying image with '{}'...".format(args\["model"\]))
logits = model(image)
probabilities = torch.nn.Softmax(dim=-1)(logits)
sortedProba = torch.argsort(probabilities, dim=-1, descending=True)

# loop over the predictions and display the rank-5 predictions and
# corresponding probabilities to our terminal
for (i, idx) in enumerate(sortedProba\[0, :5\]):
	print("{}. {}: {:.2f}%".format
		(i, imagenetLabels\[idx.item()\].strip(),
		probabilities\[0, idx.item()\] \* 100))

**Line 69** performs a forward-pass of our network, resulting in the outputs of the network.

We pass these through the `Softmax` function on **Line 70** to obtain the predicted probabilities for each of the possible 1,000 class labels the `model` was trained on.

**Line 71** then sorts the probabilities in descending order with higher probabilities at the _front_ of the list.

We then display the top-5 predicted class labels and corresponding probabilities to our terminal on **Lines 75-78** by:

*   Looping over the top-5 predictions
*   Looking up the name of the class label using our `imagenetLabels` dictionary
*   Displaying the predicted probability

Our final code block draws the top-1 (i.e., top predicted label) on our output image:

\# draw the top prediction on the image and display the image to
# our screen
(label, prob) = (imagenetLabels\[probabilities.argmax().item()\],
	probabilities.max().item())
cv2.putText(orig, "Label: {}, {:.2f}%".format(label.strip(), prob \* 100),
	(10, 30), cv2.FONT\_HERSHEY\_SIMPLEX, 0.8, (0, 0, 255), 2)
cv2.imshow("Classification", orig)
cv2.waitKey(0)

The result is then displayed to our screen.

### **Image classification with PyTorch results**

We are now ready to apply image classification with PyTorch!

Be sure to access the **_“Downloads”_** section of this tutorial to retrieve the source code and example images.

From there, try classifying an input image using the following command:

$ python classify\_image.py --image images/boat.png
\[INFO\] loading vgg16...
\[INFO\] loading image...
\[INFO\] loading ImageNet labels...
\[INFO\] classifying image with 'vgg16'...
0. wreck: 99.99%
1. seashore, coast, seacoast, sea-coast: 0.01%
2. pirate, pirate\_ship: 0.00%
3. breakwater, groin, groyne, mole, bulwark, seawall, jetty: 0.00%
4. sea\_lion: 0.00%

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

**Figure 3:** Using PyTorch and VGG16 to classify an input image.

It appears that Captain Jack Sparrow is stranded on the beach! And sure enough, the VGG16 network is able to correctly classify the input image as a “wreck” (i.e., shipwreck) with 99.99% probability.

It’s also interesting to see that “seashore” is the second top prediction from the model — this prediction is also accurate, due to the boat being on the beach.

Let’s try a different image, this time using the DenseNet model:

$ python classify\_image.py --image images/bmw.png --model densenet
\[INFO\] loading densenet...
\[INFO\] loading image...
\[INFO\] loading ImageNet labels...
\[INFO\] classifying image with 'densenet'...
0. convertible: 96.61%
1. sports\_car, sport\_car: 2.25%
2. car\_wheel: 0.45%
3. beach\_wagon, station\_wagon, wagon, estate\_car, beach\_waggon, station\_waggon, waggon: 0.22%
4. racer, race\_car, racing\_car: 0.13%

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

**Figure 4:** Applying DenseNet and PyTorch to classify an image.

The top prediction from DenseNet is “convertible” with 96.61% accuracy. The second top prediction, “sports car” is also accurate.

This image contains Jemma, my family’s beagle:

$ python classify\_image.py --image images/jemma.png --model resnet
\[INFO\] loading resnet...
\[INFO\] loading image...
\[INFO\] loading ImageNet labels...
\[INFO\] classifying image with 'resnet'...
0. beagle: 95.98%
1. bluetick: 1.46%
2. Walker\_hound, Walker\_foxhound: 1.11%
3. English\_foxhound: 0.45%
4. maraca: 0.25%

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

**Figure 5:** Utilizing ResNet and PyTorch to correctly classify an input image.

Here we are using the ResNet architecture to classify our input image. Jemma is a “beagle” (a type of dog), which ResNet accurately predicts with 95.98% probability.

Interestingly, a “bluetick,” “walker hound,” and “English foxhound” are all types of dogs belonging to the “hound” family — all of these would be reasonable predictions from the model.

Let’s take a look at one final example:

$ python classify\_image.py --image images/soccer\_ball.jpg --model inception
\[INFO\] loading inception...
\[INFO\] loading image...
\[INFO\] loading ImageNet labels...
\[INFO\] classifying image with 'inception'...
0. soccer\_ball: 100.00%
1. volleyball: 0.00%
2. sea\_urchin: 0.00%
3. rugby\_ball: 0.00%
4. silky\_terrier, Sydney\_silky: 0.00%

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

**Figure 6:** Using Inception and PyTorch to make predictions on an input image.

Our Inception model correctly classifies the input image as “soccer ball” with 100% probability.

Image classification allows us to assign one or more labels to an input image; however, it tells us _nothing_ about _where in the image_ the object resides.

To determine _where_ in an input image a given object is, we need to apply **object detection:**

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

**Figure 7:** Object detection can not only tell us _what_ is in an image but also _where_ the object is.

Just like we have pre-trained networks for image classification, we also have pre-trained networks for object detection as well. Next week you’ll learn how to use PyTorch to detect objects in images using specialized object detection networks.

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

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

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

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

That’s _not_ the case.

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

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

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

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

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

**Summary**
-----------

In this tutorial, you learned how to perform image classification using PyTorch. Specifically, we utilized popular pre-trained network architectures, including:

*   VGG16
*   VGG19
*   Inception
*   DenseNet
*   ResNet

These models were trained by the researchers responsible for inventing and proposing the novel architectures listed above. After training was complete, these researchers saved the model weights to disk and then published them for other researchers, students, and developers to learn from and use in their own projects.

While the models are free to use, make sure you check any terms/conditions associated with them, as some models are not free to use in commercial applications (typically entrepreneurs in the AI space get around this restriction by training the models themselves rather than using the pre-trained weights provided by the original authors).

Stay tuned for next week’s blog post, where you’ll learn how to perform object detection using PyTorch.

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

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

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

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