"""The inference API: checkpoint plus image, no dataloader and no training config. None of this was possible before. Every path went through a config and a paired-mask dataloader, and ``EMTestDatasetMultiClass`` intersects image and mask filenames — so an image with no ground-truth mask produced *zero* samples. Preprocessing lived inside ``__getitem__``, so any inference path had to reimplement it, and drift between the two degrades predictions in a way that looks like a model problem. The test that matters most is :meth:`TestPreprocessing.test_matches_the_dataset_chain`: if the predictor preprocesses differently from training, everything downstream is quietly wrong. """ import shutil from pathlib import Path import numpy as np import pytest import torch from affmae.config import load_config from affmae.inference import AFFMAE, SegmentationResult from affmae.models.registry import get_model_spec REPO = Path(__file__).resolve().parents[1] CONFIG = "configs/aff_base_finetune_512_fpw.yaml" SMALL = 256 # a valid AFF geometry that is bearable on CPU @pytest.fixture(scope="module") def checkpoint(tmp_path_factory): """An untrained checkpoint plus its config, laid out as training would.""" directory = tmp_path_factory.mktemp("ckpt") cfg = load_config(CONFIG) cfg.img_size = SMALL model = get_model_spec(cfg.model_type).build_segmentation(cfg) path = directory / "best_model.pth" torch.save({"model_state_dict": model.state_dict(), "epoch": 0}, path) # The training scripts copy the config into the experiment directory, which # is what lets from_checkpoint work with no --config. text = (REPO / CONFIG).read_text().replace("img_size: 512", f"img_size: {SMALL}") (directory / "config.yaml").write_text(text) return path @pytest.fixture(scope="module") def predictor(checkpoint): return AFFMAE.from_checkpoint(str(checkpoint), device="cpu") @pytest.fixture def image(): return (np.random.RandomState(0).rand(300, 220) * 255).astype(np.uint8) class TestPreprocessing: def test_matches_the_dataset_chain(self, image): """Bit-identical to what training feeds the model. The dataset runs ToImage -> ConvertImageDtype(float32) -> Resize before normalizing; ConvertImageDtype is what maps uint8 0..255 into 0..1. Omitting it makes every value ~255x too large, which is exactly the kind of silent drift this test exists to prevent. """ from torchvision.transforms import v2 from affmae.data.finetune_dataset import IMAGE_MEAN, IMAGE_STD from affmae.data.preprocess import apply_clahe, preprocess_image mine = preprocess_image(image, img_size=64) array = apply_clahe(image)[..., None] expected = v2.Compose([ v2.ToImage(), v2.ConvertImageDtype(torch.float32), v2.Resize((64, 64), antialias=True, interpolation=v2.InterpolationMode.BILINEAR), ])(array).float().sub_(IMAGE_MEAN).div_(IMAGE_STD).unsqueeze(0) torch.testing.assert_close(mine, expected, rtol=0, atol=0) def test_output_shape_and_dtype(self, image): from affmae.data.preprocess import preprocess_image out = preprocess_image(image, img_size=128) assert out.shape == (1, 1, 128, 128) assert out.dtype == torch.float32 def test_handles_high_bit_depth(self): """Micrographs are often 12- or 16-bit; CLAHE needs 8-bit input.""" from affmae.data.preprocess import apply_clahe deep = (np.random.RandomState(1).rand(64, 64) * 65535).astype(np.uint16) assert apply_clahe(deep).dtype == np.uint8 def test_handles_rgb_by_converting_to_grey(self): from affmae.data.preprocess import apply_clahe rgb = (np.random.RandomState(2).rand(32, 32, 3) * 255).astype(np.uint8) assert apply_clahe(rgb).ndim == 2 def test_multichannel_mask_becomes_labels(self): from affmae.data.preprocess import multichannel_mask_to_labels mask = np.zeros((2, 8, 8)) mask[0, :4] = 255 mask[1, 4:] = 255 labels = multichannel_mask_to_labels(mask) assert labels[0, 0] == 1 and labels[7, 0] == 2 def test_missing_file_is_a_clear_error(self): from affmae.data.preprocess import load_image with pytest.raises(FileNotFoundError, match="image not found"): load_image("/nonexistent/image.tif") class TestFromCheckpoint: def test_finds_the_config_beside_the_checkpoint(self, predictor): assert predictor.img_size == SMALL assert predictor.num_classes >= 2 assert predictor.device.type == "cpu" def test_missing_checkpoint_is_a_clear_error(self): with pytest.raises(FileNotFoundError, match="checkpoint not found"): AFFMAE.from_checkpoint("/nonexistent/model.pth") def test_missing_config_names_what_to_pass(self, checkpoint, tmp_path): orphan = tmp_path / "alone.pth" shutil.copy(checkpoint, orphan) with pytest.raises(FileNotFoundError, match="no config given"): AFFMAE.from_checkpoint(str(orphan)) def test_unavailable_device_downgrades(self, checkpoint): """Asking for a device you do not have must warn, not crash.""" loaded = AFFMAE.from_checkpoint(str(checkpoint), device="mps") assert loaded.device.type in ("mps", "cuda", "cpu") class TestPredict: def test_accepts_a_numpy_array(self, predictor, image): result = predictor.segment(image) assert isinstance(result, SegmentationResult) assert result.labels.shape == (SMALL, SMALL) assert result.logits.shape[0] == predictor.num_classes def test_accepts_a_file_path(self, predictor, image, tmp_path): from PIL import Image path = tmp_path / "img.png" Image.fromarray(image).save(path) result = predictor.segment(str(path)) assert result.source.endswith("img.png") def test_accepts_a_tensor(self, predictor): result = predictor.segment(torch.rand(1, 64, 64) * 255) assert result.labels.shape == (SMALL, SMALL) def test_labels_are_in_range(self, predictor, image): result = predictor.segment(image) assert int(result.labels.min()) >= 0 assert int(result.labels.max()) < predictor.num_classes def test_class_pixel_counts_sums_to_the_image(self, predictor, image): result = predictor.segment(image) assert sum(result.class_pixel_counts.values()) == SMALL * SMALL def test_batch_preserves_order_and_length(self, predictor, image): other = np.flipud(image).copy() results = predictor.segment_batch([image, other]) assert len(results) == 2 assert not torch.equal(results[0].labels, results[1].labels) def test_is_deterministic(self, predictor, image): first = predictor.segment(image) second = predictor.segment(image) torch.testing.assert_close(first.logits, second.logits, rtol=0, atol=0) def test_does_not_leave_the_model_in_train_mode(self, predictor, image): """A renderer used to call model.train() on the way out.""" predictor.segment(image) assert not predictor.model.training class TestSaveOverlay: def test_writes_a_figure_without_ground_truth(self, predictor, image, tmp_path): out = predictor.segment(image).save_overlay(str(tmp_path / "o.png")) assert Path(out).stat().st_size > 0 def test_honours_a_viz_config(self, predictor, image, tmp_path): from affmae.viz import VizConfig result = predictor.segment(image) low = tmp_path / "low.png" high = tmp_path / "high.png" result.save_overlay(str(low), config=VizConfig(dpi=72)) result.save_overlay(str(high), config=VizConfig(dpi=200)) assert high.stat().st_size != low.stat().st_size class TestTokenLayout: def test_returns_positions_per_stage(self, predictor, image): rendered, positions = predictor.token_layout(image) assert rendered.shape == (1, SMALL, SMALL) assert len(positions) >= 2 assert all(p.shape[-1] == 2 for p in positions) def test_token_count_decreases_across_stages(self, predictor, image): """The point of adaptive downsampling.""" _, positions = predictor.token_layout(image) counts = [p.shape[0] for p in positions] assert counts == sorted(counts, reverse=True), counts assert counts[0] > counts[-1] class TestReconstruct: def test_declines_clearly_on_a_segmentation_checkpoint(self, predictor, image): """A finetuned model has replaced the MAE head; say so.""" with pytest.raises(NotImplementedError, match="pretraining checkpoint"): predictor.reconstruct(image) class TestDemoWiring: """The Gradio app is a thin shell over the predictor; check the seams.""" def test_interface_builds(self, predictor): pytest.importorskip("gradio") from affmae.demo import build_interface assert build_interface(predictor) is not None def test_segmentation_tab_returns_an_image_and_a_table(self, predictor, image): pytest.importorskip("gradio") from affmae.demo import _predictor_cache, _segment resolve = _predictor_cache(predictor, predictor.device) path, text = _segment(resolve, predictor.img_size, image, 100, 0.6) assert path and Path(path).stat().st_size > 0 assert "class" in text def test_segmentation_tab_reports_a_missing_resolution(self, predictor, image): """A resolution with no checkpoint on disk must explain itself, not raise. The picker only offers downloaded resolutions, but the API accepts any value, so the handler is the last line of defence. """ pytest.importorskip("gradio") from affmae.demo import _predictor_cache, _segment resolve = _predictor_cache(predictor, predictor.device) path, text = _segment(resolve, 999, image, 100, 0.6) assert path is None assert "999" in text def test_tabs_handle_no_upload(self, predictor): pytest.importorskip("gradio") from affmae.demo import (_predictor_cache, _reconstruct, _segment, _tokens) resolve = _predictor_cache(predictor, predictor.device) for handler in (_segment, _tokens): _, message = handler(resolve, predictor.img_size, None, 100, 0.6) assert "Upload" in message _, message = _reconstruct(predictor, None, 0.5) assert "Upload" in message def test_the_gpu_hook_wraps_every_model_handler(self, predictor): """ZeroGPU needs each handler wrapped; a missed one would run on CPU. The Space passes spaces.GPU in here. If a tab were wired with a bare lambda again, that tab would silently lose its GPU on the Space. """ pytest.importorskip("gradio") from affmae.demo import build_interface wrapped = [] def fake_gpu(fn): wrapped.append(fn.__name__) return fn assert build_interface(predictor, gpu=fake_gpu) is not None assert sorted(wrapped) == ["on_batch", "on_reconstruct", "on_segment", "on_tokens"], wrapped def test_the_gpu_hook_is_optional(self, predictor): """Local use passes nothing and must behave exactly as before.""" pytest.importorskip("gradio") from affmae.demo import build_interface assert build_interface(predictor) is not None def test_reconstruction_tab_degrades_gracefully(self, predictor, image): pytest.importorskip("gradio") from affmae.demo import _reconstruct path, message = _reconstruct(predictor, image, 0.5) assert path is None assert "pretraining" in message def test_batch_tab_returns_a_zip(self, predictor, image, tmp_path): pytest.importorskip("gradio") import zipfile from PIL import Image from affmae.demo import _batch paths = [] for index in range(2): path = tmp_path / f"i{index}.png" Image.fromarray(image).save(path) paths.append(str(path)) archive, message = _batch(predictor, paths, 100, 0.6) assert archive and zipfile.is_zipfile(archive) with zipfile.ZipFile(archive) as bundle: names = bundle.namelist() assert len(names) == 4, names # mask + overlay per image assert "2 image(s)" in message def test_batch_tab_handles_no_files(self, predictor): pytest.importorskip("gradio") from affmae.demo import _batch archive, message = _batch(predictor, [], 100, 0.6) assert archive is None and "Upload" in message class TestModes: """Inference skips backward-only work; training modes retain gradients.""" def test_defaults_to_inference(self, predictor): from affmae.ops.policy import Mode assert predictor.mode is Mode.INFERENCE def test_inference_mode_freezes_parameters(self, checkpoint): model = AFFMAE.from_checkpoint(str(checkpoint), device="cpu", mode="inference") assert not any(p.requires_grad for p in model.model.parameters()) def test_training_modes_keep_parameters_trainable(self, checkpoint): for mode in ("finetune", "pretrain"): model = AFFMAE.from_checkpoint(str(checkpoint), device="cpu", mode=mode) assert any(p.requires_grad for p in model.model.parameters()), mode def test_inference_mode_refuses_a_gradient_with_a_useful_message(self, checkpoint): from affmae.ops.policy import InferenceOnlyError model = AFFMAE.from_checkpoint(str(checkpoint), device="cpu", mode="inference") images = torch.randn(1, 1, SMALL, SMALL, requires_grad=True) with pytest.raises(InferenceOnlyError, match="mode='finetune'"): model(images) def test_finetune_mode_allows_a_backward_pass(self, checkpoint): model = AFFMAE.from_checkpoint(str(checkpoint), device="cpu", mode="finetune") images = torch.randn(1, 1, SMALL, SMALL) outputs = model(images) logits = outputs[-1] if isinstance(outputs, (list, tuple)) else outputs logits.sum().backward() assert any(p.grad is not None for p in model.model.parameters()) def test_rejects_an_unknown_mode(self, checkpoint): with pytest.raises(ValueError, match="unknown mode"): AFFMAE.from_checkpoint(str(checkpoint), device="cpu", mode="fast") def test_mode_policy_does_not_select_component_backends(self): from affmae.config import load_config from affmae.ops.policy import KernelPolicy, Mode for mode in Mode: cfg = load_config(str(REPO / CONFIG)) before = (cfg.decoder_deform_backend, cfg.cluster_attention_backend) KernelPolicy.for_mode(mode).apply_to_config(cfg) after = (cfg.decoder_deform_backend, cfg.cluster_attention_backend) assert before == after, f"{mode} changed the algorithm: {before} -> {after}" @pytest.mark.parametrize("alias", ["auto", "fused", "csr_cached", "csr_knn_cached"]) def test_decoder_fused_aliases(self, checkpoint, alias): loaded = AFFMAE.from_checkpoint( str(checkpoint), device="cpu", decoder_deform_backend=alias) decoder = loaded.model.cross_attention_decoder assert decoder.deform_backend == "csr_knn_cached" def test_backend_arguments_override_the_yaml(self, checkpoint): loaded = AFFMAE.from_checkpoint( str(checkpoint), device="cpu", cluster_attention_backend="torch", decoder_deform_backend="unfused") first_attention = loaded.model.encoder.layers[0].blocks[0].attn assert first_attention.backend == "torch" assert loaded.model.cross_attention_decoder.deform_backend == "unfused" def test_rejects_unknown_component_backends(self, checkpoint): with pytest.raises(ValueError, match="neighbourhood-attention backend"): AFFMAE.from_checkpoint( str(checkpoint), device="cpu", cluster_attention_backend="mystery") with pytest.raises(ValueError, match="decoder deform backend"): AFFMAE.from_checkpoint( str(checkpoint), device="cpu", decoder_deform_backend="mystery") def test_inference_enables_the_free_knn_cache(self): """Caching the KNN table is bit-identical, so there is no reason not to.""" from affmae.config import load_config from affmae.ops.policy import KernelPolicy, Mode cfg = load_config(str(REPO / CONFIG)) KernelPolicy.for_mode(Mode.INFERENCE).apply_to_config(cfg) assert cfg.decoder_knn_cache is True def test_capabilities_are_introspectable(self, predictor): """Ask, rather than call and catch.""" assert "segment" in predictor.capabilities assert predictor.can_segment is True assert isinstance(predictor.can_reconstruct, bool) def test_from_model_wraps_without_copying(self, predictor): wrapped = AFFMAE.from_model(predictor.model, img_size=SMALL, num_classes=predictor.num_classes, mode="inference") assert wrapped.model is predictor.model PRETRAIN_CONFIG = "configs/aff_base_pretrain_0.4ds_0.5mask_last_local.yaml" @pytest.fixture(scope="module") def pretrain_checkpoint(tmp_path_factory): """An untrained *pretraining* checkpoint plus its config.""" directory = tmp_path_factory.mktemp("mae_ckpt") cfg = load_config(PRETRAIN_CONFIG) cfg.img_size = SMALL model = get_model_spec(cfg.model_type).build_pretrain(cfg) path = directory / "ckpt.pth" torch.save({"model_state_dict": model.state_dict(), "epoch": 0}, path) text = (REPO / PRETRAIN_CONFIG).read_text().replace( "img_size: 512", f"img_size: {SMALL}") (directory / "config.yaml").write_text(text) return path class TestPretrainingCheckpoints: """``from_checkpoint`` used to hardcode ``build_segmentation``. So :meth:`AFFMAE.reconstruct` was unreachable from a checkpoint even though :attr:`capabilities` advertised it -- the only way in was ``from_model``. """ def test_a_pretraining_checkpoint_can_reconstruct(self, pretrain_checkpoint): model = AFFMAE.from_checkpoint(str(pretrain_checkpoint), device="cpu") assert model.can_reconstruct assert not model.can_segment def test_task_can_be_forced(self, pretrain_checkpoint): model = AFFMAE.from_checkpoint(str(pretrain_checkpoint), device="cpu", task="pretrain") assert model.can_reconstruct def test_an_unknown_task_is_rejected(self, pretrain_checkpoint): with pytest.raises(ValueError, match="auto|segmentation|pretrain"): AFFMAE.from_checkpoint(str(pretrain_checkpoint), device="cpu", task="reconstruction") def test_reconstruction_keeps_the_visible_patches_exactly( self, pretrain_checkpoint, image): """The point of the composite: only the masked region is model output. A figure that quietly re-renders the *whole* image from the decoder would overstate reconstruction quality, so pin the visible half to the input. """ from affmae.utils.dist import unwrap_model predictor = AFFMAE.from_checkpoint(str(pretrain_checkpoint), device="cpu") result = predictor.reconstruct(image, mask_ratio=0.5) model = unwrap_model(predictor.model) original = model.patchify(result.original[None]) composite = model.patchify(result.reconstructions[-1][None]) blanked = model.patchify(result.masked[None]) visible = blanked.abs().sum(-1) > 0 assert 0 < int(visible.sum()) < visible.numel() assert torch.equal(original[visible], composite[visible]) def test_masked_token_layout_is_sparser_than_dense(self, pretrain_checkpoint, image): """``mask_ratio`` shows the tokens the encoder actually received.""" predictor = AFFMAE.from_checkpoint(str(pretrain_checkpoint), device="cpu") _, dense = predictor.token_layout(image) _, sparse = predictor.token_layout(image, mask_ratio=0.5) assert len(dense) == len(sparse) for stage, (d, s) in enumerate(zip(dense, sparse)): assert s.shape[0] < d.shape[0], f"stage {stage} was not sparsened" def test_the_mask_comes_from_the_model_not_a_copy(self, pretrain_checkpoint, image): """Layout and reconstruction share ``mask_and_embed``. They were separate code paths, so a figure could show tokens from one mask beside a reconstruction from another. """ from affmae.utils.dist import unwrap_model predictor = AFFMAE.from_checkpoint(str(pretrain_checkpoint), device="cpu") model = unwrap_model(predictor.model) assert hasattr(model, "mask_and_embed") tensor, _ = predictor._prepare(image) embedded = model.mask_and_embed(tensor) kept = embedded["ids_keep"].shape[1] hidden = embedded["ids_masked"].shape[1] assert kept + hidden == embedded["img_patches"].shape[1] def test_mask_ratio_is_rejected_for_a_segmentation_checkpoint(self, predictor, image): with pytest.raises(NotImplementedError, match="mask_ratio"): predictor.token_layout(image, mask_ratio=0.5) def test_reconstruction_tokens_sit_on_visible_patches(self, pretrain_checkpoint, image): """The tokens and the masked image must come from the same mask. ``render_examples.py`` used to take the reconstruction from ``reconstruct()`` and the tokens from a second ``token_layout(mask_ratio=...)`` call. Each draws its own Perlin mask, so the figure showed one mask's tokens over another mask's image: measured, only ~50% of tokens landed on a patch the image showed as visible, which is chance. Taking both from one result makes it 100%. """ from affmae.utils.dist import unwrap_model predictor = AFFMAE.from_checkpoint(str(pretrain_checkpoint), device="cpu") model = unwrap_model(predictor.model) result = predictor.reconstruct(image, mask_ratio=0.5) assert result.locations, "no per-stage locations to check" visible = model.patchify(result.masked[None])[0].abs().sum(-1) > 0 grid = SMALL // model.encoder_patch_size for name, positions in zip(result.stage_names, result.locations): linear = (positions[:, 1].round().long().clamp(0, grid - 1) * grid + positions[:, 0].round().long().clamp(0, grid - 1)) on_visible = visible[linear].float().mean().item() assert on_visible == 1.0, ( f"{name}: only {on_visible:.1%} of tokens are on visible " f"patches; the tokens and the masked image disagree.") def test_a_second_token_layout_call_draws_a_different_mask(self, pretrain_checkpoint, image): """Pin the footgun itself, so the docstring warning stays true.""" predictor = AFFMAE.from_checkpoint(str(pretrain_checkpoint), device="cpu") _, first = predictor.token_layout(image, mask_ratio=0.5) _, second = predictor.token_layout(image, mask_ratio=0.5) assert not torch.equal(first[0], second[0]), ( "two token_layout calls produced the same mask; if masking became " "deterministic, the warning on token_layout should be removed.")