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!
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