MSA-TransUNet/transforms.py at main · hhza2222/MSA-TransUNet · GitHub

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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
import numpy as np
import random
import torch
from torchvision import transforms as T
import torch.nn.functional as F
from scipy.ndimage import rotate, median_filter

class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, input_data, label_data, mask):
        for t in self.transforms:
            input_data, label_data, mask = t(input_data, label_data, mask)
        return input_data, label_data, mask


class OceanPadding(object):
    def __init__(self, land_fill=np.nan):
        self.land_fill = land_fill

    def __call__(self, input_data, label_data, mask):
        input_data[:, ~mask] = np.nan
        input_channel_means = np.nanmean(input_data, axis=(1, 2), keepdims=True)
        input_data = np.where(mask, np.nan_to_num(input_data, nan=input_channel_means), np.nan)
        ocean_salinity = label_data[mask]
        salinity_mean = np.nanmean(ocean_salinity)
        label_data = np.where(mask, np.nan_to_num(label_data, nan=salinity_mean), self.land_fill)
        return input_data, label_data, mask


class RandomHorizontalFlip(object):
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, input_data, label_data, mask):
        if np.random.rand() < self.flip_prob:
            input_data = np.flip(input_data, axis=2)
            label_data = np.flip(label_data, axis=1)
            mask = np.flip(mask, axis=1)
        return input_data, label_data, mask


class RandomRotate(object):
    def __init__(self, rotate_prob):
        self.rotate_prob = rotate_prob

    def __call__(self, input_data, label_data, mask):
        if np.random.rand() < self.rotate_prob:
            input_data = rotate(input_data, 180, axes=(1, 2), reshape=False, order=1)
            label_data = rotate(label_data, 180, axes=(0, 1), reshape=False, order=0)
            mask = rotate(mask, 180, axes=(0, 1), reshape=False, order=0)
        return input_data, label_data, mask


class ToTensor(object):
    def __call__(self, input_data, label_data, mask):
        input_data = torch.from_numpy(input_data.astype(np.float32))  # (C, H, W)
        label_data = torch.from_numpy(label_data.astype(np.float32)).unsqueeze(0)  # (1, H, W)
        mask = torch.tensor(mask.copy(), dtype=torch.bool).unsqueeze(0)  # (1, H, W)
        return input_data, label_data, mask


class PadToFixedSize(object):
    def __init__(self, size=224):
        self.size = size

    def __call__(self, input_data, label_data, mask):
        h, w = input_data.shape[-2:]
        pad_h = max(self.size - h, 0)
        pad_w = max(self.size - w, 0)
        pad_h = int(pad_h / 2)
        pad_w = int(pad_w / 2)

        input_data = F.pad(input_data, (pad_w, pad_w, pad_h, pad_h), mode='reflect')
        label_data = F.pad(label_data, (pad_w, pad_w, pad_h, pad_h), mode='reflect')
        mask = F.pad(mask, (pad_w, pad_w, pad_h, pad_h), mode='constant', value=False)
        return input_data, label_data, mask


class MaskedNormalize(object):
    def __init__(self, mean, std, fill_nan=0.0):
        self.mean = mean
        self.std = std
        self.fill_nan = fill_nan

    def __call__(self, input_data, label_data, mask):
        mean = self.mean
        std = self.std

        label_mean = mean[-1]
        label_std = std[-1]
        label_data0 = label_data.clone()
        label_data0[~torch.isnan(label_data0)] = (label_data0[~torch.isnan(label_data0)] - label_mean) / label_std
        label_data = label_data0.clone()
        label_data[torch.isnan(label_data)] = 0.0

        for c in range(input_data.shape[0]):
            c_data = input_data[c].clone()
            input_mean = mean[c]
            input_std = std[c]
            c_data[~torch.isnan(c_data)] = (c_data[~torch.isnan(c_data)] - input_mean) / input_std
            c_data[torch.isnan(c_data)] = 0.0
            input_data[c] = c_data.clone()
        return input_data, label_data, mask