Spaces:
Sleeping
Sleeping
File size: 2,712 Bytes
b20c769 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import json
from pathlib import Path
import torch
import torch.multiprocessing
from einops import rearrange
from torch.utils.data import Dataset
from ..preprocess import normalize_bands
torch.multiprocessing.set_sharing_strategy("file_system")
class PASTISDataset(Dataset):
def __init__(self, path_to_splits: Path, split: str, norm_operation, augmentation, partition):
with (Path(__file__).parents[0] / Path("configs") / Path("pastis.json")).open("r") as f:
config = json.load(f)
# NOTE: I imputed bands for this dataset before saving the tensors, so no imputation is necessary
assert split in ["train", "val", "valid", "test"]
if split == "val":
split = "valid"
self.band_info = config["band_info"]
self.split = split
self.augmentation = augmentation
self.norm_operation = norm_operation
torch_obj = torch.load(path_to_splits / f"pastis_{split}.pt")
self.images = torch_obj["images"] # (N, 12, 13, 64, 64)
self.months = torch_obj["months"] - 1 # subtract 1 for zero-indexing , shape (N, 12)
self.labels = torch_obj["targets"] # (N, 64, 64)
if (partition != "default") and (split == "train"):
with open(path_to_splits / f"{partition}_partition.json", "r") as json_file:
subset_indices = json.load(json_file)
self.images = self.images[subset_indices]
self.months = self.months[subset_indices]
self.labels = self.labels[subset_indices]
def __len__(self):
return self.images.shape[0]
def __getitem__(self, idx):
images = self.images[idx] # (12, 13, 64, 64)
months = self.months[idx] # (12)
labels = self.labels[idx] # (64, 64)
assert images.shape[0] == 12
# normalize one timestep at a time
normed_images = []
for i in range(12):
# sorry for the ugly code
single_timestep_image = rearrange(images[i], "c h w -> h w c").numpy()
normed_image = torch.tensor(
normalize_bands(single_timestep_image, self.norm_operation, self.band_info)
)
normed_images.append(normed_image)
normed_images = torch.stack(normed_images) # (12, 64, 64, 13)
normed_images = rearrange(normed_images, "t h w c -> h w t c") # (64, 64, 12, 13)
assert normed_images.shape[-2] == 12
assert normed_images.shape[-1] == 13
# important note: augmentation for timeseries is not supported
# there is obviously a better way to do this but oh well, I'll remember it
return {"s2": normed_images, "target": labels, "months": months}
|