| """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_) |
| |
| 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() |
|
|
