datamodule/sevir_lightning_module.py

from typing import Tuple,Literal,Optional
import os
import datetime

import numpy as np
import torch
from torch.utils.data import random_split,DataLoader

from lightning import LightningDataModule,seed_everything

from .sevir_const import (
    SEVIR_ROOT_DIR,SEVIR_ROOT_LR_DIR,
    SEVIR_CATALOG,SEVIR_LR_CATALOG,SEVIR_LR_FILTERED_CATALOG,
    SEVIR_DATA_DIR,SEVIR_LR_DATA_DIR,
    SEVIR_RAW_SEQ_LEN,SEVIR_LR_RAW_SEQ_LEN,
    SEVIR_INTERVAL_REAL_TIME,SEVIR_LR_INTERVAL_REAL_TIME,
    SEVIR_H_W_SIZE,SEVIR_LR_H_W_SIZE
)
from .sevir_torch_dataset import SEVIRTorchDataset

class SEVIRLightningDataModule(LightningDataModule):

    def __init__(
        self,
        seq_len: int = 25,stride: int = 12,
        sample_mode: str = "sequent",
        layout: str = "NTHWC",
        
        output_type = np.float32,
        preprocess: bool = True,
        rescale_method: str = "01",
        verbose: bool = False,
        aug_mode: str = "0",
        ret_contiguous: bool = True,
        
        # region Datamodule Config
        dataset_name:Literal['sevir','sevirlr'] = "sevir",
        sevir_dir: str = None,
        start_date: Tuple[int] = None,
        train_test_split_date: Tuple[int] = (2019, 6, 1),
        end_date: Tuple[int] = None,
        val_ratio: float = 0.1,
        batch_size: int = 1,
        num_workers: int = 1,
        seed: int = 0,
        # endregion
    ):
        super(SEVIRLightningDataModule, self).__init__()
        
        # region Data Formatting
        self.seq_len = seq_len
        self.stride = stride
        self.sample_mode = sample_mode
        
        assert layout[0] == "N"
        self.layout = layout.replace("N", "")
        # endregion
        
        # region Operation Config
        self.output_type = output_type
        self.preprocess = preprocess
        self.rescale_method = rescale_method
        self.verbose = verbose
        self.aug_mode = aug_mode
        self.ret_contiguous = ret_contiguous
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.seed = seed
        # endregion
        
        self._config_sevir(sevir_dir,dataset_name)
        
        # train val test split
        self.start_date = datetime.datetime(*start_date) \
            if start_date is not None else None
        self.train_test_split_date = datetime.datetime(*train_test_split_date) \
            if train_test_split_date is not None else None
        self.end_date = datetime.datetime(*end_date) \
            if end_date is not None else None
        self.val_ratio = val_ratio
        
    def _config_sevir(self,sevir_dir:Optional[str], dataset_name:Literal['sevir','sevirlr']):
        assert dataset_name in ['sevir','sevirlr','filtered_sevirlr'], f"Unknown dataset configuration {dataset_name}"
        sevir_dir = os.path.abspath(sevir_dir) if sevir_dir is not None else None
        if dataset_name == "sevir":
            if sevir_dir is None:
                sevir_dir = SEVIR_ROOT_DIR
            catalog_path = SEVIR_CATALOG
            raw_data_dir = SEVIR_DATA_DIR
            raw_seq_len = SEVIR_RAW_SEQ_LEN
            interval_real_time = SEVIR_INTERVAL_REAL_TIME
            img_height,img_width = SEVIR_H_W_SIZE
        elif dataset_name == "sevirlr":
            if sevir_dir is None:
                sevir_dir = SEVIR_ROOT_LR_DIR
            catalog_path = SEVIR_LR_CATALOG
            raw_data_dir = SEVIR_LR_DATA_DIR
            raw_seq_len = SEVIR_LR_RAW_SEQ_LEN
            interval_real_time = SEVIR_LR_INTERVAL_REAL_TIME
            img_height,img_width = SEVIR_LR_H_W_SIZE
        elif dataset_name == "filtered_sevirlr":
            if sevir_dir is None:
                sevir_dir = SEVIR_ROOT_LR_DIR
            catalog_path = SEVIR_LR_FILTERED_CATALOG
            raw_data_dir = SEVIR_LR_DATA_DIR
            raw_seq_len = SEVIR_LR_RAW_SEQ_LEN
            interval_real_time = SEVIR_LR_INTERVAL_REAL_TIME
            img_height,img_width = SEVIR_LR_H_W_SIZE
        else:
            raise ValueError(f"Wrong dataset name {dataset_name}. Must be 'sevir' or 'sevirlr'.")
        self.dataset_name = dataset_name
        self.sevir_dir = sevir_dir
        self.catalog_path = catalog_path
        self.raw_data_dir = raw_data_dir
        self.raw_seq_len = raw_seq_len
        self.interval_real_time = interval_real_time
        self.img_height = img_height
        self.img_width = img_width

    def prepare_data(self) -> None:
        if os.path.exists(self.sevir_dir):
            # Further check
            assert os.path.exists(self.catalog_path), f"CATALOG.csv not found! Should be located at {self.catalog_path}"
            assert os.path.exists(self.raw_data_dir), f"SEVIR data not found! Should be located at {self.raw_data_dir}"
        else:
            raise NotImplementedError(f'Data not available in specified directory {self.sevir_dir}')

    def setup(self, stage = None) -> None:
        seed_everything(seed=self.seed)
        if stage in (None, "fit"):
            sevir_train_val = SEVIRTorchDataset(
                sevir_catalog=self.catalog_path,
                sevir_data_dir=self.raw_data_dir,
                raw_seq_len=self.raw_seq_len,
                split_mode="uneven",
                shuffle=True,
                seq_len=self.seq_len,
                stride=self.stride,
                sample_mode=self.sample_mode,
                layout=self.layout,
                start_date=self.start_date,
                end_date=self.train_test_split_date,
                output_type=self.output_type,
                preprocess=self.preprocess,
                rescale_method=self.rescale_method,
                verbose=self.verbose,
                aug_mode=self.aug_mode,
                ret_contiguous=self.ret_contiguous,
            )
            self.sevir_train, self.sevir_val = random_split(
                dataset=sevir_train_val,
                lengths=[1 - self.val_ratio, self.val_ratio],
                generator=torch.Generator().manual_seed(self.seed)
            )
        if stage in (None, "test"):
            self.sevir_test = SEVIRTorchDataset(
                sevir_catalog=self.catalog_path,
                sevir_data_dir=self.raw_data_dir,
                raw_seq_len=self.raw_seq_len,
                split_mode="uneven",
                shuffle=False,
                seq_len=self.seq_len,
                stride=self.stride,
                sample_mode=self.sample_mode,
                layout=self.layout,
                start_date=self.train_test_split_date,
                end_date=self.end_date,
                output_type=self.output_type,
                preprocess=self.preprocess,
                rescale_method=self.rescale_method,
                verbose=self.verbose,
                aug_mode="0",
                ret_contiguous=self.ret_contiguous,
            )

    def train_dataloader(self):
        return DataLoader(
            self.sevir_train,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers
        )

    def val_dataloader(self):
        return DataLoader(
            self.sevir_val,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers
        )

    def test_dataloader(self):
        return DataLoader(
            self.sevir_test,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers
        )

    @property
    def num_train_samples(self):
        return len(self.sevir_train)

    @property
    def num_val_samples(self):
        return len(self.sevir_val)

    @property
    def num_test_samples(self):
        return len(self.sevir_test)

if __name__ == "__main__":
    dm = SEVIRLightningDataModule(
        seq_len=22,
        sample_mode='sequent',
        stride=3,
        layout='NTHWC',
        aug_mode='2',
        ret_contiguous=False,
        dataset_name='sevirlr',
        start_date=None,
        train_test_split_date=[2019, 6, 1],
        end_date=None,
        val_ratio=0.1,
        num_workers=8
    )
    dm.prepare_data()
    dm.setup()
    for batch in dm.train_dataloader():
        print(batch[0].shape)
        break
    # print(dm.num_train_samples)
    # print(dm.num_val_samples)
    # print(dm.num_test_samples)