Spaces:
Sleeping
Sleeping
File size: 3,422 Bytes
5d31c67 4a5c225 | 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 | import os
import numpy as np
import pytest
import torch
from cellquant.model import CellUNet, N_OUT, load_checkpoint, pad_to_multiple, save_checkpoint
WEIGHTS = os.environ.get("CELLQUANT_WEIGHTS", os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "weights", "cellquant.pt"))
def test_pad_to_multiple():
x = torch.zeros(1, 3, 70, 100)
xp, (ph, pw) = pad_to_multiple(x, 32)
assert xp.shape[-2:] == (96, 128) and (ph, pw) == (26, 28)
x2, (ph2, pw2) = pad_to_multiple(torch.zeros(1, 3, 64, 64), 32)
assert x2.shape[-2:] == (64, 64) and (ph2, pw2) == (0, 0)
def test_model_forward_and_checkpoint_roundtrip(tmp_path):
m = CellUNet("resnet34", pretrained=False).eval()
x = torch.rand(1, 3, 64, 96)
with torch.no_grad():
y = m(x)
assert y.shape == (1, N_OUT, 64, 96)
p = tmp_path / "m.pt"
save_checkpoint(m, str(p), {"note": "test"})
m2, meta = load_checkpoint(str(p))
assert meta["note"] == "test" and meta["encoder"] == "resnet34"
with torch.no_grad():
assert torch.allclose(m2(x), y, atol=1e-5)
@pytest.mark.skipif(not os.path.exists(WEIGHTS), reason="trained weights not present")
def test_predictor_deterministic_and_tiling_consistent():
from cellquant.infer import Predictor
rng = np.random.default_rng(0)
img = (rng.random((300, 420, 3)) * 60).astype(np.uint8)
p = Predictor(WEIGHTS, device="cpu", tta=1)
a = p.probs(img)
b = p.probs(img)
assert all(np.array_equal(a[k], b[k]) for k in a)
assert a["cell"].shape == (3, 300, 420)
# force tiling and compare with the whole-image pass
pt = Predictor(WEIGHTS, device="cpu", tta=1, tile=256, overlap=64, max_whole=128)
c = pt.probs(img)
diff = np.abs(c["cell"] - a["cell"]).mean()
assert diff < 0.05, diff
@pytest.mark.skipif(not os.path.exists(WEIGHTS), reason="trained weights not present")
def test_tta_outputs_valid_probabilities():
from cellquant.infer import Predictor
img = (np.random.default_rng(1).random((128, 160, 3)) * 60).astype(np.uint8)
for tta in (4, 8):
pr = Predictor(WEIGHTS, device="cpu", tta=tta).probs(img)
assert np.allclose(pr["cell"].sum(0), 1, atol=1e-4) and np.allclose(pr["nos"].sum(0), 1, atol=1e-4)
def test_multiscale_tta_keeps_shapes_and_matches_single_scale():
from cellquant.model import CellUNet
from cellquant.infer import Predictor
import numpy as np
img = (np.random.default_rng(3).random((70, 90, 3)) * 120).astype(np.uint8)
m = CellUNet("resnet34", pretrained=False).eval()
base = Predictor(m, device="cpu", tta=1).probs(img)
same = Predictor(m, device="cpu", tta=1, scales=(1.0,)).probs(img)
assert all(np.array_equal(base[k], same[k]) for k in base)
ms = Predictor(m, device="cpu", tta=4, scales=(0.8, 1.0, 1.25)).probs(img)
assert set(ms) == set(base) and all(ms[k].shape == base[k].shape for k in base)
assert np.allclose(ms["cell"].sum(0), 1, atol=1e-4) and 0 <= ms["dist"].min() and ms["dist"].max() <= 1
assert not np.array_equal(ms["cell"], base["cell"]) # the scales changed the average
# the tiled path takes the scales too
tiled = Predictor(m, device="cpu", tta=1, scales=(0.8, 1.0), tile=64, overlap=16, max_whole=32).probs(img)
assert tiled["dist"].shape == (70, 90)
import pytest
with pytest.raises(AssertionError):
Predictor(m, device="cpu", scales=(0.1,))
|