I am doing RandomHorizontalFlip and RandomVerticalFlip, though they are being applied to the image and mask separately.
import numpy as np
import pandas as pd
import torch
from torch import nn
from torch import optim
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
from pathlib import Path
from PIL import Image
from torchvision import tv_tensors
from torchvision.transforms.v2 import Compose, RandomVerticalFlip, RandomHorizontalFlip,ToDtype
masks = np.load(‘path/to/masks’)
train_images = Path(‘traindirectory’)
train_imgs = [x for x in train_images.iterdir() if x.is_file()]
train_imgs = sorted(train_imgs)
train_imgs = [Image.open(x) for x in train_imgs]
for i in range(len(train_imgs)):
image = train_imgs[i]
image = image.transform(
image.size,
Image.AFFINE,
(1, 0, -30, 0, 1, -30),
fillcolor=(0, 0, 0)
)
train_imgs[i] = image
test_images = Path(‘testdirectory’)
test_imgs = [x for x in test_images.iterdir() if x.is_file()]
test_imgs = sorted(test_imgs)
test_imgs = [Image.open(x) for x in test_imgs]
class SegmentationDataset(Dataset):
def init(self, subset, images, transforms, masks=None):
super().init()
self.images = images
self.masks = masks
self.subset = subset
self.transforms = transforms
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
if self.subset == 'test':
image = self.images[idx]
return image
else:
image = tv_tensors.Image(self.images[idx])
mask = tv_tensors.Mask(torch.Tensor(self.masks[idx]))
out = self.transforms({ 'mask':mask, 'image': image})
return out
train_er = SegmentationDataset(‘train’, train_imgs, Compose([RandomVerticalFlip(), RandomHorizontalFlip() , ToDtype(torch.float32, scale=True)]), masks =masks )
img = train_er[0][‘image’].permute(1, 2, 0).numpy()
msk = train_er[0][‘mask’]
plt.imshow(img)
plt.imshow(msk, alpha=0.5, cmap=‘plasma’)
this is the output I get:
