| """Reconstruct validated sensor NPZ files in bounded device batches.""" |
|
|
| import importlib.util |
| 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("dofa_model", ROOT / "model" / "dofa.py") |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module.DOFA |
|
|
|
|
| def expected_wavelengths(modality): |
| if modality.get("wavelength_mode") == "synthetic_uniform": |
| return modality["wavelength_start"] + np.arange(modality["channels"], dtype=np.float32) * modality["wavelength_step"] |
| return np.asarray(modality["wavelengths"], dtype=np.float32) |
|
|
|
|
| def scalar(archive, key, default=None): |
| if key not in archive: |
| if default is not None: |
| return default |
| raise ValueError(f"NPZ is missing required metadata: {key}") |
| if archive[key].ndim != 0: |
| raise ValueError(f"NPZ metadata {key} must be a scalar") |
| return archive[key].item() |
|
|
|
|
| def validate_archive(archive, path, name, data_config): |
| if "images" not in archive or "wavelengths" not in archive: |
| raise ValueError(f"{path}: images and wavelengths are required") |
| images, wavelengths = archive["images"], archive["wavelengths"] |
| modality_config = data_config["modalities"][name] |
| if images.ndim != 4 or images.dtype != np.float32: |
| raise ValueError(f"{path}: images must be float32 NCHW") |
| expected_shape = (modality_config["channels"], data_config["image_size"], data_config["image_size"]) |
| if images.shape[1:] != expected_shape: |
| raise ValueError(f"{path}: expected [N,{expected_shape[0]},224,224], got {images.shape}") |
| expected = expected_wavelengths(modality_config) |
| if wavelengths.shape != (modality_config["channels"],) or not np.issubdtype(wavelengths.dtype, np.floating) or not np.isfinite(wavelengths).all() or not np.allclose(wavelengths, expected, rtol=1e-5, atol=1e-6): |
| raise ValueError(f"{path}: wavelengths do not match configured sensor wavelengths") |
| modality = str(scalar(archive, "modality")) |
| protocol = str(scalar(archive, "protocol")) |
| source = str(scalar(archive, "data_source", "unknown")) |
| if modality != name or protocol != data_config["protocol"]: |
| raise ValueError(f"{path}: modality/protocol metadata does not match config") |
| data_range = float(scalar(archive, "data_range", data_config.get("data_range"))) |
| if not np.isfinite(data_range) or data_range <= 0: |
| raise ValueError(f"{path}: PSNR requires a positive data_range metadata or config value") |
| return images, wavelengths.astype("float32", copy=False), protocol, source, data_range |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text(encoding="utf-8")) |
| checkpoint_path = ROOT / config["paths"]["checkpoint"] |
| if not checkpoint_path.exists(): |
| raise FileNotFoundError("Missing checkpoint. Run `python scripts/train.py` first.") |
| device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu") |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| if checkpoint.get("protocol") != config["data"]["protocol"]: |
| raise ValueError("Checkpoint and configured protocols do not match") |
| model = load_model_class()(**config["model"]).to(device) |
| model.load_state_dict(checkpoint["model"]) |
| model.eval() |
| output_dir = ROOT / config["paths"]["inference_dir"] |
| output_dir.mkdir(parents=True, exist_ok=True) |
| data_root = ROOT / config["data"]["root"] |
| batch_size = config["runtime"]["inference_batch_size"] |
| torch.manual_seed(config["seed"]) |
| for modality in config["data"]["modalities"]: |
| path = data_root / f"test_{modality}.npz" |
| if not path.exists(): |
| raise FileNotFoundError(f"Missing test data: {path.relative_to(ROOT)}") |
| archive = np.load(path) |
| images, wavelengths, protocol, source, data_range = validate_archive( |
| archive, path, modality, config["data"]) |
| reconstructions = np.empty_like(images) |
| masks = np.empty((len(images), model.num_patches), dtype=bool) |
| wavelength_tensor = torch.from_numpy(wavelengths).to(device) |
| with torch.inference_mode(): |
| for start in range(0, len(images), batch_size): |
| stop = min(start + batch_size, len(images)) |
| batch = torch.from_numpy(images[start:stop]).to(device) |
| output = model(batch, wavelength_tensor) |
| reconstructions[start:stop] = output["reconstruction"].cpu().numpy() |
| masks[start:stop] = output["mask"].cpu().numpy() |
| del batch, output |
| target = output_dir / f"{modality}_reconstruction.npz" |
| np.savez_compressed(target, inputs=images, reconstructions=reconstructions, masks=masks, |
| wavelengths=wavelengths, modality=np.asarray(modality), |
| data_source=np.asarray(source), protocol=np.asarray(protocol), |
| data_range=np.asarray(data_range, dtype=np.float32)) |
| print(f"output={target.relative_to(ROOT)} channels={images.shape[1]} batch_size={batch_size}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
