| |
| from __future__ import annotations |
| import argparse |
| from pathlib import Path |
| import sys |
| import numpy as np |
| import torch |
| import yaml |
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(PROJECT_ROOT)) |
| from scripts.data_loader import MRMSDataset |
| from model.model_factory import build_model, load_checkpoint |
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="NowcastNet inference") |
| parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml")) |
| parser.add_argument("--data-dir", help="MRMS event directory; defaults to data.data_dir in config.yaml") |
| parser.add_argument("--checkpoint") |
| parser.add_argument("--output-dir") |
| parser.add_argument("--device", default="auto") |
| parser.add_argument("--height", type=int) |
| parser.add_argument("--width", type=int) |
| parser.add_argument("--ngf", type=int) |
| args = parser.parse_args() |
| cfg = yaml.safe_load(Path(args.config).read_text()) |
| mc, dc, ic = cfg["model"], cfg["data"], cfg["inference"] |
| height, width = args.height or dc["image_height"], args.width or dc["image_width"] |
| if args.ngf: mc["ngf"] = args.ngf |
| mc["img_height"], mc["img_width"] = height, width |
| device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available() else ("cpu" if args.device == "auto" else args.device)) |
| checkpoint = Path(args.checkpoint) if args.checkpoint else PROJECT_ROOT / mc["checkpoint_dir"] / f"{mc.get('checkpoint_prefix', 'model_bak')}.pth" |
| if not checkpoint.is_file(): |
| raise FileNotFoundError(f"Checkpoint not found: {checkpoint}. Run scripts/train.py first or pass --checkpoint.") |
| model = build_model(mc, device); load_checkpoint(model, checkpoint, device); model.eval() |
| data_dir = Path(args.data_dir) if args.data_dir else PROJECT_ROOT / dc["data_dir"] |
| ds = MRMSDataset(data_dir, height, width, dc["total_length"], "test") |
| out = Path(args.output_dir) if args.output_dir else PROJECT_ROOT / ic["output_dir"] |
| out.mkdir(parents=True, exist_ok=True) |
| with torch.no_grad(): |
| for item in ds: |
| frames = item["radar_frames"].unsqueeze(0).to(device=device, dtype=torch.float32) |
| pred = model(frames).squeeze(0).squeeze(-1).cpu().numpy() |
| target = frames[0, mc["input_length"]:, :, :, 0].cpu().numpy() |
| np.save(out / f"{item['event']}_pred.npy", pred) |
| np.save(out / f"{item['event']}_input.npy", frames[0, :mc["input_length"], :, :, 0].cpu().numpy()) |
| np.save(out / f"{item['event']}_target.npy", target) |
| print(item["event"], pred.shape, float(pred.min()), float(pred.max())) |
|
|
| if __name__ == "__main__": main() |
|
|
