File size: 4,096 Bytes
387a20d 7a2d30b 387a20d 7a2d30b 387a20d 7a2d30b 387a20d 7a2d30b 387a20d 7a2d30b 387a20d 7a2d30b 387a20d 7a2d30b 387a20d | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | 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()
|