Working with image data from memory in Pytorch

Question

Question:

If I have an array in memory with dims (n, height, width, channels) and I want to get a Pytorch classifier to feed them forward and give me an array with class predictions for each of the n images in the array, how do I do that?

Background:

I am working with a computer vision problem where I modify some images using pre-existing code and want to send the modified images into a Pytorch Classifier CNN (not developed or controlled by me). I am accustomed to Tensorflow/Keras more than Pytorch.

With Tensorflow/Keras models you can give them a bunch of images in a numpy array and it'll go ahead and feed them forward through the model.

PS:

A colleague suggested saving all the images to disk first, then reading them in with DataLoader but that is so unnecessary when I already have the images in memory.

Sorry if it's a dumb question, I tried to find a solution elsewhere but obviously haven't had much success.

Answer 1

You can create a custom DataLoader function which takes the images in memory and returns tensors which can be fed directly to the model without having to save them on disk first. A very simple implementation can be:

def images_to_tensor(images):
    #images is numpy array of shape N,H,W,C
    #normalizes images between -1 and 1, comment this if you want to normalize images between 0 and 1
    images = (images.astype(np.float32) - 127.5)/128.0
    #to normalize image from 0 to 1 uncomment the line below
    #images = (images.astype(np.float32))/255.0
    #changes numpy array to tensors
    images = torch.from_numpy(images).permute(0, 3, 1, 2)
    #to convert cpu tensors to cuda uncomment the line below
    #images = images.to("cuda")
    return images

You can then use the function to convert your images to tensors and pass them to the classification model to get the output predictions.