File size: 6,562 Bytes
8bfc737 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import numpy as np
if __package__ in (None, ""):
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from script.utils import DEFAULT_CONFIG, load_config
SPLIT_OFFSETS = {"train": 0, "val": 100_000, "test": 200_000}
# Matches the official Earthformer SEVIR config
# (scripts/cuboid_transformer/sevir/earthformer_sevir_v1.yaml):
# dataset.img_height/img_width = 384, in_len = 13, out_len = 12,
# seq_len = 25, interval_real_time = 5, sample_mode = "sequent",
# stride = 12, metrics_list = ['csi', 'pod', 'sucr', 'bias'],
# threshold_list = [16, 74, 133, 160, 181, 219].
SEVIR_VIL_REFERENCE = {
"dataset": "SEVIR VIL",
"spatial_shape": [384, 384, 1],
"sequence": {"input_frames": 13, "output_frames": 12},
"seq_len": 25,
"sample_mode": "sequent",
"stride": 12,
"frame_interval_minutes": 5,
"thresholds": [16, 74, 133, 160, 181, 219],
"layout": "train/val/test NPZ splits; each holds inputs [N,13,H,W,1] and targets [N,12,H,W,1]",
}
def generate_sequence(height: int, width: int, input_length: int, output_length: int, seed: int) -> np.ndarray:
"""Generate a continuous, temporally coherent 25-frame SEVIR-like VIL window.
The window is sampled exactly like official SEVIR "sequent" mode: a single
continuous sequence of seq_len = input_length + output_length frames is
produced and later split into the first 13 input and the last 12 target
frames, matching the official 13 -> 12 (65 -> 60 minute) task at 5-minute
intervals.
"""
rng = np.random.default_rng(seed)
total = input_length + output_length
yy, xx = np.mgrid[:height, :width]
background = rng.normal(0.0, 0.008, (height, width)).astype(np.float32)
background = (background + np.roll(background, 1, 0) + np.roll(background, 1, 1)) / 3.0
cell_count = int(rng.integers(2, 5))
cells = []
for _ in range(cell_count):
cells.append(
(
rng.uniform(0.15 * width, 0.85 * width),
rng.uniform(0.15 * height, 0.85 * height),
rng.uniform(-0.45, 0.45),
rng.uniform(-0.45, 0.45),
rng.uniform(max(1.2, width / 18), max(2.0, width / 8)),
rng.uniform(max(1.2, height / 18), max(2.0, height / 8)),
rng.uniform(0.45, 0.95),
rng.uniform(-0.035, 0.035),
rng.uniform(0, np.pi),
)
)
frames = np.empty((total, height, width, 1), dtype=np.float32)
for time in range(total):
frame = np.maximum(background * (0.8 + 0.2 * np.sin(time / 5)), 0.0)
for cx, cy, vx, vy, sx, sy, amplitude, growth, angle in cells:
dx, dy = xx - (cx + vx * time), yy - (cy + vy * time)
ca, sa = np.cos(angle), np.sin(angle)
xr, yr = ca * dx + sa * dy, -sa * dx + ca * dy
scale = np.clip(1.0 + growth * time, 0.55, 1.6)
intensity = amplitude * np.exp(-0.5 * ((xr / (sx * scale)) ** 2 + (yr / (sy * scale)) ** 2))
lifecycle = np.clip(1.0 + growth * time, 0.35, 1.25)
frame += intensity.astype(np.float32) * lifecycle
noise = rng.normal(0.0, 0.004, (height, width)).astype(np.float32)
frames[time, ..., 0] = np.clip(frame + noise, 0.0, 1.0)
return frames
def generate_split(samples: int, height: int, width: int, input_length: int, output_length: int, seed: int) -> tuple[np.ndarray, np.ndarray]:
sequences = np.stack(
[generate_sequence(height, width, input_length, output_length, seed + index) for index in range(samples)]
)
return sequences[:, :input_length], sequences[:, input_length:]
def parse_args() -> argparse.Namespace:
config_parser = argparse.ArgumentParser(add_help=False)
config_parser.add_argument("--config", default=str(DEFAULT_CONFIG))
config_args, _ = config_parser.parse_known_args()
defaults = load_config(config_args.config)
data = defaults["data"]
parser = argparse.ArgumentParser(description="Generate deterministic synthetic SEVIR-like VIL sequences")
parser.add_argument("--config", default=config_args.config)
parser.add_argument("--output-dir", default=data["data_dir"])
parser.add_argument("--height", type=int, default=int(data["height"]))
parser.add_argument("--width", type=int, default=int(data["width"]))
parser.add_argument("--train-samples", type=int, default=int(data["train_samples"]))
parser.add_argument("--val-samples", type=int, default=int(data["val_samples"]))
parser.add_argument("--test-samples", type=int, default=int(data["test_samples"]))
parser.add_argument("--seed", type=int, default=int(defaults["train"]["seed"]))
return parser.parse_args()
def main() -> None:
args = parse_args()
config = load_config(args.config)
data = config["data"]
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
split_sizes = {"train": args.train_samples, "val": args.val_samples, "test": args.test_samples}
for split, samples in split_sizes.items():
if samples <= 0:
raise ValueError(f"{split} samples must be positive")
inputs, targets = generate_split(
samples,
args.height,
args.width,
int(data["input_length"]),
int(data["output_length"]),
args.seed + SPLIT_OFFSETS[split],
)
np.savez_compressed(output_dir / f"{split}.npz", inputs=inputs, targets=targets)
metadata = {
"synthetic": True,
"official_sevir": False,
"description": "Deterministic synthetic SEVIR-like VIL; not official SEVIR data",
"protocol": "synthetic_sevir",
"reference": SEVIR_VIL_REFERENCE,
"seq_len": int(data["input_length"]) + int(data["output_length"]),
"sample_mode": "sequent",
"stride": SEVIR_VIL_REFERENCE["stride"],
"frame_interval_minutes": int(data["frame_interval_minutes"]),
"input_frames": int(data["input_length"]),
"output_frames": int(data["output_length"]),
"shape": [args.height, args.width, 1],
"normalization": "unit [0,1] float32",
"seed": args.seed,
"splits": split_sizes,
}
(output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
print(json.dumps({"output_dir": str(output_dir), "metadata": metadata}, indent=2))
if __name__ == "__main__":
main()
|