| import torch |
| from torch import nn |
| import pytorch_lightning as pl |
| from torch.utils.data import DataLoader |
| from torchvision.transforms.functional import center_crop |
|
|
| from pathlib import Path |
| import numpy as np |
|
|
| from pyphoon2.DigitalTyphoonDataset import DigitalTyphoonDataset |
|
|
|
|
| class TyphoonDataModule(pl.LightningDataModule): |
| """Typhoon Dataset Module using lightning architecture""" |
| def __init__( |
| self, |
| dataroot, |
| batch_size, |
| num_workers, |
| split_by, |
| load_data, |
| dataset_split, |
| standardize_range, |
| downsample_size, |
| cropped, |
| corruption_ceiling_pct=100, |
| ): |
| super().__init__() |
|
|
| self.batch_size = batch_size |
| self.num_workers = num_workers |
|
|
| data_path = Path(dataroot) |
| self.images_path = str(data_path / "image") + "/" |
| self.track_path = str(data_path / "metadata") + "/" |
| self.metadata_path = str(data_path / "metadata.json") |
| self.load_data = load_data |
| self.split_by = split_by |
|
|
| self.dataset_split = dataset_split |
| self.standardize_range = standardize_range |
| self.downsample_size = downsample_size |
| self.cropped = cropped |
|
|
| self.corruption_ceiling_pct = corruption_ceiling_pct |
|
|
| def setup(self, stage): |
| |
| dataset = DigitalTyphoonDataset( |
| str(self.images_path), |
| str(self.track_path), |
| str(self.metadata_path), |
| "grade", |
| load_data_into_memory=self.load_data, |
| filter_func=self.image_filter, |
| transform_func=self.transform_func, |
| spectrum="Infrared", |
| verbose=False, |
| ) |
|
|
| self.train_set, self.val_set, _ = dataset.random_split( |
| self.dataset_split, split_by=self.split_by |
| ) |
|
|
| def train_dataloader(self): |
| return DataLoader( |
| self.train_set, |
| batch_size=self.batch_size, |
| num_workers=self.num_workers, |
| shuffle=True, |
| ) |
|
|
| def val_dataloader(self): |
| return DataLoader( |
| self.val_set, |
| batch_size=self.batch_size, |
| num_workers=self.num_workers, |
| shuffle=False, |
| ) |
|
|
| def image_filter(self, image): |
| return ( |
| (image.grade() < 7) |
| and (image.year() != 2023) |
| and (100.0 <= image.long() <= 180.0) |
| ) |
|
|
| def transform_func(self, image_batch): |
| """transform function applied on the images for pre-processing""" |
| image_batch = np.clip( |
| image_batch, self.standardize_range[0], self.standardize_range[1] |
| ) |
| image_batch = (image_batch - self.standardize_range[0]) / ( |
| self.standardize_range[1] - self.standardize_range[0] |
| ) |
| if self.downsample_size != (512, 512): |
| image_batch = torch.Tensor(image_batch) |
| if self.cropped: |
| image_batch = center_crop(image_batch, (224, 224)) |
| else: |
| image_batch = torch.reshape( |
| image_batch, [1, 1, image_batch.size()[0], image_batch.size()[1]] |
| ) |
| image_batch = nn.functional.interpolate( |
| image_batch, |
| size=self.downsample_size, |
| mode="bilinear", |
| align_corners=False, |
| ) |
| image_batch = torch.reshape( |
| image_batch, [image_batch.size()[2], image_batch.size()[3]] |
| ) |
| image_batch = image_batch.numpy() |
| return image_batch |
|
|
