Spaces:
Running on Zero
Running on Zero
| 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) | |
| 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 | |
| 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,)) | |