File size: 3,494 Bytes
aed6f6f cf7a05c aed6f6f cf7a05c aed6f6f cf7a05c aed6f6f | 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 | """Generate tiny Satlas-compatible independent train/test NPZ archives."""
import json
from pathlib import Path
import numpy as np
import yaml
ROOT = Path(__file__).resolve().parents[1]
def make_split(samples, config, seed):
rng = np.random.default_rng(seed); d = config["data"]; size = d["highres_size"]
highres = rng.random((samples, d["highres_time_steps"], d["highres_channels"], size, size), dtype=np.float32)
lowres = rng.random((samples, d["lowres_time_steps"], d["lowres_channels"],
d["lowres_size"], d["lowres_size"]), dtype=np.float32)
valid_highres_times = np.ones((samples, d["highres_time_steps"]), dtype=np.bool_)
valid_lowres_times = np.ones((samples, d["lowres_time_steps"]), dtype=np.bool_)
# Exercise independent missing-observation masks even with a single smoke sample.
valid_highres_times[:, -1] = False
valid_lowres_times[:, -2:] = False
highres[~valid_highres_times] = 0
lowres[~valid_lowres_times] = 0
yy, xx = np.mgrid[:size, :size]; segmentation = np.empty((samples, size, size), dtype=np.int64)
point = np.zeros((samples, 1, size, size), np.float32); polygon = np.zeros_like(point); polyline = np.zeros_like(point)
for index in range(samples):
segmentation[index] = (xx >= size // 2) + 2 * (yy >= size // 2)
center = size // 4 + index * 8; point[index, 0, center-2:center+3, center-2:center+3] = 1
polygon[index, 0, size//3:2*size//3, size//3:2*size//3] = 1
polyline[index, 0, np.arange(size), (np.arange(size) + index * 8) % size] = 1
last_valid = valid_highres_times.sum(1) - 1
regression = np.stack([0.7 * highres[index, last_valid[index], :1] + 0.3 * highres[index, 0, 1:2]
for index in range(samples)]).astype(np.float32)
return dict(highres_images=highres, lowres_images=lowres,
valid_highres_times=valid_highres_times, valid_lowres_times=valid_lowres_times,
sample_ids=np.asarray([f"sample-{seed}-{index}" for index in range(samples)]), segmentation=segmentation,
regression=regression, point=point, polygon=polygon, polyline=polyline,
property=np.arange(samples, dtype=np.int64) % 3,
classification=(np.arange(samples, dtype=np.int64) + 1) % 3)
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()); d = config["data"]
output = ROOT / d["root"]; output.mkdir(parents=True, exist_ok=True)
for split, count, seed in (("train", d["train_samples"], config["seed"]), ("test", d["test_samples"], config["seed"] + 1)):
np.savez_compressed(output / f"{split}.npz", **make_split(count, config, seed),
source=np.asarray("synthetic"), protocol=np.asarray(d["protocol"]))
(output / "format.json").write_text(json.dumps({
"protocol": d["protocol"], "source": "synthetic",
"highres_images": "float32 [N,4,3,512,512] NAIP RGB",
"lowres_images": "float32 [N,8,9,512,512] Sentinel-2 multispectral",
"valid_highres_times": "bool [N,4], false slots are zero-filled",
"valid_lowres_times": "bool [N,8], false slots are zero-filled",
"dense_targets": "segmentation [N,512,512]; regression/point/polygon/polyline [N,1,512,512]",
"global_targets": "property/classification int64 [N]"
}, indent=2) + "\n")
print(f"created {output / 'train.npz'} and {output / 'test.npz'}")
if __name__ == "__main__": main()
|