| import os |
| import rasterio |
| import torch |
| from torchgeo.datasets import NonGeoDataset |
| from torch.utils.data import DataLoader |
| from torchgeo.datamodules import NonGeoDataModule |
|
|
| from methan_text_dataset import MethaneTextDataset |
|
|
| class MethaneTextDataModule(NonGeoDataModule): |
| """ |
| A DataModule for handling MethaneClassificationDataset |
| """ |
|
|
| def __init__( |
| self, |
| data_root: str, |
| paths: list, |
| captions: list, |
| batch_size: int = 8, |
| num_workers: int = 0, |
| train_transform: callable = None, |
| val_transform: callable = None, |
| test_transform: callable = None, |
| **kwargs |
| ): |
| super().__init__(MethaneTextDataset, batch_size, num_workers, **kwargs) |
|
|
| self.data_root = data_root |
| self.paths = paths |
| self.captions = captions |
| self.train_transform = train_transform |
| self.val_transform = val_transform |
| self.test_transform = test_transform |
|
|
| def setup(self, stage: str = None): |
| if stage in ("fit", "train"): |
| self.train_dataset = MethaneTextDataset( |
| root_dir=self.data_root, |
| paths=self.paths, |
| captions=self.captions, |
| transform=self.train_transform, |
| ) |
| if stage in ("fit", "validate", "val"): |
| self.val_dataset = MethaneTextDataset( |
| root_dir=self.data_root, |
| paths=self.paths, |
| captions=self.captions, |
| transform=self.val_transform, |
| ) |
| if stage in ("test", "predict"): |
| self.test_dataset = MethaneTextDataset( |
| root_dir=self.data_root, |
| paths=self.paths, |
| captions=self.captions, |
| transform=self.test_transform, |
| ) |
|
|
| def train_dataloader(self): |
| return DataLoader( |
| self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, drop_last=True |
| ) |
|
|
| def val_dataloader(self): |
| return DataLoader( |
| self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True |
| ) |
|
|
| def test_dataloader(self): |
| return DataLoader( |
| self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True |
| ) |
