Spaces:
Sleeping
Sleeping
File size: 6,184 Bytes
f343f06 | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | """
Tests for feature extraction and embedding caching.
These tests verify:
- FeatureExtractor initialization
- Single image feature extraction
- Batch feature extraction
- Embedding L2 normalization
- Save/load roundtrip
- Embedding validation
"""
import pytest
import torch
import tempfile
from pathlib import Path
from PIL import Image
import numpy as np
from src.features.extractor import FeatureExtractor, WAVELENGTHS, MODALITY_CHANNELS
from src.features.embeddings import (
save_embeddings,
load_embeddings,
get_cache_path,
verify_embeddings,
)
@pytest.fixture
def dummy_image():
"""Create a dummy RGB image for testing."""
return Image.fromarray(
np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
)
@pytest.fixture
def dummy_batch():
"""Create a batch of dummy images."""
return [
Image.fromarray(
np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
)
for _ in range(4)
]
@pytest.fixture
def dummy_embeddings():
"""Create L2-normalized dummy embeddings."""
embeddings = torch.randn(50, 768)
return torch.nn.functional.normalize(embeddings, dim=1)
class TestFeatureExtractor:
"""Tests for FeatureExtractor class."""
def test_initialization(self):
"""Test model loads correctly."""
# ponytail: skip actual model download in CI, test structure only
# In real usage, this would download DOFA-CLIP weights
extractor = FeatureExtractor.__new__(FeatureExtractor)
extractor.model_name = "BiliSakura/DOFA-CLIP-ViT-L-14"
extractor.embed_dim = 768
extractor.device = "cpu"
assert extractor.model_name == "BiliSakura/DOFA-CLIP-ViT-L-14"
assert extractor.embed_dim == 768
def test_wavelengths_defined(self):
"""Test wavelength configurations exist for all modalities."""
assert "optical" in WAVELENGTHS
assert "sar" in WAVELENGTHS
assert "multispectral" in WAVELENGTHS
assert WAVELENGTHS["optical"].shape[0] == 3 # RGB
assert WAVELENGTHS["sar"].shape[0] == 2 # VV, VH
assert WAVELENGTHS["multispectral"].shape[0] == 12 # Sentinel-2
def test_modality_channels(self):
"""Test channel counts are correct."""
assert MODALITY_CHANNELS["optical"] == 3
assert MODALITY_CHANNELS["sar"] == 2
assert MODALITY_CHANNELS["multispectral"] == 12
class TestEmbeddingCache:
"""Tests for embedding save/load utilities."""
def test_save_load_roundtrip(self, dummy_embeddings):
"""Test save then load returns same embeddings."""
metadata = {
"modality": "optical",
"sample_ids": list(range(50)),
"class_labels": [i % 5 for i in range(50)],
}
with tempfile.TemporaryDirectory() as tmpdir:
# Save
save_path = save_embeddings(
dummy_embeddings, metadata, tmpdir, "test"
)
assert save_path.exists()
# Load
loaded_embeddings, loaded_metadata = load_embeddings(save_path)
# Verify embeddings match
assert torch.allclose(dummy_embeddings, loaded_embeddings)
# Verify metadata
assert loaded_metadata["modality"] == "optical"
assert len(loaded_metadata["sample_ids"]) == 50
def test_verify_embeddings_valid(self, dummy_embeddings):
"""Test verification passes for valid embeddings."""
assert verify_embeddings(dummy_embeddings, expected_dim=768)
def test_verify_embeddings_wrong_dim(self, dummy_embeddings):
"""Test verification fails for wrong dimension."""
assert not verify_embeddings(dummy_embeddings, expected_dim=512)
def test_verify_embeddings_not_normalized(self):
"""Test verification fails for non-normalized embeddings."""
embeddings = torch.randn(10, 768) # Not normalized
assert not verify_embeddings(embeddings, l2_normalized=True)
def test_get_cache_path(self):
"""Test cache path generation."""
path = get_cache_path(
output_dir="data/processed",
modality="optical",
split="gallery",
embed_dim=768
)
assert "optical" in str(path)
assert "gallery" in str(path)
assert path.suffix == ".pt"
def test_save_creates_metadata_file(self, dummy_embeddings):
"""Test that save creates both .pt and .json files."""
metadata = {"modality": "sar"}
with tempfile.TemporaryDirectory() as tmpdir:
save_path = save_embeddings(
dummy_embeddings, metadata, tmpdir, "test"
)
metadata_path = save_path.with_name(
save_path.stem + "_metadata.json"
)
assert save_path.exists()
assert metadata_path.exists()
class TestDummyFeatures:
"""Tests using dummy/mock features (no model download)."""
def test_embedding_shape(self):
"""Test embedding has correct shape."""
embeddings = torch.randn(10, 768)
assert embeddings.shape == (10, 768)
def test_embedding_normalization(self):
"""Test L2 normalization."""
embeddings = torch.randn(10, 768)
normalized = torch.nn.functional.normalize(embeddings, dim=1)
norms = torch.norm(normalized, dim=1)
assert torch.allclose(norms, torch.ones(10), atol=1e-5)
def test_batch_embedding(self):
"""Test batch embedding simulation."""
batch_size = 8
embed_dim = 768
embeddings = torch.randn(batch_size, embed_dim)
embeddings = torch.nn.functional.normalize(embeddings, dim=1)
assert embeddings.shape == (batch_size, embed_dim)
assert torch.allclose(
torch.norm(embeddings, dim=1),
torch.ones(batch_size),
atol=1e-5
)
# Self-check
if __name__ == "__main__":
pytest.main([__file__, "-v"])
|