Example: code snipped for local image files

[andife]

Hello,
very interesting package.

Can someone provide a code snipped for a
"Add implementation for len and getitem methods in dataset.py."

My dataset is a classification problem, and the files are structured
like

Project/group1/*bmp
Project/group2/*bmp

Thank you

Andreas

[davidecarnevali]

The following dataset.py is working for me:

import torch.utils.data
import torchvision
from torchvision import dataset
import cv2
cv2.setNumThreads(0)

dataset = datasets.ImageFolder("/Project/")
images = dataset.imgs

class SearchDataset(torch.utils.data.Dataset):

def __init__(self, images_filepaths=images, transform=None):
    self.images_filepaths = images_filepaths
    self.transform = transform

def __len__(self):
    return len(self.images_filepaths)

def __getitem__(self, index):
    path, label = self.images_filepaths[index]
    image = cv2.imread(path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    if self.transform is not None:
        transformed = self.transform(image=image)
        image = transformed["image"]

    return image, label

Images should be all of the same size. If they are not, you can resize them by adding to __getitem__:
image = cv2.resize(image, (height, width))