| import argparse |
| from pathlib import Path |
| import sys |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from model.spectralgpt import SpectralGPT |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Run SpectralGPT reconstruction") |
| parser.add_argument("--config", default="conf/config.yaml") |
| parser.add_argument("--checkpoint") |
| parser.add_argument("--batch-size", type=int) |
| args = parser.parse_args() |
| with open(args.config, encoding="utf-8") as handle: |
| config = yaml.safe_load(handle) |
| requested = config["runtime"]["device"] |
| device = torch.device("cuda" if torch.cuda.is_available() and requested != "cpu" else "cpu") |
| torch.manual_seed(config["runtime"]["seed"]) |
| stage = config["stages"][-1] |
| model = SpectralGPT(image_size=stage["image_size"], **config["model"]).to(device) |
| checkpoint_path = args.checkpoint or config["training"]["checkpoint"] |
| if not Path(checkpoint_path).exists(): |
| raise FileNotFoundError( |
| f"Missing checkpoint: {checkpoint_path}. Run `python scripts/train.py` first." |
| ) |
| checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| if checkpoint.get("stage") != stage["name"] or checkpoint.get("image_size") != stage["image_size"]: |
| raise ValueError("Checkpoint is not the configured final stage2 checkpoint") |
| model.load_state_dict(checkpoint["model"]) |
| model.eval() |
| data_path = Path(config["data"]["test_path"]) |
| if not data_path.exists(): |
| raise FileNotFoundError( |
| f"Missing inference data: {data_path}. Run `python scripts/fake_data.py` first." |
| ) |
| with np.load(data_path) as data: |
| images = data["images"].copy() |
| data_source = str(data["data_source"]) if "data_source" in data.files else "unknown" |
| protocol = str(data["protocol"]) if "protocol" in data.files else "unknown" |
| normalization = str(data["normalization"]) if "normalization" in data.files else "unknown" |
| scale_factors = data["scale_factors"].copy() if "scale_factors" in data.files else np.ones(len(images), np.float32) |
| stored_stage = str(data["stage"]) if "stage" in data.files else "unknown" |
| expected = (config["model"]["in_channels"], stage["image_size"], stage["image_size"]) |
| if images.dtype != np.float32 or images.ndim != 4 or tuple(images.shape[1:]) != expected: |
| raise ValueError(f"Expected float32 stage2 test [N,{','.join(map(str, expected))}], got {images.dtype} {images.shape}") |
| if stored_stage != stage["name"]: |
| raise ValueError(f"Expected test stage {stage['name']}, got {stored_stage}") |
| batch_size = args.batch_size or config["training"]["batch_size"] |
| collected = {name: [] for name in ("reconstruction", "prediction_image", "mask", "mask_image")} |
| with torch.inference_mode(): |
| for start in range(0, len(images), batch_size): |
| batch = torch.from_numpy(images[start:start + batch_size]).to(device) |
| output = model(batch) |
| for name in collected: |
| collected[name].append(output[name].cpu().numpy()) |
| output_dir = Path(config["runtime"]["output_dir"]) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output_dir / "reconstruction.npz", |
| inputs=images, |
| reconstructions=np.concatenate(collected["reconstruction"]), |
| predictions=np.concatenate(collected["prediction_image"]), |
| masks=np.concatenate(collected["mask"]), |
| mask_images=np.concatenate(collected["mask_image"]), |
| data_source=np.asarray(data_source), |
| protocol=np.asarray(protocol), normalization=np.asarray(normalization), |
| scale_factors=scale_factors, stage=np.asarray(stored_stage)) |
| print( |
| f"saved: {output_dir / 'reconstruction.npz'} " |
| f"data_source={data_source} protocol={protocol}" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
