datamodule.py

"""
LightningDataModule to load Earth Observation data from GeoTIFF files using
rasterio.
"""

from collections import defaultdict
from pathlib import Path
from typing import List, Literal

import lightning as L
import numpy as np
import torch
import torchdata
import yaml
from box import Box
from einops import rearrange
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.sampler import Sampler
from torchvision.transforms import v2


class EODataset(Dataset):
    """Reads different Earth Observation data sources from a directory."""

    def __init__(
        self, chips_path: List[Path], size: int, platforms: list, metadata: Box
    ) -> None:
        super().__init__()
        self.chips_path = chips_path
        self.size = size
        self.transforms = {}

        # Generate transforms for each platform using a helper function
        for platform in platforms:
            mean = list(metadata[platform].bands.mean.values())
            std = list(metadata[platform].bands.std.values())
            self.transforms[platform] = self.create_transforms(mean, std)

    def create_transforms(self, mean, std):
        return v2.Compose(
            [
                v2.RandomHorizontalFlip(p=0.5),
                v2.RandomVerticalFlip(p=0.5),
                v2.RandomCrop(size=(self.size, self.size)),
                v2.Normalize(mean=mean, std=std),
            ]
        )

    def __len__(self):
        return len(self.chips_path)

    def __getitem__(self, idx):
        chip_path = self.chips_path[idx]
        with np.load(chip_path, allow_pickle=False) as chip:
            pixels = torch.from_numpy(chip["pixels"].astype(np.float32))
            platform = chip_path.parent.name
            pixels = self.transforms[platform](pixels)

            # Prepare additional information
            additional_info = {
                "platform": platform,
                "time": torch.tensor(
                    np.hstack((chip["week_norm"], chip["hour_norm"])),
                    dtype=torch.float32,
                ),
                "latlon": torch.tensor(
                    np.hstack((chip["lat_norm"], chip["lon_norm"])), dtype=torch.float32
                ),
            }

        return {"pixels": pixels, **additional_info}


class ClaySampler(Sampler):
    def __init__(self, dataset, platforms, batch_size):
        self.dataset = dataset
        self.platforms = platforms
        self.batch_size = batch_size

        self.cubes_per_platform = {platform: [] for platform in platforms}
        for idx, chip_path in enumerate(self.dataset.chips_path):
            platform = chip_path.parent.name
            self.cubes_per_platform[platform].append(idx)

    def __iter__(self):
        cubes_per_platform_per_epoch = {}
        rng = np.random.default_rng()
        # Shuffle and adjust sizes
        max_len = max(len(indices) for indices in self.cubes_per_platform.values())
        for platform in self.platforms:
            indices = self.cubes_per_platform[platform]
            rng.shuffle(indices)
            repeated_indices = np.tile(indices, (max_len // len(indices) + 1))[:max_len]
            cubes_per_platform_per_epoch[platform] = repeated_indices

        # Create batches such that we return one platform per batch in cycle
        # Ignore the last batch if it is incomplete
        for i in range(0, max_len, self.batch_size):
            for platform in self.platforms:
                batch = cubes_per_platform_per_epoch[platform][i : i + self.batch_size]
                if len(batch) == self.batch_size:
                    yield batch

    def __len__(self):
        return len(self.dataset.chips_path) // self.batch_size


def batch_collate(batch):
    """Collate function for DataLoader.

    Merge the first two dimensions of the input tensors.
    """
    d = defaultdict(list)
    for item in batch:
        d["pixels"].append(item["pixels"])
        d["time"].append(item["time"])
        d["latlon"].append(item["latlon"])
        d["platform"].append(item["platform"])
    return {
        "pixels": rearrange(d["pixels"], "b1 b2 c h w -> (b1 b2) c h w"),
        "time": rearrange(d["time"], "b1 b2 t -> (b1 b2) t"),
        "latlon": rearrange(d["latlon"], "b1 b2 ll -> (b1 b2) ll"),
        "platform": d["platform"],
    }


class ClayDataModule(L.LightningDataModule):
    def __init__(  # noqa: PLR0913
        self,
        data_dir: str = "data",
        size: int = 224,
        metadata_path: str = "configs/metadata.yaml",
        platforms: list = [
            "landsat-c2l1",
            "landsat-c2l2-sr",
            "linz",
            "naip",
            "sentinel-1-rtc",
            "sentinel-2-l2a",
        ],
        batch_size: int = 10,
        num_workers: int = 8,
    ):
        super().__init__()
        self.data_dir = data_dir
        self.size = size
        self.platforms = platforms
        self.metadata = Box(yaml.safe_load(open(metadata_path)))
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.split_ratio = 0.8

    def setup(self, stage: Literal["fit", "predict"] | None = None) -> None:
        # Get list of GeoTIFF filepaths from s3 bucket or data/ folder
        if self.data_dir.startswith("s3://"):
            dp = torchdata.datapipes.iter.IterableWrapper(iterable=[self.data_dir])
            chips_path = list(dp.list_files_by_s3(masks="*.npz"))
        else:  # if self.data_dir is a local data path
            chips_path = sorted(list(Path(self.data_dir).glob("**/*.npz")))
            chips_platform = [chip.parent.parent.name for chip in chips_path]
            # chips_platform = [chip.parent.parent.name for chip in chips_path]
        print(f"Total number of chips: {len(chips_path)}")

        if stage == "fit":
            trn_paths, val_paths = train_test_split(
                chips_path,
                test_size=(1 - self.split_ratio),
                stratify=chips_platform,
                shuffle=True,
            )

            self.trn_ds = EODataset(
                chips_path=trn_paths,
                size=self.size,
                platforms=self.platforms,
                metadata=self.metadata,
            )
            self.trn_sampler = ClaySampler(
                dataset=self.trn_ds,
                platforms=self.platforms,
                batch_size=self.batch_size,
            )
            self.val_ds = EODataset(
                chips_path=val_paths,
                size=self.size,
                platforms=self.platforms,
                metadata=self.metadata,
            )
            self.val_sampler = ClaySampler(
                dataset=self.val_ds,
                platforms=self.platforms,
                batch_size=self.batch_size,
            )

        elif stage == "predict":
            self.prd_ds = EODataset(
                chips_path=chips_path,
                platform=self.platform,
                metadata_path=self.metadata_path,
            )

    def train_dataloader(self):
        return DataLoader(
            self.trn_ds,
            num_workers=self.num_workers,
            batch_sampler=self.trn_sampler,
            collate_fn=batch_collate,
            pin_memory=True,
            prefetch_factor=4,
        )

    def val_dataloader(self):
        return DataLoader(
            self.val_ds,
            num_workers=self.num_workers,
            batch_sampler=self.val_sampler,
            collate_fn=batch_collate,
            pin_memory=True,
            prefetch_factor=4,
        )

    def predict_dataloader(self):
        return DataLoader(
            dataset=self.prd_ds,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=False,
        )