File size: 2,820 Bytes
aed6f6f cf7a05c aed6f6f cf7a05c 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 | """Run checkpoint-backed SatlasNet inference on a test NPZ."""
import argparse
import importlib.util
from pathlib import Path
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml"); parser.add_argument("--data", type=Path)
parser.add_argument("--checkpoint", type=Path); parser.add_argument("--output-dir", type=Path)
parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto"); args = parser.parse_args()
config = yaml.safe_load(args.config.read_text()); checkpoint_path = args.checkpoint or ROOT / config["paths"]["checkpoint"]
if not checkpoint_path.is_file(): raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}")
spec = importlib.util.spec_from_file_location("satlaspretrain", ROOT / "model/satlaspretrain.py")
module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module)
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
model = module.SatlasPretrain(**config["model"]); model.load_state_dict(checkpoint["model"])
use_cuda = torch.cuda.is_available() and args.device != "cpu"
if args.device == "cuda" and not use_cuda: raise RuntimeError("CUDA requested but unavailable")
device = torch.device("cuda" if use_cuda else "cpu"); model.to(device).eval()
data_path = args.data or ROOT / config["data"]["root"] / "test.npz"; archive = np.load(data_path)
module.validate_npz(archive, config)
with torch.inference_mode():
outputs = model(torch.from_numpy(archive["highres_images"]).to(device),
torch.from_numpy(archive["lowres_images"]).to(device),
torch.from_numpy(archive["valid_highres_times"]).to(device),
torch.from_numpy(archive["valid_lowres_times"]).to(device))
predictions = {}
for name, value in outputs.items():
if name in ("segmentation", "property", "classification"): value = value.softmax(1)
elif name != "regression": value = value.sigmoid()
predictions[name] = value.cpu().numpy().astype(np.float32)
output = args.output_dir or ROOT / config["paths"]["inference_dir"]; output.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output / "predictions.npz", **predictions, checkpoint=np.asarray(str(checkpoint_path)),
sample_ids=archive["sample_ids"],
source=archive["source"],
protocol=archive["protocol"] if "protocol" in archive else np.asarray("provided_npz"))
print(f"inference={output / 'predictions.npz'} checkpoint={checkpoint_path}")
if __name__ == "__main__": main()
|