File size: 2,158 Bytes
d218927 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import pytorch_lightning as pl
from torchvision.datasets import VOCSegmentation
from torchvision.transforms import transforms
from torch.utils.data import DataLoader, random_split
class SegmentationDataModule(pl.LightningDataModule):
def __init__(self, data_dir: str, config):
super().__init__()
self.config = config
# Transformación para la imagen y la máscara
image_transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor()
])
mask_transform = transforms.Compose([
transforms.Resize((128, 128)), # Puedes añadir más transformaciones si lo deseas
transforms.ToTensor(),
lambda x: x.long()
])
self.transform = DualTransform(image_transform, mask_transform)
self.data_dir = data_dir
def prepare_data(self):
# Descargar el dataset (si es necesario)
VOCSegmentation(root=self.data_dir, year='2012', image_set='trainval', download=False)
def setup(self, stage=None):
# Inicializa el dataset
self.dataset = VOCSegmentation(root=self.data_dir, year='2012', image_set='trainval', transforms=self.transform)
# Dividir el dataset y asignar a sets de entrenamiento/validación
train_len = int(0.8 * len(self.dataset))
val_len = len(self.dataset) - train_len
self.train_dataset, self.val_dataset = random_split(self.dataset, [train_len, val_len])
def train_dataloader(self):
return DataLoader(self.train_dataset, batch_size=self.config.batch_size, shuffle=self.config.shuffle, num_workers=self.config.num_workers)
def val_dataloader(self):
return DataLoader(self.val_dataset, batch_size=self.config.batch_size, shuffle=False, num_workers=self.config.num_workers)
# Puedes añadir un test_dataloader si lo deseas
class DualTransform:
def __init__(self, image_transform, mask_transform):
self.image_transform = image_transform
self.mask_transform = mask_transform
def __call__(self, image, mask):
return self.image_transform(image), self.mask_transform(mask)
|