Class SitsDataset

data/data.py:14–132

class SitsDataset(Dataset):
    def __init__(self,
                 path,
                 split,
                 domain_shift,
                 num_channels,
                 num_classes,
                 img_size,
                 true_size,
                 train_length):
        super(SitsDataset, self).__init__()
        self.path = path
        self.image_folder, self.gt_folder = join(path, split if domain_shift else 'train'), join(path, 'labels')
        self.split = split
        self.domain_shift = domain_shift
        self.num_channels = num_channels
        self.num_classes = num_classes
        self.img_size = img_size
        self.true_size = true_size
        self.train_length = train_length
        self.monthly_dates = get_monthly_dates_dict()
        self.gt, self.sits_ids = self.load_ground_truth(split)
        self.collate_fn = collate_fn
        self.mean, self.std, self.month_list = None, None, None  # Needs to be defined in subclass

    def __len__(self):
        if self.split == 'train':
            return len(self.sits_ids) * ((self.true_size // self.img_size) ** 2 - 4)
        elif self.domain_shift:
            return len(self.sits_ids) * (self.true_size // self.img_size) ** 2
        else:
            return len(self.sits_ids) * 2

    def __getitem__(self, i):
        """Returns an item from the dataset.
        Args:
            i (int): index of the item
        Returns:
            dict: dictionary with keys "data", "gt", "positions", "idx"
        Shapes:
            data: T x C x H x W
            gt: T x H x W
            positions: T
            idx: 1
        """
        if self.split == 'train':
            num_patches_per_sits = (self.true_size // self.img_size) ** 2 - 4
            sits_number = i // num_patches_per_sits
            patch_loc_i, patch_loc_j = None, None
            months = self.get_random_months(sits_number)
        elif self.domain_shift:
            num_patches_per_sits = (self.true_size // self.img_size) ** 2
            sits_number = i // num_patches_per_sits
            patch_loc_i = (i % num_patches_per_sits) // (self.true_size // self.img_size)
            patch_loc_j = (i % num_patches_per_sits) % (self.true_size // self.img_size)
            months = list(range(24))
        else:
            num_patches_per_sits = 2