text_plain_output_1.png
from datasets import load_dataset

from torch.utils.data import DataLoader

from transformers import ViTImageProcessor

import torch

from datasets import load_dataset
train_ds = load_dataset('miladfa7/5-Flower-Types-Classification-Dataset')

train_ds = train_ds['train'].train_test_split(test_size=0.15)
train_data = train_ds['train']
test_data = train_ds['test']

label = list(set(train_data['label']))
id2label = {id: label for id, label in enumerate(label)}
label2id = {label: id for id, label in id2label.items()}

from transformers import ViTImageProcessor
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')

from torchvision.transforms import CenterCrop, Compose, Normalize, RandomHorizontalFlip, RandomResizedCrop, Resize, ToTensor
image_mean, image_std = (processor.image_mean, processor.image_std)
size = processor.size['height']
normalize = Normalize(mean=image_mean, std=image_std)
_train_transforms = Compose([Resize((size, size)), RandomHorizontalFlip(), ToTensor(), normalize])
_val_transforms = Compose([Resize((size, size)), ToTensor(), normalize])
def train_transforms(examples):
    examples['pixel_values'] = [_train_transforms(image.convert('RGB')) for image in examples['image']]
    return examples
def val_transforms(examples):
    examples['pixel_values'] = [_val_transforms(image.convert('RGB')) for image in examples['image']]
    return examples

train_data.set_transform(train_transforms)
test_data.set_transform(val_transforms)

from torch.utils.data import DataLoader
import torch
def collate_fn(examples):
    pixel_values = torch.stack([example['pixel_values'] for example in examples])
    labels = torch.tensor([label2id[example['label']] for example in examples])
    return {'pixel_values': pixel_values, 'labels': labels}
train_dataloader = DataLoader(train_data, collate_fn=collate_fn, batch_size=4)
test_dataloader = DataLoader(test_data, collate_fn=collate_fn, batch_size=4)

batch = next(iter(train_dataloader))

batch = next(iter(test_dataloader))
for k, v in batch.items():
    if isinstance(v, torch.Tensor):
        print(k, v.shape)