albumentations-team / autoalbument

AutoML for image augmentation. AutoAlbument uses the Faster AutoAugment algorithm to find optimal augmentation policies. Documentation - https://albumentations.ai/docs/autoalbument/
https://albumentations.ai/docs/autoalbument/
MIT License
203 stars 20 forks source link

Example: code snipped for local image files #23

Open andife opened 3 years ago

andife commented 3 years ago

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 commented 3 years ago

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