feat: add dataset and transforms for TIR triplets

src/data/dataset.py

@@ -0,0 +1,37 @@
+import os
+
+import torch
+from torch.utils.data import Dataset
+
+from src.data.transforms import augment_triplet
+
+
+class GOESTripletDataset(Dataset):
+    """PyTorch Dataset for loading pre-processed Satellite TIR triplets.
+
+    Expects data to be stored as `.pt` files containing tensors of shape [3, 1, H, W],
+    representing Past, Present, and Future Brightness Temperature frames.
+    """
+
+    def __init__(self, data_dir: str, augment: bool = True):
+        self.data_dir = data_dir
+        self.augment = augment
+        self.triplet_files = sorted(
+            [f for f in os.listdir(data_dir) if f.endswith(".pt")]
+        )
+
+    def __len__(self) -> int:
+        return len(self.triplet_files)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        file_path = os.path.join(self.data_dir, self.triplet_files[idx])
+        triplet = torch.load(file_path)  # [3, 1, H, W]
+
+        img0 = triplet[0]  # Past
+        gt = triplet[1]  # Ground Truth / Present
+        img1 = triplet[2]  # Future
+
+        if self.augment:
+            img0, img1, gt = augment_triplet(img0, img1, gt)
+
+        return img0, img1, gt

src/data/transformes.py

@@ -0,0 +1,51 @@
+import random
+
+import torch
+import torchvision.transforms.functional as TF
+
+
+def augment_triplet(img0: torch.Tensor, img1: torch.Tensor, gt: torch.Tensor, crop_size: int = 256) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Applies identical spatial augmentations to all three frames in a triplet.
+    
+    Ensures that the cloud motions remain physically consistent across the temporal sequence.
+    
+    Args:
+        img0: Past frame tensor [C, H, W]
+        img1: Future frame tensor [C, H, W]
+        gt: Ground truth intermediate frame tensor [C, H, W]
+        crop_size: Final output dimension for height and width.
+        
+    Returns:
+        Tuple of augmented tensors (img0, img1, gt).
+    """
+    _, h, w = img0.shape
+    
+    # Random Spatial Cropping
+    top = random.randint(0, h - crop_size)
+    left = random.randint(0, w - crop_size)
+    
+    img0 = TF.crop(img0, top, left, crop_size, crop_size)
+    img1 = TF.crop(img1, top, left, crop_size, crop_size)
+    gt = TF.crop(gt, top, left, crop_size, crop_size)
+    
+    # Random Horizontal Flip
+    if random.random() > 0.5:
+        img0 = TF.hflip(img0)
+        img1 = TF.hflip(img1)
+        gt = TF.hflip(gt)
+        
+    # Random Vertical Flip
+    if random.random() > 0.5:
+        img0 = TF.vflip(img0)
+        img1 = TF.vflip(img1)
+        gt = TF.vflip(gt)
+        
+    # Random Dihedral Transformations
+    angles = [0, 90, 180, 270]
+    angle = random.choice(angles)
+    if angle > 0:
+        img0 = TF.rotate(img0, angle)
+        img1 = TF.rotate(img1, angle)
+        gt = TF.rotate(gt, angle)
+        
+    return img0, img1, gt

tests/test_dataset.py

@@ -0,0 +1,24 @@
+import torch
+import pytest
+from src.data.transforms import augment_triplet
+from src.data.dataset import GOESTripletDataset
+
+def test_dataset(tmp_path):
+    # Setup mock data
+    tensor_data = torch.zeros((3, 1, 512, 512))
+    torch.save(tensor_data, tmp_path / "triplet_001.pt")
+    
+    dataset = GOESTripletDataset(str(tmp_path), augment=False)
+    assert len(dataset) == 1
+    
+    img0, img1, gt = dataset[0]
+    assert img0.shape == (1, 512, 512)
+    assert gt.shape == (1, 512, 512)
+
+def test_augment_triplet():
+    img0 = torch.zeros((1, 512, 512))
+    img1 = torch.zeros((1, 512, 512))
+    gt = torch.zeros((1, 512, 512))
+    
+    a_img0, a_img1, a_gt = augment_triplet(img0, img1, gt, crop_size=256)
+    assert a_img0.shape == (1, 256, 256)