File size: 2,707 Bytes
439c523 | 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 | #!/usr/bin/env python3
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()
|