How to modify ImageFolder in pytorch to return a tensor of a different shape?

Question

How to extend torch.datasets.ImageFolder in pytorch to return a tensor of a different shape?

It currently returns: torch.Size([1, 3, 256, 256]). I want to return [1, 10, 3, 256, 256].

I have a directory with multiple images separated into folders. Each folder has up to 3000 images. I would like to modify the getitem function so that it returns bags of images, where each bag contains 10 images.

Thank you!

Answer 1

A possible option may be to split your dataset into files of 10 images per file and then in your __getitem__(self,idx) method you can iterate 10 images at a time using the file that corresponds to idx, concatenate them and return that concatenated tensor. so for example (and make your own adjustments based on your init, etc..) Given a directory with this form:

- all_images
  - images_0
     - im_0
     - im_1
     - ...
     - im_9
  - images_1
     - ...
  - ...
  - images_n

then

def __init__(self,file="all_images"):
    self.images_file = file
        

def __getitem__(self,idx):
    ret_tensor = torch.tensor([])
    images = [image for image in os.listdir(f"{self.images_file}/images_{idx}")]
    for image in images:
        ret_tensor = torch.cat((ret_tensor,torch.load(image)),1)
    return ret_tensor