use custom dataset and augmentation

[ramdhan1989]

hi there! is there any easy way to run for custom dataset and augmentation? if they are not supported yet, kindly need your advise which part of your code that need to be modified?

thank you

[tieckit]

This is my code for using custom dataset.

utils.py

@attr.s(auto_attribs=True, slots=True)
class FashionDataset(DatasetBase):
    transform_train: Callable[[Any], torch.Tensor] = imagenet_default_transform
    transform_test: Callable[[Any], torch.Tensor] = imagenet_default_transform

    def configure_train(self):
        assert os.path.exists(self.data_path)
        return CustomDataset(self.data_path, split="train", transform=self.transform_train)

    def configure_validation(self):
        assert os.path.exists(self.data_path)
        return CustomDataset(self.data_path, split="val", transform=self.transform_test)

custom.py (new file)

class CustomDataset(Dataset):
    def __init__(self, csv_file, split="train", transform=None):
        self.data = pd.read_csv(csv_file)
        self.data = self.data[self.data["split"] == split]
        if split == "val":
            self.data = self.data.sample(frac=1, random_state=42)
        self.transform = transform
        self.image_paths = self.data["image_path"].to_numpy()
        self.labels = self.data["label"].to_numpy()
        self.image2ram = True
        if self.image2ram:
            self.images = []
            for img_path in tqdm(self.image_paths, desc="Loading images to RAM", total=len(self.image_paths)):
                image = Image.open(img_path).convert("RGB")
                self.images += [image]

        del self.data

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

    def __getitem__(self, idx):
        # Open image file
        if self.image2ram:
            image = self.images[idx]
        else:
            img_path = self.image_paths[idx]
            image = Image.open(img_path).convert("RGB")

        # Get label
        label = torch.tensor(self.labels[idx], dtype=torch.long)

        # Apply transforms if any
        if self.transform:
            image = self.transform(image)

        return image, label