| """Run SkySense inference and save arrays for evaluation.""" |
|
|
| import importlib.util |
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def load_model_class(): |
| spec = importlib.util.spec_from_file_location("skysense_model", ROOT / "model" / "skysense.py") |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module.SkySense |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Run batched SkySense segmentation inference") |
| parser.add_argument("--batch-size", type=int) |
| args = parser.parse_args() |
| with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle: |
| config = yaml.safe_load(handle) |
| checkpoint_path = ROOT / config["paths"]["checkpoint"] |
| if not checkpoint_path.exists(): |
| raise FileNotFoundError( |
| f"Missing checkpoint: {checkpoint_path.relative_to(ROOT)}. " |
| "Run `python scripts/train.py` first." |
| ) |
| use_accelerator = torch.cuda.is_available() and config["runtime"].get("device", "auto") != "cpu" |
| device = torch.device("cuda" if use_accelerator else "cpu") |
| checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| SkySense = load_model_class() |
| model = SkySense( |
| **config["model"], |
| hr_channels=config["data"]["hr_channels"], |
| s2_channels=config["data"]["s2_channels"], |
| s1_channels=config["data"]["s1_channels"], |
| num_classes=config["data"]["num_classes"], |
| ).to(device) |
| model.load_state_dict(checkpoint["model"]) |
| model.eval() |
| test_path = ROOT / config["data"]["root"] / "test.npz" |
| if not test_path.exists(): |
| raise FileNotFoundError( |
| f"Missing inference data: {test_path.relative_to(ROOT)}. " |
| "Run `python scripts/fake_data.py` first." |
| ) |
| archive = np.load(test_path) |
| keys = ["hr", "s2", "s1", "dates_hr", "dates_s2", "dates_s1", "region"] |
| arrays = {key: archive[key] for key in keys} |
| expected = { |
| "hr": (config["data"]["hr_timesteps"], config["data"]["hr_channels"], config["data"]["hr_size"], config["data"]["hr_size"]), |
| "s2": (config["data"]["s2_timesteps"], config["data"]["s2_channels"], config["data"]["s2_size"], config["data"]["s2_size"]), |
| "s1": (config["data"]["s1_timesteps"], config["data"]["s1_channels"], config["data"]["s1_size"], config["data"]["s1_size"]), |
| "dates_hr": (config["data"]["hr_timesteps"],), |
| "dates_s2": (config["data"]["s2_timesteps"],), |
| "dates_s1": (config["data"]["s1_timesteps"],), |
| "region": (), |
| } |
| sample_count = len(arrays["hr"]) |
| for key, shape in expected.items(): |
| if len(arrays[key]) != sample_count or tuple(arrays[key].shape[1:]) != shape: |
| raise ValueError(f"Invalid test {key} shape {arrays[key].shape}; expected [N,{','.join(map(str, shape))}]") |
| for key in ("hr", "s2", "s1"): |
| if not np.issubdtype(arrays[key].dtype, np.floating): |
| raise TypeError(f"{key} must use a floating dtype") |
| for key in ("dates_hr", "dates_s2", "dates_s1", "region"): |
| if arrays[key].dtype != np.int64: |
| raise TypeError(f"{key} must use int64") |
| if any(np.any((arrays[key] < 0) | (arrays[key] > 364)) for key in ("dates_hr", "dates_s2", "dates_s1")): |
| raise ValueError("Test dates must be in [0, 364]") |
| if np.any((arrays["region"] < 0) | (arrays["region"] >= config["model"]["num_regions"])): |
| raise ValueError("Test region IDs are out of range") |
| labels = archive["labels"] |
| if labels.dtype != np.int64 or labels.shape != (sample_count, config["data"]["hr_size"], config["data"]["hr_size"]): |
| raise ValueError("Test labels must be int64 [N,hr_size,hr_size]") |
| batch_size = args.batch_size or config["train"]["batch_size"] |
| predictions = [] |
| all_probabilities = [] |
| with torch.inference_mode(): |
| for start in range(0, len(arrays["hr"]), batch_size): |
| tensors = {key: torch.from_numpy(value[start:start + batch_size]).to(device) |
| for key, value in arrays.items()} |
| output = model(tensors["hr"], tensors["s2"], tensors["s1"], tensors["dates_hr"], tensors["dates_s2"], tensors["dates_s1"], tensors["region"]) |
| probabilities = output["logits"].softmax(dim=1).cpu().numpy() |
| all_probabilities.append(probabilities) |
| predictions.append(probabilities.argmax(axis=1)) |
| output_dir = ROOT / config["paths"]["inference_dir"] |
| output_dir.mkdir(parents=True, exist_ok=True) |
| np.save(output_dir / "predictions.npy", np.concatenate(predictions)) |
| np.save(output_dir / "probabilities.npy", np.concatenate(all_probabilities)) |
| np.save(output_dir / "targets.npy", labels) |
| data_source = str(archive["data_source"]) if "data_source" in archive.files else "unknown" |
| protocol = str(archive["protocol"]) if "protocol" in archive.files else "unknown" |
| np.savez(output_dir / "metadata.npz", data_source=data_source, protocol=protocol) |
| print( |
| f"output={output_dir.relative_to(ROOT)} samples={len(archive['hr'])} " |
| f"data_source={data_source} protocol={protocol}" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
