feat(feature-compressor): add DINOv2 feature extraction and compression pipeline

mini-nav/tests/test_extractor.py

@@ -0,0 +1,152 @@
+"""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