torchgeo/tests/data/globbiomass/data.py

64 строки
1.7 KiB
Python
Executable File

#!/usr/bin/env python3
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
import hashlib
import os
import zipfile
import numpy as np
import rasterio
np.random.seed(0)
SIZE = 64
files = {
'agb': ['N00E020_agb.tif', 'N00E020_agb_err.tif'],
'gsv': ['N00E020_gsv.tif', 'N00E020_gsv_err.tif'],
}
def create_file(path: str, dtype: str, num_channels: int) -> None:
profile = {}
profile['driver'] = 'GTiff'
profile['dtype'] = dtype
profile['count'] = num_channels
profile['crs'] = 'epsg:4326'
profile['transform'] = rasterio.transform.from_bounds(0, 0, 1, 1, 1, 1)
profile['height'] = SIZE
profile['width'] = SIZE
profile['compress'] = 'lzw'
profile['predictor'] = 2
Z = np.random.randint(
np.iinfo(profile['dtype']).max, size=(1, SIZE, SIZE), dtype=profile['dtype']
)
with rasterio.open(path, 'w', **profile) as src:
src.write(Z)
if __name__ == '__main__':
for measurement, file_paths in files.items():
zipfilename = f'N00E020_{measurement}.zip'
files_to_zip = []
for path in file_paths:
# remove old data
if os.path.exists(path):
os.remove(path)
# Create mask file
create_file(path, dtype='int32', num_channels=1)
files_to_zip.append(path)
# Compress data
with zipfile.ZipFile(zipfilename, 'w') as zip:
for file in files_to_zip:
zip.write(file, arcname=file)
# Compute checksums
with open(zipfilename, 'rb') as f:
md5 = hashlib.md5(f.read()).hexdigest()
print(f'{zipfilename}: {md5}')