Class DynamicEarthNet

data/data.py:180–229

class DynamicEarthNet(SitsDataset):
    def __init__(
            self,
            path,
            split="train",
            domain_shift=False,
            num_channels=4,
            num_classes=7,
            img_size=128,
            true_size=1024,
            train_length=6,
            date_aug_range=2
    ):

        super(DynamicEarthNet, self).__init__(path=path,
                                              split=split,
                                              domain_shift=domain_shift,
                                              num_channels=num_channels,
                                              num_classes=num_classes,
                                              img_size=img_size,
                                              true_size=true_size,
                                              train_length=train_length)
        """Initializes the dataset.
        Args:
            path (str): path to the dataset
            split (str): split to use (train, val, test)
            domain_shift (bool): if val/test, whether we are in a domain shift setting or not
        """
        self.date_aug_range = date_aug_range
        self.monthly_dates = get_monthly_dates_dict()
        self.gt, self.sits_ids = self.load_ground_truth(split)
        self.month_list = [list(range(24)) for _ in range(self.gt.shape[0])]
        self.mean = torch.tensor([83.1029, 80.7615, 69.3328, 133.8648], dtype=torch.float16).reshape(4, 1, 1)
        self.std = torch.tensor([33.2714, 25.5288, 23.9868, 30.5591], dtype=torch.float16).reshape(4, 1, 1)

    def load_data(self, sits_number, sits_id, months, curr_sits_path):
        data = torch.zeros((len(months), self.num_channels, self.true_size, self.true_size), dtype=torch.float16)
        days = [self.random_date_augmentation(month) for month in months]
        name_rgb = [f'{sits_id}_{day}_rgb.jpeg' for day in days]
        name_infra = [f'{sits_id}_{day}_infra.jpeg' for day in days]
        for d, (n_rgb, n_infra) in enumerate(zip(name_rgb, name_infra)):
            data[d, :3] = torchvision.io.read_image(join(curr_sits_path, n_rgb))
            data[d, 3] = torchvision.io.read_image(join(curr_sits_path, n_infra))
        return data, days

    def random_date_augmentation(self, month):
        if self.split == 'train':
            return max(0, random.randint(0, self.date_aug_range * 2) - self.date_aug_range + self.monthly_dates[month])
        else:
            return self.monthly_dates[month]


def collate_fn(batch):