File size: 963 Bytes
29faa9d | 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 | """Real-artifact sanity. Skipped unless the checkpoint is present.
When the user drops their trained best_model.pt at the canonical path, this
test runs automatically — catches class-order or input-shape drift.
"""
from __future__ import annotations
from pathlib import Path
import numpy as np
import pytest
from PIL import Image
from src.models import mri_dl_2d
REAL_CKPT = Path("data/processed/mri_dl_2d/best_model.pt")
@pytest.mark.skipif(not REAL_CKPT.exists(), reason="real MRI checkpoint not present")
def test_real_checkpoint_loads_and_predicts(tmp_path):
model = mri_dl_2d.load(REAL_CKPT)
arr = (np.random.RandomState(0).rand(170, 170, 3) * 255).astype(np.uint8)
img = tmp_path / "scan.png"
Image.fromarray(arr).save(str(img))
result = mri_dl_2d.predict_image(model, img)
assert result["label_text"] in mri_dl_2d.CLASS_TO_IDX
s = sum(p["probability"] for p in result["probabilities"])
assert abs(s - 1.0) < 1e-5
|