mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-03-12 12:25:32 +08:00
chore(tests): remove obsolete test files for compressor and extractor
This commit is contained in:
@@ -1,99 +0,0 @@
|
|||||||
"""Tests for PoolNetCompressor module."""
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
import torch
|
|
||||||
from feature_compressor.core.compressor import PoolNetCompressor
|
|
||||||
|
|
||||||
|
|
||||||
class TestPoolNetCompressor:
|
|
||||||
"""Test suite for PoolNetCompressor class."""
|
|
||||||
|
|
||||||
def test_compressor_init(self):
|
|
||||||
"""Test PoolNetCompressor initializes with correct parameters."""
|
|
||||||
# This test will fail until we implement the module
|
|
||||||
|
|
||||||
compressor = PoolNetCompressor(
|
|
||||||
input_dim=1024,
|
|
||||||
compression_dim=256,
|
|
||||||
top_k_ratio=0.5,
|
|
||||||
hidden_ratio=2.0,
|
|
||||||
dropout_rate=0.1,
|
|
||||||
use_residual=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert compressor.input_dim == 1024
|
|
||||||
assert compressor.compression_dim == 256
|
|
||||||
assert compressor.top_k_ratio == 0.5
|
|
||||||
|
|
||||||
def test_compressor_forward_shape(self):
|
|
||||||
"""Test output shape is [batch, compression_dim]."""
|
|
||||||
|
|
||||||
compressor = PoolNetCompressor(
|
|
||||||
input_dim=1024,
|
|
||||||
compression_dim=256,
|
|
||||||
top_k_ratio=0.5,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Simulate DINOv2 output: batch=2, seq_len=257 (CLS+256 patches), dim=1024
|
|
||||||
x = torch.randn(2, 257, 1024)
|
|
||||||
out = compressor(x)
|
|
||||||
|
|
||||||
assert out.shape == (2, 256), f"Expected (2, 256), got {out.shape}"
|
|
||||||
|
|
||||||
def test_attention_scores_shape(self):
|
|
||||||
"""Test attention scores have shape [batch, seq_len]."""
|
|
||||||
|
|
||||||
compressor = PoolNetCompressor(input_dim=1024, compression_dim=256)
|
|
||||||
|
|
||||||
x = torch.randn(2, 257, 1024)
|
|
||||||
scores = compressor._compute_attention_scores(x)
|
|
||||||
|
|
||||||
assert scores.shape == (2, 257), f"Expected (2, 257), got {scores.shape}"
|
|
||||||
|
|
||||||
def test_top_k_selection(self):
|
|
||||||
"""Test that only top_k_ratio tokens are selected."""
|
|
||||||
|
|
||||||
compressor = PoolNetCompressor(
|
|
||||||
input_dim=1024, compression_dim=256, top_k_ratio=0.5
|
|
||||||
)
|
|
||||||
|
|
||||||
x = torch.randn(2, 257, 1024)
|
|
||||||
pooled = compressor._apply_pooling(x, compressor._compute_attention_scores(x))
|
|
||||||
|
|
||||||
# With top_k_ratio=0.5, should select 50% of tokens (int rounds down)
|
|
||||||
expected_k = 128 # int(257 * 0.5) = 128
|
|
||||||
assert pooled.shape[1] == expected_k, (
|
|
||||||
f"Expected seq_len={expected_k}, got {pooled.shape[1]}"
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_residual_connection(self):
|
|
||||||
"""Test residual adds input contribution to output."""
|
|
||||||
|
|
||||||
compressor = PoolNetCompressor(
|
|
||||||
input_dim=1024,
|
|
||||||
compression_dim=256,
|
|
||||||
use_residual=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
x = torch.randn(2, 257, 1024)
|
|
||||||
out1 = compressor(x)
|
|
||||||
|
|
||||||
# Residual should affect output
|
|
||||||
assert out1 is not None
|
|
||||||
assert out1.shape == (2, 256)
|
|
||||||
|
|
||||||
def test_gpu_device(self):
|
|
||||||
"""Test model moves to GPU correctly if available."""
|
|
||||||
|
|
||||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
||||||
|
|
||||||
compressor = PoolNetCompressor(
|
|
||||||
input_dim=1024,
|
|
||||||
compression_dim=256,
|
|
||||||
device=device,
|
|
||||||
)
|
|
||||||
|
|
||||||
x = torch.randn(2, 257, 1024).to(device)
|
|
||||||
out = compressor(x)
|
|
||||||
|
|
||||||
assert out.device.type == device
|
|
||||||
@@ -1,152 +0,0 @@
|
|||||||
"""Tests for DINOv2FeatureExtractor module."""
|
|
||||||
|
|
||||||
import json
|
|
||||||
import tempfile
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import pytest
|
|
||||||
import torch
|
|
||||||
from feature_compressor.core.extractor import DINOv2FeatureExtractor
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
|
|
||||||
class TestDINOv2FeatureExtractor:
|
|
||||||
"""Test suite for DINOv2FeatureExtractor class."""
|
|
||||||
|
|
||||||
def test_extractor_init(self):
|
|
||||||
"""Test DINOv2FeatureExtractor initializes correctly."""
|
|
||||||
|
|
||||||
extractor = DINOv2FeatureExtractor()
|
|
||||||
|
|
||||||
assert extractor.model is not None
|
|
||||||
assert extractor.processor is not None
|
|
||||||
assert extractor.compressor is not None
|
|
||||||
|
|
||||||
def test_single_image_processing(self):
|
|
||||||
"""Test processing a single image."""
|
|
||||||
|
|
||||||
extractor = DINOv2FeatureExtractor()
|
|
||||||
|
|
||||||
# Create a simple test image
|
|
||||||
img_array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
|
||||||
img = Image.fromarray(img_array)
|
|
||||||
|
|
||||||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
|
|
||||||
img.save(f.name)
|
|
||||||
result = extractor.process_image(f.name)
|
|
||||||
|
|
||||||
assert "original_features" in result
|
|
||||||
assert "compressed_features" in result
|
|
||||||
assert "metadata" in result
|
|
||||||
|
|
||||||
# Check shapes
|
|
||||||
assert result["original_features"].shape[0] == 1 # batch=1
|
|
||||||
assert result["compressed_features"].shape == (1, 256)
|
|
||||||
assert "compression_ratio" in result["metadata"]
|
|
||||||
|
|
||||||
def test_output_structure(self):
|
|
||||||
"""Test output structure contains expected keys."""
|
|
||||||
|
|
||||||
extractor = DINOv2FeatureExtractor()
|
|
||||||
|
|
||||||
# Create test image
|
|
||||||
img_array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
|
||||||
img = Image.fromarray(img_array)
|
|
||||||
|
|
||||||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
|
|
||||||
img.save(f.name)
|
|
||||||
result = extractor.process_image(f.name)
|
|
||||||
|
|
||||||
required_keys = [
|
|
||||||
"original_features",
|
|
||||||
"compressed_features",
|
|
||||||
"pooled_features",
|
|
||||||
"metadata",
|
|
||||||
]
|
|
||||||
for key in required_keys:
|
|
||||||
assert key in result, f"Missing key: {key}"
|
|
||||||
|
|
||||||
metadata_keys = [
|
|
||||||
"compression_ratio",
|
|
||||||
"processing_time",
|
|
||||||
"feature_norm",
|
|
||||||
"device",
|
|
||||||
]
|
|
||||||
for key in metadata_keys:
|
|
||||||
assert key in result["metadata"], f"Missing metadata key: {key}"
|
|
||||||
|
|
||||||
def test_feature_saving(self):
|
|
||||||
"""Test saving features to disk."""
|
|
||||||
|
|
||||||
extractor = DINOv2FeatureExtractor()
|
|
||||||
|
|
||||||
# Create test image
|
|
||||||
img_array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
|
||||||
img = Image.fromarray(img_array)
|
|
||||||
|
|
||||||
with tempfile.TemporaryDirectory() as tmpdir:
|
|
||||||
tmpdir = Path(tmpdir)
|
|
||||||
|
|
||||||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
|
|
||||||
img.save(f.name)
|
|
||||||
result = extractor.process_image(f.name)
|
|
||||||
|
|
||||||
# Save features
|
|
||||||
json_path = tmpdir / "features.json"
|
|
||||||
from feature_compressor.utils.feature_utils import (
|
|
||||||
save_features_to_json,
|
|
||||||
)
|
|
||||||
|
|
||||||
save_features_to_json(
|
|
||||||
result["compressed_features"], json_path, result["metadata"]
|
|
||||||
)
|
|
||||||
|
|
||||||
assert json_path.exists()
|
|
||||||
|
|
||||||
# Verify file can be loaded
|
|
||||||
with open(json_path) as f:
|
|
||||||
data = json.load(f)
|
|
||||||
assert "features" in data
|
|
||||||
assert "metadata" in data
|
|
||||||
|
|
||||||
def test_batch_processing(self):
|
|
||||||
"""Test batch processing of multiple images."""
|
|
||||||
|
|
||||||
extractor = DINOv2FeatureExtractor()
|
|
||||||
|
|
||||||
# Create multiple test images
|
|
||||||
images = []
|
|
||||||
with tempfile.TemporaryDirectory() as tmpdir:
|
|
||||||
tmpdir = Path(tmpdir)
|
|
||||||
|
|
||||||
for i in range(3):
|
|
||||||
img_array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
|
||||||
img = Image.fromarray(img_array)
|
|
||||||
img_path = tmpdir / f"test_{i}.jpg"
|
|
||||||
img.save(img_path)
|
|
||||||
images.append(str(img_path))
|
|
||||||
|
|
||||||
results = extractor.process_batch(str(tmpdir), batch_size=2)
|
|
||||||
|
|
||||||
assert len(results) == 3
|
|
||||||
for result in results:
|
|
||||||
assert result["compressed_features"].shape == (1, 256)
|
|
||||||
|
|
||||||
def test_gpu_handling(self):
|
|
||||||
"""Test GPU device handling."""
|
|
||||||
|
|
||||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
||||||
extractor = DINOv2FeatureExtractor(device=device)
|
|
||||||
|
|
||||||
assert extractor.device.type == device
|
|
||||||
|
|
||||||
# Create test image
|
|
||||||
img_array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
|
||||||
img = Image.fromarray(img_array)
|
|
||||||
|
|
||||||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
|
|
||||||
img.save(f.name)
|
|
||||||
result = extractor.process_image(f.name)
|
|
||||||
|
|
||||||
assert result["metadata"]["device"] == device
|
|
||||||
Reference in New Issue
Block a user