| """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() |
|
|
