Ehzoahis/BigEarthNet

1import os2import json3import shutil4import string5import tarfile6import tifffile7import datasets8 9import numpy as np10import pandas as pd11 12from tqdm import tqdm13 14class_sets = {15        19: [16            'Urban fabric',17            'Industrial or commercial units',18            'Arable land',19            'Permanent crops',20            'Pastures',21            'Complex cultivation patterns',22            'Land principally occupied by agriculture, with significant areas of natural vegetation',23            'Agro-forestry areas',24            'Broad-leaved forest',25            'Coniferous forest',26            'Mixed forest',27            'Natural grassland and sparsely vegetated areas',28            'Moors, heathland and sclerophyllous vegetation',29            'Transitional woodland, shrub',30            'Beaches, dunes, sands',31            'Inland wetlands',32            'Coastal wetlands',33            'Inland waters',34            'Marine waters',35        ],36        43: [37            'Continuous urban fabric',38            'Discontinuous urban fabric',39            'Industrial or commercial units',40            'Road and rail networks and associated land',41            'Port areas',42            'Airports',43            'Mineral extraction sites',44            'Dump sites',45            'Construction sites',46            'Green urban areas',47            'Sport and leisure facilities',48            'Non-irrigated arable land',49            'Permanently irrigated land',50            'Rice fields',51            'Vineyards',52            'Fruit trees and berry plantations',53            'Olive groves',54            'Pastures',55            'Annual crops associated with permanent crops',56            'Complex cultivation patterns',57            'Land principally occupied by agriculture, with significant areas of natural vegetation',58            'Agro-forestry areas',59            'Broad-leaved forest',60            'Coniferous forest',61            'Mixed forest',62            'Natural grassland',63            'Moors and heathland',64            'Sclerophyllous vegetation',65            'Transitional woodland/shrub',66            'Beaches, dunes, sands',67            'Bare rock',68            'Sparsely vegetated areas',69            'Burnt areas',70            'Inland marshes',71            'Peatbogs',72            'Salt marshes',73            'Salines',74            'Intertidal flats',75            'Water courses',76            'Water bodies',77            'Coastal lagoons',78            'Estuaries',79            'Sea and ocean',80        ],81    }82 83label_converter = {84        0: 0,85        1: 0,86        2: 1,87        11: 2,88        12: 2,89        13: 2,90        14: 3,91        15: 3,92        16: 3,93        18: 3,94        17: 4,95        19: 5,96        20: 6,97        21: 7,98        22: 8,99        23: 9,100        24: 10,101        25: 11,102        31: 11,103        26: 12,104        27: 12,105        28: 13,106        29: 14,107        33: 15,108        34: 15,109        35: 16,110        36: 16,111        38: 17,112        39: 17,113        40: 18,114        41: 18,115        42: 18,116    }117 118S2_MEAN = [752.40087073, 884.29673756, 1144.16202635, 1297.47289228, 1624.90992062, 2194.6423161, 2422.21248945, 2517.76053101, 2581.64687018, 2645.51888987, 2368.51236873, 1805.06846033]119S2_STD = [1108.02887453, 1155.15170768, 1183.6292542, 1368.11351514, 1370.265037, 1355.55390699, 1416.51487101, 1474.78900051, 1439.3086061, 1582.28010962, 1455.52084939, 1343.48379601]120 121S1_MEAN = [-12.54847273, -20.19237134]122S1_STD = [5.25697717, 5.91150917]123 124parts = [f"a{letter}" for letter in string.ascii_lowercase]125parts.extend([f"b{letter}" for letter in string.ascii_lowercase[:8]]) 126 127class BigEarthNetDataset(datasets.GeneratorBasedBuilder):128    VERSION = datasets.Version("1.0.0")129 130    DATA_URL = [131        f"https://huggingface.co/datasets/GFM-Bench/BigEarthNet/resolve/main/data/bigearthnet_part_{part}"132        for part in parts133    ]134 135    metadata = {136        "s2c": {137            "bands":["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B9", "B11", "B12"],138            "channel_wv": [442.7, 492.4, 559.8, 664.6, 704.1, 740.5, 782.8, 832.8, 864.7, 945.1, 1613.7, 2202.4],139            "mean": S2_MEAN,140            "std": S2_STD141        },142        "s1": {143            "bands": ["VV", "VH"],144            "channel_wv": [5500, 5700],145            "mean": S1_MEAN,146            "std": S1_STD147        }148    }149 150    SIZE = HEIGHT = WIDTH = 120151 152    NUM_CLASSES = 19153 154    spatial_resolution = 10155 156    def __init__(self, *args, **kwargs):157        self.class2idx = {c: i for i, c in enumerate(class_sets[43])}158 159        super().__init__(*args, **kwargs)160 161    def _info(self):162        metadata = self.metadata163        metadata['size'] = self.SIZE164        metadata['num_classes'] = self.NUM_CLASSES165        metadata['spatial_resolution'] = self.spatial_resolution166        return datasets.DatasetInfo(167            description=json.dumps(metadata),168            features=datasets.Features({169                "optical": datasets.Array3D(shape=(12, self.HEIGHT, self.WIDTH), dtype="float32"),170                "radar": datasets.Array3D(shape=(2, self.HEIGHT, self.WIDTH), dtype="float32"),171                "optical_channel_wv": datasets.Sequence(datasets.Value("float32")),172                "radar_channel_wv": datasets.Sequence(datasets.Value("float32")),173                "label": datasets.Sequence(datasets.Value("float32"), length=self.NUM_CLASSES),174                "spatial_resolution": datasets.Value("int32"),175            }),176        )177 178    def _split_generators(self, dl_manager):179        print(dl_manager.download_config.cache_dir)180        # Ensure cache directory is set181        if dl_manager.download_config.cache_dir is None:182            return []183        if isinstance(self.DATA_URL, list):184            try:185                downloaded_files = dl_manager.download(self.DATA_URL)186                print(f"downloaded files: {downloaded_files}")187                combined_file = os.path.join(dl_manager.download_config.cache_dir, "combined.tar.gz") 188                print(f"copying files to {combined_file}")189                target_dir = os.path.dirname(combined_file)190                os.makedirs(target_dir, exist_ok=True)  # Create only the directory                191                with open(combined_file, 'wb') as outfile:192                    counter = 0 193                    for part_file in tqdm(downloaded_files, desc="Copying files", unit="file"):194                        # print(f"copying {counter}-th file: {part_file}")195                        with open(part_file, 'rb') as infile:196                            shutil.copyfileobj(infile, outfile)197                        counter += 1198                print(f"extacting from {combined_file}")199                # data_dir = dl_manager.extract(combined_file)200                data_dir = os.path.join(dl_manager.download_config.cache_dir, "extracted")201                os.makedirs(data_dir, exist_ok=True)202                with tarfile.open(combined_file, "r:gz") as tar:203                    tar.extractall(path=data_dir)204                os.remove(combined_file)205                print(f"data_dir: {data_dir}")206            except Exception as e:207                print(f"exception: {e}, so setting data_dir to None")208                data_dir = None209        else:210            data_dir = dl_manager.download_and_extract(self.DATA_URL)211 212        return [213            datasets.SplitGenerator(214                name="train",215                gen_kwargs={216                    "split": 'train',217                    "data_dir": data_dir, 218                },219            ),220            datasets.SplitGenerator(221                name="val",222                gen_kwargs={223                    "split": 'val',224                    "data_dir": data_dir,225                },226            ),227            datasets.SplitGenerator(228                name="test",229                gen_kwargs={230                    "split": 'test',231                    "data_dir": data_dir,232                },233            )234        ]235 236    def _generate_examples(self, split, data_dir):237        optical_channel_wv = np.array(self.metadata["s2c"]["channel_wv"])238        radar_channel_wv = np.array(self.metadata["s1"]["channel_wv"])239        spatial_resolution = self.spatial_resolution240 241        data_dir = os.path.join(data_dir, "BigEarthNet")242        metadata = pd.read_csv(os.path.join(data_dir, "metadata.csv"))243        metadata = metadata[metadata["split"] == split].reset_index(drop=True)244 245        for index, row in metadata.iterrows():246            optical_path = os.path.join(data_dir, row.optical_path)247            optical = self._read_image(optical_path).astype(np.float32) # CxHxW248 249            radar_path = os.path.join(data_dir, row.radar_path)250            radar = self._read_image(radar_path).astype(np.float32)251 252            label_path = os.path.join(data_dir, row.label_path)253            label = self._load_label(label_path)254 255            sample = {256                "optical": optical,257                "radar": radar,258                "optical_channel_wv": optical_channel_wv,259                "radar_channel_wv": radar_channel_wv,260                "label": label,261                "spatial_resolution": spatial_resolution,262            }263 264            yield f"{index}", sample265    266    def _load_label(self, label_path):267        with open(label_path) as f:268            labels = json.load(f)['labels']269        indices =[self.class2idx[label] for label in labels]270        indices_optional = [label_converter.get(idx) for idx in indices]271        indices = [idx for idx in indices_optional if idx is not None]272        label = np.zeros(19, dtype=np.int64)273        label[indices] = 1274        return label275    276    def _read_image(self, image_path):277        """Read tiff image from image_path278        Args:279            image_path: 280                Image path to read from281 282        Return:283            image: 284                C, H, W numpy array image285        """286        image = tifffile.imread(image_path)287        image = np.transpose(image, (2, 0, 1))288 289        return image