File size: 5,285 Bytes
1558db5 2d7eead 1558db5 2d7eead 1558db5 2d7eead 1558db5 2d7eead 1558db5 2d7eead 1558db5 2d7eead 1558db5 | 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 | """Run SkySense inference and save arrays for evaluation."""
import importlib.util
import argparse
from pathlib import Path
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
def load_model_class():
spec = importlib.util.spec_from_file_location("skysense_model", ROOT / "model" / "skysense.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module.SkySense
def main():
parser = argparse.ArgumentParser(description="Run batched SkySense segmentation inference")
parser.add_argument("--batch-size", type=int)
args = parser.parse_args()
with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
config = yaml.safe_load(handle)
checkpoint_path = ROOT / config["paths"]["checkpoint"]
if not checkpoint_path.exists():
raise FileNotFoundError(
f"Missing checkpoint: {checkpoint_path.relative_to(ROOT)}. "
"Run `python scripts/train.py` first."
)
use_accelerator = torch.cuda.is_available() and config["runtime"].get("device", "auto") != "cpu"
device = torch.device("cuda" if use_accelerator else "cpu")
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
SkySense = load_model_class()
model = SkySense(
**config["model"],
hr_channels=config["data"]["hr_channels"],
s2_channels=config["data"]["s2_channels"],
s1_channels=config["data"]["s1_channels"],
num_classes=config["data"]["num_classes"],
).to(device)
model.load_state_dict(checkpoint["model"])
model.eval()
test_path = ROOT / config["data"]["root"] / "test.npz"
if not test_path.exists():
raise FileNotFoundError(
f"Missing inference data: {test_path.relative_to(ROOT)}. "
"Run `python scripts/fake_data.py` first."
)
archive = np.load(test_path)
keys = ["hr", "s2", "s1", "dates_hr", "dates_s2", "dates_s1", "region"]
arrays = {key: archive[key] for key in keys}
expected = {
"hr": (config["data"]["hr_timesteps"], config["data"]["hr_channels"], config["data"]["hr_size"], config["data"]["hr_size"]),
"s2": (config["data"]["s2_timesteps"], config["data"]["s2_channels"], config["data"]["s2_size"], config["data"]["s2_size"]),
"s1": (config["data"]["s1_timesteps"], config["data"]["s1_channels"], config["data"]["s1_size"], config["data"]["s1_size"]),
"dates_hr": (config["data"]["hr_timesteps"],),
"dates_s2": (config["data"]["s2_timesteps"],),
"dates_s1": (config["data"]["s1_timesteps"],),
"region": (),
}
sample_count = len(arrays["hr"])
for key, shape in expected.items():
if len(arrays[key]) != sample_count or tuple(arrays[key].shape[1:]) != shape:
raise ValueError(f"Invalid test {key} shape {arrays[key].shape}; expected [N,{','.join(map(str, shape))}]")
for key in ("hr", "s2", "s1"):
if not np.issubdtype(arrays[key].dtype, np.floating):
raise TypeError(f"{key} must use a floating dtype")
for key in ("dates_hr", "dates_s2", "dates_s1", "region"):
if arrays[key].dtype != np.int64:
raise TypeError(f"{key} must use int64")
if any(np.any((arrays[key] < 0) | (arrays[key] > 364)) for key in ("dates_hr", "dates_s2", "dates_s1")):
raise ValueError("Test dates must be in [0, 364]")
if np.any((arrays["region"] < 0) | (arrays["region"] >= config["model"]["num_regions"])):
raise ValueError("Test region IDs are out of range")
labels = archive["labels"]
if labels.dtype != np.int64 or labels.shape != (sample_count, config["data"]["hr_size"], config["data"]["hr_size"]):
raise ValueError("Test labels must be int64 [N,hr_size,hr_size]")
batch_size = args.batch_size or config["train"]["batch_size"]
predictions = []
all_probabilities = []
with torch.inference_mode():
for start in range(0, len(arrays["hr"]), batch_size):
tensors = {key: torch.from_numpy(value[start:start + batch_size]).to(device)
for key, value in arrays.items()}
output = model(tensors["hr"], tensors["s2"], tensors["s1"], tensors["dates_hr"], tensors["dates_s2"], tensors["dates_s1"], tensors["region"])
probabilities = output["logits"].softmax(dim=1).cpu().numpy()
all_probabilities.append(probabilities)
predictions.append(probabilities.argmax(axis=1))
output_dir = ROOT / config["paths"]["inference_dir"]
output_dir.mkdir(parents=True, exist_ok=True)
np.save(output_dir / "predictions.npy", np.concatenate(predictions))
np.save(output_dir / "probabilities.npy", np.concatenate(all_probabilities))
np.save(output_dir / "targets.npy", labels)
data_source = str(archive["data_source"]) if "data_source" in archive.files else "unknown"
protocol = str(archive["protocol"]) if "protocol" in archive.files else "unknown"
np.savez(output_dir / "metadata.npz", data_source=data_source, protocol=protocol)
print(
f"output={output_dir.relative_to(ROOT)} samples={len(archive['hr'])} "
f"data_source={data_source} protocol={protocol}"
)
if __name__ == "__main__":
main()
|