File size: 5,329 Bytes
1d4cac8 | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """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()
|