| """Create paired low/high resolution scenes with labels for kNN evaluation.""" |
| import argparse, json |
| from pathlib import Path |
| import numpy as np, yaml |
| from PIL import Image |
| ROOT = Path(__file__).resolve().parents[1] |
|
|
| def make_split(count, cfg, rng): |
| size, target, channels, classes = cfg["input_size"], cfg["target_size"], cfg["channels"], cfg["num_classes"] |
| labels = np.arange(count, dtype=np.int64) % classes; rng.shuffle(labels) |
| gsd = rng.choice(np.asarray(cfg["gsd_values"], dtype=np.float32), count) |
| y, x = np.mgrid[:target, :target].astype(np.float32); images = np.empty((count, channels, size, size), np.float32); targets = np.empty((count, channels, target, target), np.float32) |
| for i, label in enumerate(labels): |
| pattern = np.sin((x + label*2)*np.pi*(label+1)/target) + np.cos((y-label*2)*np.pi*(label+1)/target) |
| pattern = (pattern-pattern.min())/(pattern.max()-pattern.min()) |
| scene = np.stack([np.roll(pattern, label*c, axis=c%2) for c in range(channels)]) |
| targets[i] = np.clip(scene + rng.normal(0, .02 + .01*gsd[i], scene.shape), 0, 1) |
| images[i] = np.asarray([Image.fromarray((targets[i,c]*255).astype('uint8')).resize((size,size), Image.Resampling.BOX) for c in range(channels)], dtype=np.float32)/255 |
| return images, targets, labels, gsd |
|
|
| def main(): |
| p=argparse.ArgumentParser(); p.add_argument("--config", default=str(ROOT/"conf/config.yaml")); a=p.parse_args(); cfg=yaml.safe_load(Path(a.config).read_text()); d=cfg["data"]; out=ROOT/d["root"]; out.mkdir(exist_ok=True); rng=np.random.default_rng(cfg["seed"]) |
| for split,n in (("train",d["train_samples"]),("test",d["test_samples"])): |
| images,targets,labels,gsd=make_split(n,d,rng); np.savez_compressed(out/f"{split}.npz", images=images, targets=targets, labels=labels, gsd=gsd) |
| (out/"format.json").write_text(json.dumps({"images":"float32 BCHW input resolution","targets":"float32 BCHW target resolution","labels":"int64","gsd":"metres per pixel"},indent=2)+"\n"); print("created",out/"train.npz",out/"test.npz") |
| if __name__ == "__main__": main() |
|
|
