How to load images for inference in batch

(Bibhu) #1

Hey there I am new to PyTorch. I have a inference code that predicts and classify images. I can predict and classify images one by one, can anyone please help me to classify all the images of a folder in a batch.

Directory structure:
images
… img1.jpg
… img2.jpg
… img3.jpg

Output:
elephant
lion
tiger

How can I load all the image in the folder and predict one by one.
I am using the prediction code as follows:

def predict_image(image_path):
    print("prediciton in progress")
    image = Image.open(image_path)

    transformation = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    image_tensor = transformation(image).float()
    image_tensor = image_tensor.unsqueeze_(0)

    if cuda:
        image_tensor.cuda()

    input = Variable(image_tensor)
    output = model(input)

    index = output.data.numpy().argmax()
    return index

This works for single images, if called again and again the time of execution will increase. Also I am doing inference on a CPU machine.

1 Like
(Amrit Das) #2

As per my understanding, I wrote this piece of code. To load images and predict.

data_transforms = {
    'predict': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    }

dataset = {'predict' : datasets.ImageFolder("./data", data_transforms['predict'])}
dataloader = {'predict': torch.utils.data.DataLoader(dataset['predict'], batch_size = 1, shuffle=False, num_workers=4)}

outputs = list()
since = time.time()
for inputs, labels in dataloader['predict']:
    inputs = inputs.to(device)
    output = model(inputs)
    output = output.to(device)
    index = output.data.numpy().argmax()
    print index

I hope this helps you and solves your problem :smile:

1 Like
(Bibhu) #3

Thank You :slight_smile: worked like a charm

(Ben Bowles) #4

What is a device? torch.cuda.current_device() ?

(Manuel Alejandro Diaz Zapata) #5

Usually a device is where you do your computations, for example if you use a GPU your device will be something like cuda:0 or just 0.

So what torch.cuda.current_device() does is to return the identifier of which GPU is currently being used. This can be really helpful for systems with multiple GPUs.

Docs for more information.