scripts/smoke.py

"""Contract smoke test: restore the baked degraded image, check PSNR/SSIM improve."""
from __future__ import annotations

import json
import sys

from PIL import Image

from core.io import APP_TMP_DIR, register_protected
from core.processing import process


def main() -> int:
    tif = APP_TMP_DIR / "example.tif"
    if not tif.exists():
        import scripts.make_example as me
        me.main()
    register_protected(tif)
    register_protected(APP_TMP_DIR / ".example_clean.npy")
    src = str(tif)

    deg, res, clean, rep = process(src, engine="fast", psf_sigma=2.0, strength=1.0)
    assert clean is not None, "clean ground-truth sidecar not found"
    m = rep["metrics"]
    assert m is not None, "metrics not computed"
    assert m["psnr_restored_db"] > m["psnr_degraded_db"], ("restoration did not improve PSNR", m)
    assert m["ssim_restored"] > m["ssim_degraded"], ("restoration did not improve SSIM", m)
    assert res.shape == deg.shape

    from core.viz import triptych
    Image.fromarray(triptych(deg, res, clean)).save(APP_TMP_DIR / "smoke_care.png")
    print("SMOKE OK:", json.dumps({"model": rep["model"], **m}))
    return 0


if __name__ == "__main__":
    sys.exit(main())