-
Notifications
You must be signed in to change notification settings - Fork 43
/
Copy pathopenearthmap.py
58 lines (52 loc) · 2.23 KB
/
openearthmap.py
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
from typing import Any
import torch
import albumentations as A
import kornia.augmentation as K
from torchgeo.datamodules import NonGeoDataModule
from torchgeo.transforms import AugmentationSequential
from terratorch.datasets import OpenEarthMapNonGeo
from terratorch.datamodules.utils import wrap_in_compose_is_list
MEANS = {
"BLUE": 116.628328,
"GREEN": 119.65935,
"RED": 113.385309
}
STDS = {
"BLUE": 44.668890717415586,
"GREEN": 48.282311849967364,
"RED": 54.19692448815262,
}
class OpenEarthMapNonGeoDataModule(NonGeoDataModule):
def __init__(
self,
batch_size: int = 8,
num_workers: int = 0,
data_root: str = "./",
train_transform: A.Compose | None | list[A.BasicTransform] = None,
val_transform: A.Compose | None | list[A.BasicTransform] = None,
test_transform: A.Compose | None | list[A.BasicTransform] = None,
aug: AugmentationSequential = None,
**kwargs: Any
) -> None:
super().__init__(OpenEarthMapNonGeo, batch_size, num_workers, **kwargs)
bands = kwargs.get("bands", OpenEarthMapNonGeo.all_band_names)
self.means = torch.tensor([MEANS[b] for b in bands])
self.stds = torch.tensor([STDS[b] for b in bands])
self.train_transform = wrap_in_compose_is_list(train_transform)
self.val_transform = wrap_in_compose_is_list(val_transform)
self.test_transform = wrap_in_compose_is_list(test_transform)
self.data_root = data_root
self.aug = AugmentationSequential(K.Normalize(self.means, self.stds), data_keys=["image"]) if aug is None else aug
def setup(self, stage: str) -> None:
if stage in ["fit"]:
self.train_dataset = self.dataset_class(
split="train", data_root=self.data_root, transform=self.train_transform, **self.kwargs
)
if stage in ["fit", "validate"]:
self.val_dataset = self.dataset_class(
split="val", data_root=self.data_root, transform=self.val_transform, **self.kwargs
)
if stage in ["test"]:
self.test_dataset = self.dataset_class(
split="test",data_root=self.data_root, transform=self.test_transform, **self.kwargs
)