File size: 3,051 Bytes
355f250 | 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 | """Run SatMAE masked reconstruction inference."""
import argparse
import importlib.util
from pathlib import Path
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
def parse_args():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
parser.add_argument("--data", type=Path, default=None)
parser.add_argument("--checkpoint", type=Path, default=None)
parser.add_argument("--output-dir", type=Path, default=None)
parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
parser.add_argument("--mask-ratio", type=float, default=None)
return parser.parse_args()
def main():
args = parse_args()
config = yaml.safe_load(args.config.read_text())
spec = importlib.util.spec_from_file_location("satmae", ROOT / "model/satmae.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
model_args = {
key: value for key, value in config["model"].items()
if key not in {"architecture", "runtime_profile"}
}
model = module.SatMAE(**model_args)
checkpoint_path = args.checkpoint or ROOT / config["paths"]["checkpoint"]
if not checkpoint_path.exists():
raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}")
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
model.load_state_dict(checkpoint["model"])
use_cuda = torch.cuda.is_available() and args.device != "cpu"
if args.device == "cuda" and not torch.cuda.is_available():
raise RuntimeError("CUDA was requested but is 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)
images = torch.from_numpy(archive["images"]).to(device)
timestamps = None
if "timestamps" in archive:
timestamps = torch.from_numpy(archive["timestamps"]).to(device)
with torch.inference_mode():
output = model(images, timestamps=timestamps, mask_ratio=args.mask_ratio)
output_dir = args.output_dir or ROOT / config["paths"]["inference_dir"]
output_dir.mkdir(parents=True, exist_ok=True)
payload = {
"target": output["target"].cpu().numpy(),
"prediction": output["prediction"].cpu().numpy(),
"mask": output["mask"].cpu().numpy(),
"labels": archive["labels"],
}
if timestamps is not None:
payload["timestamps"] = timestamps.cpu().numpy()
for index, (prediction, target) in enumerate(zip(
output["group_predictions"], output["group_targets"]
)):
payload[f"prediction_group_{index}"] = prediction.cpu().numpy()
payload[f"target_group_{index}"] = target.cpu().numpy()
np.savez_compressed(output_dir / "reconstruction.npz", **payload)
print("inference=", output_dir / "reconstruction.npz")
if __name__ == "__main__":
main()
|