| """Tests for local image embedding helpers.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import cast |
|
|
| import pytest |
| from PIL import Image |
|
|
| from mathvision_explorer.embeddings import ( |
| ColorStatsEmbedder, |
| IJepaImageEmbedder, |
| JepaDependencyError, |
| MissingImageError, |
| PatchInterestMap, |
| _mean_pool_features, |
| embed_record_image, |
| render_patch_interest_heatmap, |
| render_patch_interest_overlay, |
| ) |
|
|
|
|
| def test_color_stats_embedder_returns_six_normalized_features(tmp_path: Path) -> None: |
| """The baseline image embedder returns RGB means and standard deviations.""" |
|
|
| image_path = tmp_path / "red.png" |
| Image.new("RGB", (4, 4), color=(255, 0, 0)).save(image_path) |
|
|
| vector = ColorStatsEmbedder().embed_image(image_path) |
|
|
| assert vector == (1.0, 0.0, 0.0, 0.0, 0.0, 0.0) |
|
|
|
|
| def test_embed_record_image_requires_path() -> None: |
| """Embedding a record without an image path fails loudly.""" |
|
|
| with pytest.raises(MissingImageError): |
| embed_record_image(None, ColorStatsEmbedder()) |
|
|
|
|
| def test_ijepa_embedder_reports_missing_optional_dependencies() -> None: |
| """I-JEPA is optional and explains how to install its dependencies.""" |
|
|
| try: |
| IJepaImageEmbedder(model_id="unused") |
| except JepaDependencyError as error: |
| assert "sync-ijepa" in str(error) |
| except OSError: |
| pytest.skip("I-JEPA dependencies are installed but model files are unavailable.") |
| else: |
| pytest.skip("I-JEPA dependencies and model files are available in this environment.") |
|
|
|
|
| def test_mean_pool_features_preserves_last_dimension() -> None: |
| """Pooling should average token axes but keep hidden features.""" |
|
|
| fake_features = _FakeFeatures(ndim=4) |
|
|
| pooled = _mean_pool_features(fake_features) |
|
|
| assert pooled is fake_features |
| assert fake_features.mean_dim == (1, 2) |
|
|
|
|
| def test_patch_interest_heatmap_matches_image_size() -> None: |
| """Patch interest scores render into a transparent overlay at image resolution.""" |
|
|
| interest_map = PatchInterestMap( |
| scores=((0.0, 1.0), (0.5, 0.75)), |
| image_size=(12, 8), |
| ) |
|
|
| heatmap = render_patch_interest_heatmap(interest_map) |
|
|
| assert heatmap.mode == "RGBA" |
| assert heatmap.size == (12, 8) |
| pixel = cast(tuple[int, int, int, int], heatmap.getpixel((1, 1))) |
| assert pixel[3] > 0 |
|
|
|
|
| def test_patch_interest_overlay_preserves_source_size(tmp_path: Path) -> None: |
| """The rendered overlay can be shown directly by Streamlit.""" |
|
|
| image_path = tmp_path / "source.png" |
| Image.new("RGB", (10, 10), color=(20, 30, 40)).save(image_path) |
| interest_map = PatchInterestMap(scores=((1.0,),), image_size=(10, 10)) |
|
|
| overlay = render_patch_interest_overlay(image_path, interest_map) |
|
|
| assert overlay.mode == "RGBA" |
| assert overlay.size == (10, 10) |
| assert overlay.getpixel((5, 5)) != (20, 30, 40, 255) |
|
|
|
|
| class _FakeFeatures: |
| """Tiny tensor-like object for testing pooling dimensions without torch.""" |
|
|
| def __init__(self, *, ndim: int) -> None: |
| self.ndim = ndim |
| self.mean_dim: tuple[int, ...] | None = None |
|
|
| def mean(self, *, dim: tuple[int, ...]) -> _FakeFeatures: |
| self.mean_dim = dim |
| return self |
|
|
