Image Classification in 5 simple steps. Using pre-trained DenseNet model in PyTorch

download.py

#[STEP-2] : Downloading an example Image and imagenet output classes

# Download an example image from the pytorch website
!wget "https://github.com/pytorch/hub/raw/master/images/dog.jpg" -O "dog.jpg"

# Download ImageNet labels
!wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt

# Read the image classes
with open("imagenet_classes.txt", "r") as f:
    img_classes = [s.strip() for s in f.readlines()]

print("Total Image Classes :{}".format(len(img_classes)))

print(img_classes[:5]) #displaying first five classes

import.py

#[STEP-1] : Importing all the libraries
import torch                                  #importing torch library
from torchvision import transforms,models     #importing transforms for applying transformations
from PIL import Image                         #Image module helps to load image ,etc.

model.py

#[STEP-4] : Loading densenet model and changing the device

# model 
model = models.densenet121(pretrained=True)
model.eval()  # sets the model in evaluation mode 

# For more info about evaluation:
# https://stackoverflow.com/questions/60018578/what-does-model-eval-do-in-pytorch


# move the input and model to GPU for speed if available
if torch.cuda.is_available():
    input_batch = input_batch.to('cuda')
    model.to('cuda')

result.py

#[STEP-5] : Getting final result/output.
with torch.no_grad():
    output = model(input_batch)

# Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes
print(output.shape)

# The output has unnormalized scores. To get probabilities, you can run a softmax on it.
probabilities = torch.nn.functional.softmax(output[0], dim=0)

# Show top categories per image
k = 5
top5_prob, top5_catid = torch.topk(probabilities, k)

for id,prob in zip(top5_catid,top5_prob):
    print(f'Image class: {img_classes[id]}\t----- {prob * 100:.4f} %')

transform.py

#[STEP-3] : Applying transformations to image

# transforming input image

filename = "dog.jpg"
input_image = Image.open(filename)

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model