Image Classification Using Pre-Trained Model

PyTorch - Image Classification Using Pre-Trained Model

In this lesson, you will learn how to use a pre-trained deep learning model from torchvision to classify objects in an image. We will use SqueezeNet 1.1 (similar to the original tutorial) or ResNet, which are trained on the large ImageNet dataset.

We will cover the steps: loading an image, preprocessing it to match the model’s input requirements, loading the pre-trained model, performing inference, and interpreting the results.

Importing Libraries

First, import the necessary Python packages:

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image # Pillow library for image loading
import requests
import numpy as np
import json
import os

Ensure you have Pillow installed (pip install Pillow) and requests (pip install requests).

Setting up Model and Preprocessing

Pre-trained models expect input images to be processed in a specific way (e.g., resized, cropped, normalized). torchvision.transforms provides tools for this.

# Load a pre-trained model (e.g., SqueezeNet 1.1 or ResNet18)
# model = models.squeezenet1_1(pretrained=True)
model = models.resnet18(pretrained=True)
model.eval() # Set the model to evaluation mode (disables dropout, etc.)

# Define the image transformations
# Models trained on ImageNet usually expect 224x224 images
# Normalization values are standard for ImageNet
preprocess = transforms.Compose([
    transforms.Resize(256),         # Resize smaller edge to 256
    transforms.CenterCrop(224),     # Crop 224x224 from the center
    transforms.ToTensor(),          # Convert image to PyTorch Tensor (scales to [0, 1])
    transforms.Normalize(           # Normalize with ImageNet mean and std
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

print("Model and preprocessing pipeline ready.")

Loading and Processing an Image

Let’s load an example image. You can download one or use a local file. We’ll use a local file named image.jpg placed in the same directory as the script (or provide the full path).

image_path = 'image.jpg' # Make sure this image exists

if not os.path.exists(image_path):
    print(f"Error: Image file not found at {image_path}")
    print("Please download an image and save it as image.jpg or update the path.")
    # Example download (requires internet)
    # try:
    #     url = "https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg"
    #     img_data = requests.get(url).content
    #     with open(image_path, 'wb') as handler:
    #         handler.write(img_data)
    #     print(f"Downloaded example image to {image_path}")
    # except Exception as e:
    #     print(f"Could not download example image: {e}")
    #     exit()
else:
    print(f"Loading image from {image_path}")

try:
    input_image = Image.open(image_path).convert('RGB') # Ensure image is RGB
except FileNotFoundError:
    print(f"Error: Cannot open image file at {image_path}")
    exit()

# Apply the preprocessing transformations
input_tensor = preprocess(input_image)

# Create a mini-batch (add batch dimension)
# Models expect input shape [batch_size, channels, height, width]
input_batch = input_tensor.unsqueeze(0)

print("Image loaded and preprocessed.")
print("Input tensor shape:", input_batch.shape) # Should be [1, 3, 224, 224]

Performing Inference

Now, feed the processed image tensor to the model. We use torch.no_grad() because we are not training, which saves memory and computation.

with torch.no_grad(): # Disables gradient calculation
    output = model(input_batch)

# The output contains raw scores (logits) for each class
print("Inference complete.")
print("Output tensor shape:", output.shape) # Shape [1, 1000] for ImageNet models

Interpreting the Results

The model output is a tensor of raw scores (logits) for 1000 ImageNet classes. To get probabilities, we apply the Softmax function. Then, we find the class with the highest probability.

# Apply Softmax to get probabilities
probabilities = torch.nn.functional.softmax(output[0], dim=0)

# Get the top prediction
top1_prob, top1_idx = torch.topk(probabilities, 1)

predicted_idx = top1_idx[0].item()
confidence = top1_prob[0].item()

print(f"\nPredicted Index: {predicted_idx}")
print(f"Confidence: {confidence:.4f}")

The index (e.g., 207 for ‘golden retriever’ in ImageNet) isn’t very informative. We need to map this index to a human-readable class name.

Mapping Index to Class Name

We can download the standard ImageNet class index mapping.

# Download the ImageNet class index mapping
imagenet_classes_url = "https://raw.githubusercontent.com/pytorch/examples/main/imagenet/imagenet_classes.txt"
# Alternative: Use a cached JSON mapping if available
# imagenet_classes_url = "https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json"

try:
    response = requests.get(imagenet_classes_url)
    response.raise_for_status() # Raise an exception for bad status codes
    # Simple text format: each line is a class name
    labels = response.text.split('\n')
    # Remove empty lines if any
    labels = [label for label in labels if label]
    if predicted_idx < len(labels):
        predicted_class_name = labels[predicted_idx]
        print(f"Predicted Class Name: {predicted_class_name}")
    else:
        print(f"Error: Predicted index {predicted_idx} is out of bounds for the loaded labels (length {len(labels)}). Using index directly.")
        predicted_class_name = f"Class Index {predicted_idx}"

except requests.exceptions.RequestException as e:
    print(f"Could not download or process class labels: {e}")
    predicted_class_name = f"Class Index {predicted_idx}"

print(f"\nFinal Prediction: {predicted_class_name} (Confidence: {confidence:.4f})")

Now you should see the actual name of the predicted object in your image, along with the model’s confidence.

You have successfully used a pre-trained PyTorch model for image classification. You can easily adapt this code to use different models from torchvision.models or classify other images by changing the image_path.