Pytorch: Add information to images in image prediction

Question

I would like to add information to my current dataset. At the moment, I have six-frame sequences in folders. The DataLoader reads all 6 and uses the first 3 for predicting the last 1/2/3 (depending on how many I tell him to). This is the function for the DataLoader.

class TrainFeeder(Dataset):
    def init(self, data_set):
        super(TrainFeeder, self).init()
        self.input_data = data_set
        #print(torch.cuda.current_device())
        if torch.cuda.current_device() ==0:
            print('There are total %d sequences in trainset' % len(self.input_data))

    def getitem(self, index):
        path = self.input_data[index]
        imgs_path = sorted(glob.glob(path + '/*.png'))
        imgs = []
        for img_path in imgs_path:
            img = cv2.imread(img_path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (256,448))
            img = cv2.resize(img, (0, 0), fx=0.5, fy=0.5, interpolation=cv2.INTER_CUBIC) #has been 0.5 for official data, new is fx = 2.63 and fy = 2.84
            img_tensor = ToTensor()(img).float()
            imgs.append(img_tensor)
        imgs = torch.stack(imgs, dim=0)
        return imgs

    def len(self):
        return len(self.input_data)

Now I'd like to add one value to these images. It is a boolean, I have stored in a list in a.json in the same folder, like the six-frame-sequences. But I don't know how to add the values of the list in the.json to the tensor. Which dimension should I use? Will the system work at all, if I change the shape of the input?

Answer 1

The function getitem can return anything, so you can return a tuple instead of just images:

def __getitem__(self, index):
    path = ...
    # load your 6 images
    imgs = torch.stack( ... )
    # load your boolean metadata
    metadata = load_json_data( ... )
    # return them both
    return (imgs, metadata)

You will need to make metadata a tensor before returning it, otherwise I expect that pytorch will complain about not being able to collate (ie stack) them to make batches

"Will the system work" is a question only you can answer, since you did not provide the code of your ML model. I would bet on: "no but it won't require significant changes to work". Most likely you currently have a loop like

for imgs in dataloader:
    # do some training 
    output = model(imgs)
    ...

And you will have to make it like

for imgs, metadata in dataloader:
     # do some training
     output = model(imgs)
     ...