mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
chore(tests): remove all test files from mini-nav
This commit is contained in:
@@ -1 +0,0 @@
|
||||
"""Test suite for DINOv2 Feature Compressor."""
|
||||
@@ -1,341 +0,0 @@
|
||||
"""Tests for compressor modules (HashCompressor, Pipeline)."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from unittest.mock import Mock, patch
|
||||
from compressors import (
|
||||
BinarySign,
|
||||
HashCompressor,
|
||||
HashPipeline,
|
||||
VideoPositiveMask,
|
||||
bits_to_hash,
|
||||
create_pipeline_from_config,
|
||||
hamming_distance,
|
||||
hamming_similarity,
|
||||
hash_to_bits,
|
||||
)
|
||||
from configs import cfg_manager
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class TestHashCompressor:
|
||||
"""Test suite for HashCompressor."""
|
||||
|
||||
def test_hash_compressor_init(self):
|
||||
"""Verify HashCompressor initializes with correct dimensions."""
|
||||
compressor = HashCompressor(input_dim=1024, hash_bits=512)
|
||||
assert compressor.input_dim == 1024
|
||||
assert compressor.hash_bits == 512
|
||||
|
||||
def test_hash_compressor_forward(self):
|
||||
"""Verify forward pass produces correct output shapes."""
|
||||
compressor = HashCompressor(input_dim=1024, hash_bits=512)
|
||||
tokens = torch.randn(4, 197, 1024) # [B, N, input_dim]
|
||||
|
||||
logits, hash_codes, bits = compressor(tokens)
|
||||
|
||||
assert logits.shape == (4, 512)
|
||||
assert hash_codes.shape == (4, 512)
|
||||
assert bits.shape == (4, 512)
|
||||
# Verify bits are binary (0 or 1)
|
||||
assert torch.all((bits == 0) | (bits == 1))
|
||||
|
||||
def test_hash_compressor_encode(self):
|
||||
"""Verify encode method returns binary bits."""
|
||||
compressor = HashCompressor(input_dim=1024, hash_bits=512)
|
||||
tokens = torch.randn(2, 197, 1024)
|
||||
|
||||
bits = compressor.encode(tokens)
|
||||
|
||||
assert bits.shape == (2, 512)
|
||||
assert bits.dtype == torch.int32
|
||||
assert torch.all((bits == 0) | (bits == 1))
|
||||
|
||||
def test_hash_compressor_similarity(self):
|
||||
"""Verify compute_similarity returns correct shape."""
|
||||
compressor = HashCompressor(input_dim=1024, hash_bits=512)
|
||||
|
||||
# Create random bits
|
||||
bits1 = torch.randint(0, 2, (3, 512))
|
||||
bits2 = torch.randint(0, 2, (5, 512))
|
||||
|
||||
sim = compressor.compute_similarity(bits1, bits2)
|
||||
|
||||
assert sim.shape == (3, 5)
|
||||
|
||||
|
||||
class TestBinarySign:
|
||||
"""Test suite for BinarySign function."""
|
||||
|
||||
def test_binary_sign_forward(self):
|
||||
"""Verify BinarySign produces {-1, +1} outputs."""
|
||||
x = torch.randn(4, 512)
|
||||
result = BinarySign.apply(x)
|
||||
|
||||
assert torch.all((result == 1) | (result == -1))
|
||||
|
||||
def test_binary_sign_round_trip(self):
|
||||
"""Verify bits -> hash -> bits preserves values."""
|
||||
bits = torch.randint(0, 2, (4, 512))
|
||||
hash_codes = bits_to_hash(bits)
|
||||
bits_recovered = hash_to_bits(hash_codes)
|
||||
|
||||
assert torch.equal(bits, bits_recovered)
|
||||
|
||||
|
||||
class TestHammingMetrics:
|
||||
"""Test suite for Hamming distance and similarity."""
|
||||
|
||||
def test_hamming_distance_same_codes(self):
|
||||
"""Verify hamming distance is 0 for identical single codes."""
|
||||
bits1 = torch.randint(0, 2, (512,))
|
||||
bits2 = bits1.clone()
|
||||
|
||||
dist = hamming_distance(bits1, bits2)
|
||||
|
||||
assert dist.item() == 0
|
||||
|
||||
def test_hamming_distance_self_comparison(self):
|
||||
"""Verify hamming distance diagonal is 0 (each code compared to itself)."""
|
||||
bits = torch.randint(0, 2, (10, 512))
|
||||
|
||||
dist = hamming_distance(bits, bits)
|
||||
|
||||
# Diagonal should be 0 (distance to self)
|
||||
diagonal = torch.diag(dist)
|
||||
assert torch.all(diagonal == 0)
|
||||
|
||||
def test_hamming_distance_different(self):
|
||||
"""Verify hamming distance is correct for different codes."""
|
||||
bits1 = torch.zeros(1, 512, dtype=torch.int32)
|
||||
bits2 = torch.ones(1, 512, dtype=torch.int32)
|
||||
|
||||
dist = hamming_distance(bits1, bits2)
|
||||
|
||||
assert dist.item() == 512
|
||||
|
||||
def test_hamming_similarity(self):
|
||||
"""Verify hamming similarity is positive for similar codes."""
|
||||
hash1 = torch.ones(1, 512)
|
||||
hash2 = torch.ones(1, 512)
|
||||
|
||||
sim = hamming_similarity(hash1, hash2)
|
||||
|
||||
assert sim.item() == 512 # Max similarity
|
||||
|
||||
|
||||
class TestHashLoss:
|
||||
"""Test suite for HashLoss."""
|
||||
|
||||
def test_hash_loss_init(self):
|
||||
"""Verify HashLoss initializes with correct parameters."""
|
||||
from compressors import HashLoss
|
||||
|
||||
loss_fn = HashLoss(
|
||||
contrastive_weight=1.0,
|
||||
distill_weight=0.5,
|
||||
quant_weight=0.01,
|
||||
temperature=0.2,
|
||||
)
|
||||
|
||||
assert loss_fn.contrastive_weight == 1.0
|
||||
assert loss_fn.distill_weight == 0.5
|
||||
assert loss_fn.quant_weight == 0.01
|
||||
assert loss_fn.temperature == 0.2
|
||||
|
||||
def test_hash_loss_forward(self):
|
||||
"""Verify HashLoss computes loss correctly."""
|
||||
from compressors import HashLoss
|
||||
|
||||
loss_fn = HashLoss()
|
||||
|
||||
batch_size = 4
|
||||
hash_bits = 512
|
||||
logits = torch.randn(batch_size, hash_bits)
|
||||
hash_codes = torch.sign(logits)
|
||||
teacher_embed = torch.randn(batch_size, 1024)
|
||||
positive_mask = torch.eye(batch_size, dtype=torch.bool)
|
||||
|
||||
total_loss, components = loss_fn(
|
||||
logits=logits,
|
||||
hash_codes=hash_codes,
|
||||
teacher_embed=teacher_embed,
|
||||
positive_mask=positive_mask,
|
||||
)
|
||||
|
||||
assert "contrastive" in components
|
||||
assert "distill" in components
|
||||
assert "quantization" in components
|
||||
assert "total" in components
|
||||
|
||||
|
||||
class TestVideoPositiveMask:
|
||||
"""Test suite for VideoPositiveMask."""
|
||||
|
||||
def test_from_frame_indices(self):
|
||||
"""Verify positive mask generation from frame indices."""
|
||||
mask_gen = VideoPositiveMask(temporal_window=2)
|
||||
|
||||
frame_indices = torch.tensor([0, 1, 3, 5])
|
||||
|
||||
mask = mask_gen.from_frame_indices(frame_indices)
|
||||
|
||||
assert mask.shape == (4, 4)
|
||||
# Frame 0 and 1 should be positive (distance 1 <= 2)
|
||||
assert mask[0, 1] == True
|
||||
# Frame 0 and 3 should be negative (distance 3 > 2)
|
||||
assert mask[0, 3] == False
|
||||
|
||||
def test_from_video_ids(self):
|
||||
"""Verify positive mask generation from video IDs and frame indices."""
|
||||
mask_gen = VideoPositiveMask(temporal_window=2)
|
||||
|
||||
video_ids = torch.tensor([0, 0, 1, 1])
|
||||
frame_indices = torch.tensor([0, 1, 0, 1])
|
||||
|
||||
mask = mask_gen.from_video_ids(video_ids, frame_indices)
|
||||
|
||||
assert mask.shape == (4, 4)
|
||||
# Same video and temporally close
|
||||
assert mask[0, 1] == True # video 0, frames 0,1
|
||||
# Different video
|
||||
assert mask[0, 2] == False # video 0 vs 1
|
||||
|
||||
|
||||
class TestHashPipeline:
|
||||
"""Test suite for HashPipeline."""
|
||||
|
||||
@patch("compressors.pipeline.load_sam_model")
|
||||
@patch("compressors.pipeline.AutoModel.from_pretrained")
|
||||
@patch("compressors.pipeline.AutoImageProcessor.from_pretrained")
|
||||
def test_pipeline_init(
|
||||
self,
|
||||
mock_processor_from_pretrained,
|
||||
mock_model_from_pretrained,
|
||||
mock_load_sam_model,
|
||||
):
|
||||
"""Verify pipeline initializes all components."""
|
||||
mock_processor_from_pretrained.return_value = Mock()
|
||||
|
||||
mock_model = Mock()
|
||||
mock_model.to.return_value = mock_model
|
||||
mock_model.eval.return_value = None
|
||||
mock_model_from_pretrained.return_value = mock_model
|
||||
|
||||
mock_load_sam_model.return_value = (Mock(), Mock())
|
||||
|
||||
pipeline = HashPipeline(
|
||||
dino_model="facebook/dinov2-large",
|
||||
hash_bits=512,
|
||||
)
|
||||
|
||||
assert pipeline.dino_model == "facebook/dinov2-large"
|
||||
assert pipeline.sam_model_name == "facebook/sam2.1-hiera-large"
|
||||
assert pipeline.dino_dim == 1024
|
||||
mock_load_sam_model.assert_called_once()
|
||||
|
||||
@patch("compressors.pipeline.load_sam_model")
|
||||
@patch("compressors.pipeline.AutoModel.from_pretrained")
|
||||
@patch("compressors.pipeline.AutoImageProcessor.from_pretrained")
|
||||
def test_pipeline_hash_bits(
|
||||
self,
|
||||
mock_processor_from_pretrained,
|
||||
mock_model_from_pretrained,
|
||||
mock_load_sam_model,
|
||||
):
|
||||
"""Verify pipeline uses correct hash bits."""
|
||||
mock_processor_from_pretrained.return_value = Mock()
|
||||
|
||||
mock_model = Mock()
|
||||
mock_model.to.return_value = mock_model
|
||||
mock_model.eval.return_value = None
|
||||
mock_model_from_pretrained.return_value = mock_model
|
||||
|
||||
mock_load_sam_model.return_value = (Mock(), Mock())
|
||||
|
||||
pipeline = HashPipeline(hash_bits=256)
|
||||
assert pipeline.hash_bits == 256
|
||||
|
||||
|
||||
class TestConfigIntegration:
|
||||
"""Test suite for config integration with pipeline."""
|
||||
|
||||
@patch("compressors.pipeline.load_sam_model")
|
||||
@patch("compressors.pipeline.AutoModel.from_pretrained")
|
||||
@patch("compressors.pipeline.AutoImageProcessor.from_pretrained")
|
||||
def test_create_pipeline_from_config(
|
||||
self,
|
||||
mock_processor_from_pretrained,
|
||||
mock_model_from_pretrained,
|
||||
mock_load_sam_model,
|
||||
):
|
||||
"""Verify pipeline can be created from config."""
|
||||
mock_processor_from_pretrained.return_value = Mock()
|
||||
|
||||
mock_model = Mock()
|
||||
mock_model.to.return_value = mock_model
|
||||
mock_model.eval.return_value = None
|
||||
mock_model_from_pretrained.return_value = mock_model
|
||||
|
||||
mock_load_sam_model.return_value = (Mock(), Mock())
|
||||
|
||||
config = cfg_manager.load()
|
||||
|
||||
pipeline = create_pipeline_from_config(config)
|
||||
|
||||
assert isinstance(pipeline, HashPipeline)
|
||||
assert pipeline.hash_bits == config.model.compression_dim
|
||||
assert pipeline.sam_max_masks == config.model.sam_max_masks
|
||||
assert pipeline.sam_min_mask_area == config.model.sam_min_mask_area
|
||||
|
||||
def test_config_settings(self):
|
||||
"""Verify config contains required settings."""
|
||||
config = cfg_manager.load()
|
||||
|
||||
assert hasattr(config.model, "dino_model")
|
||||
assert hasattr(config.model, "compression_dim")
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
class TestPipelineIntegration:
|
||||
"""Integration tests for full pipeline (slow, requires model downloads)."""
|
||||
|
||||
def test_pipeline_end_to_end(self):
|
||||
"""Test full pipeline with actual models (slow test)."""
|
||||
# Skip if no GPU
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("Requires CUDA")
|
||||
|
||||
# Create a simple test image
|
||||
image = Image.new("RGB", (640, 480), color=(128, 128, 128))
|
||||
|
||||
# Initialize pipeline (will download models on first run)
|
||||
pipeline = HashPipeline(
|
||||
dino_model="facebook/dinov2-large",
|
||||
hash_bits=512,
|
||||
)
|
||||
|
||||
# Run pipeline
|
||||
hash_bits = pipeline(image)
|
||||
|
||||
# Verify output shape
|
||||
assert hash_bits.dim() == 2
|
||||
assert hash_bits.shape[1] == 512
|
||||
assert torch.all((hash_bits == 0) | (hash_bits == 1))
|
||||
|
||||
def test_extract_features(self):
|
||||
"""Test feature extraction."""
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("Requires CUDA")
|
||||
|
||||
image = Image.new("RGB", (640, 480), color=(128, 128, 128))
|
||||
|
||||
pipeline = HashPipeline(
|
||||
dino_model="facebook/dinov2-large",
|
||||
)
|
||||
|
||||
features = pipeline.extract_features(image)
|
||||
|
||||
# Should return DINO features (1024 for large)
|
||||
assert features.dim() == 2
|
||||
assert features.shape[1] == 1024
|
||||
@@ -1,167 +0,0 @@
|
||||
"""Tests for configuration system using Pydantic models."""
|
||||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from configs import (
|
||||
Config,
|
||||
ConfigError,
|
||||
ConfigManager,
|
||||
ModelConfig,
|
||||
OutputConfig,
|
||||
PoolingType,
|
||||
cfg_manager,
|
||||
load_yaml,
|
||||
save_yaml,
|
||||
)
|
||||
from pydantic import ValidationError
|
||||
|
||||
|
||||
class TestConfigModels:
|
||||
"""Test suite for Pydantic configuration models."""
|
||||
|
||||
def test_model_config_defaults(self):
|
||||
"""Verify ModelConfig creates with correct defaults."""
|
||||
config = ModelConfig()
|
||||
assert config.dino_model == "facebook/dinov2-large"
|
||||
assert config.compression_dim == 512
|
||||
assert config.device == "auto"
|
||||
|
||||
def test_model_config_validation(self):
|
||||
"""Test validation constraints for ModelConfig."""
|
||||
# Test compression_dim > 0
|
||||
with pytest.raises(ValidationError, match="greater than 0"):
|
||||
ModelConfig(compression_dim=0)
|
||||
|
||||
with pytest.raises(ValidationError, match="greater than 0"):
|
||||
ModelConfig(compression_dim=-1)
|
||||
|
||||
def test_output_config_defaults(self):
|
||||
"""Verify OutputConfig creates with correct defaults."""
|
||||
config = OutputConfig()
|
||||
output_dir = Path(__file__).parent.parent.parent / "outputs"
|
||||
|
||||
assert config.directory == output_dir
|
||||
|
||||
def test_pooling_type_enum(self):
|
||||
"""Verify PoolingType enum values."""
|
||||
assert PoolingType.ATTENTION.value == "attention"
|
||||
assert PoolingType.ATTENTION == PoolingType("attention")
|
||||
|
||||
def test_feature_compressor_config(self):
|
||||
"""Verify FeatureCompressorConfig nests all models correctly."""
|
||||
model_cfg = ModelConfig(compression_dim=512)
|
||||
out_cfg = OutputConfig(directory="/tmp/outputs")
|
||||
|
||||
config = Config(
|
||||
model=model_cfg,
|
||||
output=out_cfg,
|
||||
)
|
||||
|
||||
assert config.model.compression_dim == 512
|
||||
assert config.output.directory == Path("/tmp/outputs")
|
||||
|
||||
|
||||
class TestYamlLoader:
|
||||
"""Test suite for YAML loading and saving."""
|
||||
|
||||
def test_load_existing_yaml(self):
|
||||
"""Load config.yaml and verify values."""
|
||||
config_path = Path(__file__).parent.parent / "configs" / "config.yaml"
|
||||
config = load_yaml(config_path, Config)
|
||||
|
||||
# Verify model config
|
||||
assert config.model.dino_model == "facebook/dinov2-large"
|
||||
assert config.model.compression_dim == 256
|
||||
|
||||
# Verify output config
|
||||
output_dir = Path(__file__).parent.parent.parent / "outputs"
|
||||
|
||||
assert config.output.directory == output_dir
|
||||
|
||||
def test_load_yaml_validation(self):
|
||||
"""Test that invalid data raises ConfigError."""
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
|
||||
# Write invalid config (missing required fields)
|
||||
yaml.dump({"invalid": "data"}, f)
|
||||
temp_path = f.name
|
||||
|
||||
try:
|
||||
with pytest.raises(ConfigError, match="validation failed"):
|
||||
load_yaml(Path(temp_path), Config)
|
||||
finally:
|
||||
Path(temp_path).unlink()
|
||||
|
||||
def test_save_yaml_roundtrip(self):
|
||||
"""Create config, save to temp, verify file exists with content."""
|
||||
original = cfg_manager.load()
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
|
||||
temp_path = Path(f.name)
|
||||
|
||||
try:
|
||||
save_yaml(temp_path, original)
|
||||
|
||||
# Verify file exists and has content
|
||||
assert Path(temp_path).exists()
|
||||
with open(temp_path, "r") as f:
|
||||
content = f.read()
|
||||
assert len(content) > 0
|
||||
assert "model" in content
|
||||
assert "visualization" in content
|
||||
assert "output" in content
|
||||
finally:
|
||||
Path(temp_path).unlink()
|
||||
|
||||
def test_load_yaml_file_not_found(self):
|
||||
"""Verify FileNotFoundError raises ConfigError."""
|
||||
with pytest.raises(ConfigError, match="not found"):
|
||||
load_yaml(Path("/nonexistent/path/config.yaml"), Config)
|
||||
|
||||
|
||||
class TestConfigManager:
|
||||
"""Test suite for ConfigManager singleton with multi-config support."""
|
||||
|
||||
def test_singleton_pattern(self):
|
||||
"""Verify ConfigManager() returns same instance."""
|
||||
manager1 = ConfigManager()
|
||||
manager2 = ConfigManager()
|
||||
assert manager1 is manager2
|
||||
|
||||
def test_load_config(self):
|
||||
"""Test loading default config."""
|
||||
config = cfg_manager.load()
|
||||
|
||||
assert config is not None
|
||||
assert config.model.compression_dim == 256
|
||||
|
||||
def test_get_without_load(self):
|
||||
"""Test that get() auto-loads config if not loaded."""
|
||||
# Reset the singleton's cached config
|
||||
cfg_manager._config = None
|
||||
|
||||
# get() should auto-load
|
||||
config = cfg_manager.get()
|
||||
assert config is not None
|
||||
assert config.model.compression_dim == 256
|
||||
|
||||
def test_save_config(self):
|
||||
"""Test saving configuration to file."""
|
||||
config = Config(
|
||||
model=ModelConfig(compression_dim=512),
|
||||
output=OutputConfig(),
|
||||
)
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
|
||||
temp_path = Path(f.name)
|
||||
|
||||
try:
|
||||
cfg_manager.save(config, path=temp_path)
|
||||
loaded_config = load_yaml(temp_path, Config)
|
||||
|
||||
assert loaded_config.model.compression_dim == 512
|
||||
finally:
|
||||
if temp_path.exists():
|
||||
temp_path.unlink()
|
||||
@@ -1,50 +0,0 @@
|
||||
"""Tests for feature extraction utilities."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import AutoImageProcessor, AutoModel
|
||||
|
||||
from utils.feature_extractor import (
|
||||
extract_batch_features,
|
||||
extract_single_image_feature,
|
||||
infer_vector_dim,
|
||||
)
|
||||
|
||||
TEST_MODEL_NAME = "facebook/dinov2-base"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_and_processor():
|
||||
processor = AutoImageProcessor.from_pretrained(TEST_MODEL_NAME)
|
||||
model = AutoModel.from_pretrained(TEST_MODEL_NAME)
|
||||
model.eval()
|
||||
yield processor, model
|
||||
del model
|
||||
del processor
|
||||
|
||||
|
||||
def test_infer_vector_dim(model_and_processor):
|
||||
"""Verify infer_vector_dim returns correct dimension."""
|
||||
processor, model = model_and_processor
|
||||
sample_image = Image.new("RGB", (224, 224), color="blue")
|
||||
dim = infer_vector_dim(processor, model, sample_image)
|
||||
assert dim == 768
|
||||
|
||||
|
||||
def test_extract_single_image_feature(model_and_processor):
|
||||
"""Verify single image feature extraction."""
|
||||
processor, model = model_and_processor
|
||||
sample_image = Image.new("RGB", (224, 224), color="red")
|
||||
features = extract_single_image_feature(processor, model, sample_image)
|
||||
assert isinstance(features, list)
|
||||
assert len(features) == 768
|
||||
|
||||
|
||||
def test_extract_batch_features(model_and_processor):
|
||||
"""Verify batch feature extraction."""
|
||||
processor, model = model_and_processor
|
||||
images = [Image.new("RGB", (224, 224), color="red") for _ in range(3)]
|
||||
features = extract_batch_features(processor, model, images)
|
||||
assert isinstance(features, torch.Tensor)
|
||||
assert features.shape == (3, 768)
|
||||
@@ -1,124 +0,0 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from simulator import HabitatSimulatorConfig, create_habitat_simulator
|
||||
|
||||
|
||||
class _FakeSimulatorConfiguration:
|
||||
def __init__(self):
|
||||
self.scene_id = ""
|
||||
self.enable_physics = True
|
||||
|
||||
|
||||
class _FakeAgentConfiguration:
|
||||
def __init__(self):
|
||||
self.sensor_specifications = []
|
||||
self.action_space = {}
|
||||
|
||||
|
||||
class _FakeCameraSensorSpec:
|
||||
def __init__(self):
|
||||
self.uuid = ""
|
||||
self.sensor_type = None
|
||||
self.resolution = []
|
||||
self.position = []
|
||||
|
||||
|
||||
class _FakeActuationSpec:
|
||||
def __init__(self, amount):
|
||||
self.amount = amount
|
||||
|
||||
|
||||
class _FakeActionSpec:
|
||||
def __init__(self, name, actuation):
|
||||
self.name = name
|
||||
self.actuation = actuation
|
||||
|
||||
|
||||
class _FakeConfiguration:
|
||||
def __init__(self, sim_cfg, agent_cfgs):
|
||||
self.sim_cfg = sim_cfg
|
||||
self.agent_cfgs = agent_cfgs
|
||||
|
||||
|
||||
class _FakeSimulator:
|
||||
def __init__(self, cfg):
|
||||
self.cfg = cfg
|
||||
self.initialized_agent_id = None
|
||||
|
||||
def initialize_agent(self, agent_id):
|
||||
self.initialized_agent_id = agent_id
|
||||
return {"agent_id": agent_id}
|
||||
|
||||
|
||||
def _create_fake_habitat_module():
|
||||
return SimpleNamespace(
|
||||
SimulatorConfiguration=_FakeSimulatorConfiguration,
|
||||
CameraSensorSpec=_FakeCameraSensorSpec,
|
||||
SensorType=SimpleNamespace(COLOR="color"),
|
||||
Configuration=_FakeConfiguration,
|
||||
Simulator=_FakeSimulator,
|
||||
agent=SimpleNamespace(
|
||||
AgentConfiguration=_FakeAgentConfiguration,
|
||||
ActionSpec=_FakeActionSpec,
|
||||
ActuationSpec=_FakeActuationSpec,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_create_habitat_simulator_builds_expected_configuration():
|
||||
fake_habitat = _create_fake_habitat_module()
|
||||
config = HabitatSimulatorConfig(
|
||||
scene_path="scene.glb",
|
||||
views_per_room=8,
|
||||
image_size=128,
|
||||
sensor_height=1.25,
|
||||
move_forward_step=0.5,
|
||||
enable_physics=False,
|
||||
sensor_uuid="rgb",
|
||||
agent_id=2,
|
||||
)
|
||||
|
||||
simulator, agent = create_habitat_simulator(config, habitat_sim_module=fake_habitat)
|
||||
|
||||
assert simulator.cfg.sim_cfg.scene_id == "scene.glb"
|
||||
assert simulator.cfg.sim_cfg.enable_physics is False
|
||||
|
||||
created_agent_cfg = simulator.cfg.agent_cfgs[0]
|
||||
sensor = created_agent_cfg.sensor_specifications[0]
|
||||
assert sensor.uuid == "rgb"
|
||||
assert sensor.sensor_type == "color"
|
||||
assert sensor.resolution == [128, 128]
|
||||
assert sensor.position == [0.0, 1.25, 0.0]
|
||||
|
||||
assert created_agent_cfg.action_space["move_forward"].actuation.amount == 0.5
|
||||
assert created_agent_cfg.action_space["turn_left"].actuation.amount == 45.0
|
||||
assert created_agent_cfg.action_space["turn_right"].actuation.amount == 45.0
|
||||
|
||||
assert simulator.initialized_agent_id == 2
|
||||
assert agent == {"agent_id": 2}
|
||||
|
||||
|
||||
def test_create_habitat_simulator_validates_views_per_room():
|
||||
fake_habitat = _create_fake_habitat_module()
|
||||
config = HabitatSimulatorConfig(scene_path="scene.glb", views_per_room=0)
|
||||
|
||||
with pytest.raises(ValueError, match="views_per_room"):
|
||||
create_habitat_simulator(config, habitat_sim_module=fake_habitat)
|
||||
|
||||
|
||||
def test_create_habitat_simulator_validates_image_size():
|
||||
fake_habitat = _create_fake_habitat_module()
|
||||
config = HabitatSimulatorConfig(scene_path="scene.glb", image_size=0)
|
||||
|
||||
with pytest.raises(ValueError, match="image_size"):
|
||||
create_habitat_simulator(config, habitat_sim_module=fake_habitat)
|
||||
|
||||
|
||||
def test_create_habitat_simulator_validates_move_forward_step():
|
||||
fake_habitat = _create_fake_habitat_module()
|
||||
config = HabitatSimulatorConfig(scene_path="scene.glb", move_forward_step=0)
|
||||
|
||||
with pytest.raises(ValueError, match="move_forward_step"):
|
||||
create_habitat_simulator(config, habitat_sim_module=fake_habitat)
|
||||
@@ -1,156 +0,0 @@
|
||||
from unittest.mock import Mock
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from utils.image import segment_image, segment_image_dataset
|
||||
|
||||
|
||||
def test_segment_image_passes_pil_image_to_mask_generator() -> None:
|
||||
mock_generator = Mock(return_value={"masks": []})
|
||||
|
||||
segment_image(
|
||||
mock_generator,
|
||||
Image.new("RGBA", (16, 16), color=(255, 0, 0, 255)),
|
||||
points_per_batch=32,
|
||||
)
|
||||
|
||||
image_arg = mock_generator.call_args.args[0]
|
||||
assert isinstance(image_arg, Image.Image)
|
||||
assert image_arg.mode == "RGB"
|
||||
assert mock_generator.call_args.kwargs["points_per_batch"] == 32
|
||||
|
||||
|
||||
def test_segment_image_supports_tensor_masks_output() -> None:
|
||||
masks_tensor = torch.tensor(
|
||||
[
|
||||
[
|
||||
[1, 1, 0],
|
||||
[1, 1, 0],
|
||||
[0, 0, 0],
|
||||
],
|
||||
[
|
||||
[1, 1, 1],
|
||||
[1, 1, 1],
|
||||
[1, 1, 1],
|
||||
],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
mock_generator = Mock(return_value={"masks": masks_tensor})
|
||||
|
||||
result = segment_image(
|
||||
mock_generator,
|
||||
Image.new("RGB", (3, 3), color=(0, 0, 0)),
|
||||
min_area=3,
|
||||
max_masks=5,
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0]["area"] == 9
|
||||
assert result[0]["bbox"] == [0, 0, 3, 3]
|
||||
assert result[1]["area"] == 4
|
||||
assert result[1]["bbox"] == [0, 0, 2, 2]
|
||||
|
||||
|
||||
def test_segment_image_filters_tensor_masks_by_min_area() -> None:
|
||||
masks_tensor = torch.tensor(
|
||||
[
|
||||
[
|
||||
[1, 0, 0],
|
||||
[0, 0, 0],
|
||||
[0, 0, 0],
|
||||
],
|
||||
[
|
||||
[1, 1, 0],
|
||||
[1, 1, 0],
|
||||
[0, 0, 0],
|
||||
],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
mock_generator = Mock(return_value={"masks": masks_tensor})
|
||||
|
||||
result = segment_image(
|
||||
mock_generator,
|
||||
Image.new("RGB", (3, 3), color=(0, 0, 0)),
|
||||
min_area=2,
|
||||
max_masks=5,
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["area"] == 4
|
||||
|
||||
|
||||
def test_segment_image_dataset_returns_per_image_masks_in_order() -> None:
|
||||
first_masks = torch.tensor(
|
||||
[[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
second_masks = torch.tensor(
|
||||
[[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
mock_generator = Mock(
|
||||
return_value=[{"masks": first_masks}, {"masks": second_masks}]
|
||||
)
|
||||
images = [
|
||||
Image.new("RGB", (3, 3), color=(0, 0, 0)),
|
||||
Image.new("RGB", (3, 3), color=(0, 0, 0)),
|
||||
]
|
||||
|
||||
result = segment_image_dataset(
|
||||
mock_generator,
|
||||
images,
|
||||
min_area=2,
|
||||
max_masks=5,
|
||||
points_per_batch=16,
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0][0]["area"] == 4
|
||||
assert result[1][0]["area"] == 9
|
||||
assert mock_generator.call_count == 1
|
||||
|
||||
|
||||
def test_segment_image_dataset_falls_back_to_single_image_calls() -> None:
|
||||
call_index = {"value": 0}
|
||||
|
||||
def fake_generator(images, points_per_batch):
|
||||
if isinstance(images, list):
|
||||
raise TypeError("Batch input unsupported")
|
||||
|
||||
result_options = [
|
||||
{
|
||||
"masks": torch.tensor(
|
||||
[[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
},
|
||||
{
|
||||
"masks": torch.tensor(
|
||||
[[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
},
|
||||
]
|
||||
out = result_options[call_index["value"]]
|
||||
call_index["value"] += 1
|
||||
return out
|
||||
|
||||
images = [
|
||||
Image.new("RGB", (3, 3), color=(0, 0, 0)),
|
||||
Image.new("RGB", (3, 3), color=(0, 0, 0)),
|
||||
]
|
||||
|
||||
result = segment_image_dataset(
|
||||
fake_generator,
|
||||
images,
|
||||
min_area=2,
|
||||
max_masks=5,
|
||||
points_per_batch=16,
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0][0]["area"] == 4
|
||||
assert result[1][0]["area"] == 9
|
||||
@@ -1,238 +0,0 @@
|
||||
"""Integration tests for multi-object retrieval benchmark pipeline.
|
||||
|
||||
These tests verify the end-to-end functionality of the multi-object retrieval
|
||||
benchmark, including schema building, database population, and evaluation.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class TestMultiObjectRetrievalIntegration:
|
||||
"""Integration tests for multi-object retrieval benchmark."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model_processor(self):
|
||||
"""Create mock model and processor."""
|
||||
mock_model = Mock()
|
||||
mock_processor = Mock()
|
||||
|
||||
# Mock the feature extraction to return a fixed-size vector
|
||||
def mock_extract(processor, model, image):
|
||||
return [0.1] * 256 # 256-dim vector
|
||||
mock_processor.images = mock_extract
|
||||
|
||||
return mock_model, mock_processor
|
||||
|
||||
@pytest.fixture
|
||||
def mock_dataset(self):
|
||||
"""Create a mock dataset with images and annotations."""
|
||||
# Create mock items
|
||||
items = []
|
||||
for i in range(3):
|
||||
item = {
|
||||
"image": Image.new("RGB", (224, 224), color=(i * 50, 100, 150)),
|
||||
"image_id": f"scene_{i}",
|
||||
"objects": {
|
||||
"bbox": [[10, 10, 50, 50], [60, 60, 40, 40]],
|
||||
"category": ["object_a", "object_b"],
|
||||
"area": [2500, 1600],
|
||||
"id": [0, 1],
|
||||
},
|
||||
}
|
||||
items.append(item)
|
||||
|
||||
mock_dataset = Mock()
|
||||
mock_dataset.__len__ = Mock(return_value=len(items))
|
||||
mock_dataset.__getitem__ = lambda self, idx: items[idx]
|
||||
mock_dataset.with_format = lambda fmt: mock_dataset
|
||||
|
||||
return mock_dataset
|
||||
|
||||
def test_build_object_schema(self):
|
||||
"""Test that object schema is built correctly."""
|
||||
from benchmarks.tasks.multi_object_retrieval import _build_object_schema
|
||||
import pyarrow as pa
|
||||
|
||||
vector_dim = 256
|
||||
schema = _build_object_schema(vector_dim)
|
||||
|
||||
assert isinstance(schema, pa.Schema)
|
||||
assert "id" in schema.names
|
||||
assert "image_id" in schema.names
|
||||
assert "object_id" in schema.names
|
||||
assert "category" in schema.names
|
||||
assert "vector" in schema.names
|
||||
|
||||
# Check vector field has correct dimension
|
||||
vector_field = schema.field("vector")
|
||||
assert isinstance(vector_field.type, pa.List)
|
||||
assert vector_field.type.value_type == pa.float32()
|
||||
|
||||
@patch("benchmarks.tasks.multi_object_retrieval.load_sam_model")
|
||||
@patch("benchmarks.tasks.multi_object_retrieval.segment_image")
|
||||
def test_build_database_with_mocked_sam(
|
||||
self,
|
||||
mock_segment,
|
||||
mock_load_sam,
|
||||
mock_model_processor,
|
||||
mock_dataset,
|
||||
):
|
||||
"""Test database building with mocked SAM segmentation."""
|
||||
from benchmarks.tasks.multi_object_retrieval import (
|
||||
MultiObjectRetrievalTask,
|
||||
_build_object_schema,
|
||||
)
|
||||
|
||||
mock_model, mock_processor = mock_model_processor
|
||||
|
||||
# Mock SAM
|
||||
mock_load_sam.return_value = (Mock(), Mock())
|
||||
mock_segment.return_value = [
|
||||
{
|
||||
"segment": np.ones((224, 224), dtype=bool),
|
||||
"area": 50000,
|
||||
"bbox": [0, 0, 224, 224],
|
||||
}
|
||||
]
|
||||
|
||||
# Create task with config
|
||||
task = MultiObjectRetrievalTask(
|
||||
sam_model="facebook/sam2.1-hiera-large",
|
||||
min_mask_area=1024,
|
||||
max_masks_per_image=5,
|
||||
gamma=1.0,
|
||||
top_k_per_object=50,
|
||||
num_query_objects=3,
|
||||
)
|
||||
|
||||
# Create mock table
|
||||
mock_table = Mock()
|
||||
mock_table.schema = _build_object_schema(256)
|
||||
|
||||
# Build database (this should not raise)
|
||||
task.build_database(mock_model, mock_processor, mock_dataset, mock_table, batch_size=1)
|
||||
|
||||
# Verify table.add was called
|
||||
assert mock_table.add.called
|
||||
|
||||
@patch("benchmarks.tasks.multi_object_retrieval.load_sam_model")
|
||||
@patch("benchmarks.tasks.multi_object_retrieval.segment_image")
|
||||
def test_evaluate_with_mocked_sam(
|
||||
self,
|
||||
mock_segment,
|
||||
mock_load_sam,
|
||||
mock_model_processor,
|
||||
mock_dataset,
|
||||
):
|
||||
"""Test evaluation with mocked SAM segmentation."""
|
||||
from benchmarks.tasks.multi_object_retrieval import (
|
||||
MultiObjectRetrievalTask,
|
||||
_build_object_schema,
|
||||
)
|
||||
|
||||
mock_model, mock_processor = mock_model_processor
|
||||
|
||||
# Mock SAM
|
||||
mock_load_sam.return_value = (Mock(), Mock())
|
||||
mock_segment.return_value = [
|
||||
{
|
||||
"segment": np.ones((224, 224), dtype=bool),
|
||||
"area": 50000,
|
||||
"bbox": [0, 0, 224, 224],
|
||||
"object_id": "query_obj_0",
|
||||
}
|
||||
]
|
||||
|
||||
# Create mock table with search results
|
||||
mock_table = Mock()
|
||||
mock_table.schema = _build_object_schema(256)
|
||||
|
||||
# Mock search to return matching result
|
||||
mock_result = Mock()
|
||||
mock_result.to_polars.return_value = {
|
||||
"image_id": ["scene_0"],
|
||||
"object_id": ["scene_0_obj_0"],
|
||||
"_distance": [0.1],
|
||||
}
|
||||
|
||||
mock_table.search.return_value.select.return_value.limit.return_value = mock_result
|
||||
|
||||
# Create task
|
||||
task = MultiObjectRetrievalTask(
|
||||
sam_model="facebook/sam2.1-hiera-large",
|
||||
min_mask_area=1024,
|
||||
max_masks_per_image=5,
|
||||
gamma=1.0,
|
||||
top_k_per_object=50,
|
||||
num_query_objects=1,
|
||||
)
|
||||
|
||||
# Evaluate
|
||||
results = task.evaluate(mock_model, mock_processor, mock_dataset, mock_table, batch_size=1)
|
||||
|
||||
# Verify results structure
|
||||
assert "accuracy" in results
|
||||
assert "correct" in results
|
||||
assert "total" in results
|
||||
assert "top_k" in results
|
||||
assert results["top_k"] == 50
|
||||
|
||||
def test_task_initialization_with_config(self):
|
||||
"""Test task initialization with custom config."""
|
||||
from benchmarks.tasks.multi_object_retrieval import MultiObjectRetrievalTask
|
||||
|
||||
task = MultiObjectRetrievalTask(
|
||||
sam_model="facebook/sam2.1-hiera-small",
|
||||
min_mask_area=500,
|
||||
max_masks_per_image=3,
|
||||
gamma=0.5,
|
||||
top_k_per_object=100,
|
||||
num_query_objects=5,
|
||||
)
|
||||
|
||||
assert task.sam_model == "facebook/sam2.1-hiera-small"
|
||||
assert task.min_mask_area == 500
|
||||
assert task.max_masks_per_image == 3
|
||||
assert task.config.gamma == 0.5
|
||||
assert task.config.top_k_per_object == 100
|
||||
assert task.config.num_query_objects == 5
|
||||
|
||||
def test_task_initialization_defaults(self):
|
||||
"""Test task initialization with default config."""
|
||||
from benchmarks.tasks.multi_object_retrieval import MultiObjectRetrievalTask
|
||||
|
||||
task = MultiObjectRetrievalTask()
|
||||
|
||||
# Check defaults from BenchmarkTaskConfig
|
||||
assert task.config.gamma == 1.0
|
||||
assert task.config.top_k_per_object == 50
|
||||
assert task.config.num_query_objects == 3
|
||||
# SAM settings from ModelConfig defaults
|
||||
assert task.sam_model == "facebook/sam2.1-hiera-large"
|
||||
assert task.min_mask_area == 1024
|
||||
assert task.max_masks_per_image == 5
|
||||
|
||||
|
||||
class TestInsDetScenesDataset:
|
||||
"""Tests for InsDetScenesDataset class."""
|
||||
|
||||
def test_dataset_class_exists(self):
|
||||
"""Test that InsDetScenesDataset can be imported."""
|
||||
from data_loading.insdet_scenes import InsDetScenesDataset
|
||||
|
||||
assert InsDetScenesDataset is not None
|
||||
|
||||
@patch("data_loading.insdet_scenes.load_val_dataset")
|
||||
def test_dataset_loads_correct_split(self, mock_load):
|
||||
"""Test dataset loads correct split."""
|
||||
from data_loading.insdet_scenes import InsDetScenesDataset
|
||||
|
||||
mock_load.return_value = Mock()
|
||||
|
||||
dataset = InsDetScenesDataset("/path/to/scenes", split="easy")
|
||||
|
||||
mock_load.assert_called_once_with("/path/to/scenes", "easy")
|
||||
assert dataset.split == "easy"
|
||||
@@ -1,113 +0,0 @@
|
||||
import numpy as np
|
||||
|
||||
from compressors.object_score import (
|
||||
MaskScoringConfig,
|
||||
compute_mask_features,
|
||||
rank_masks,
|
||||
score_mask,
|
||||
select_best_mask,
|
||||
)
|
||||
|
||||
|
||||
def _rect_mask(height: int, width: int, x: int, y: int, w: int, h: int) -> np.ndarray:
|
||||
mask = np.zeros((height, width), dtype=bool)
|
||||
mask[y : y + h, x : x + w] = True
|
||||
return mask
|
||||
|
||||
|
||||
def test_compute_mask_features_core_metrics() -> None:
|
||||
mask = _rect_mask(height=20, width=20, x=5, y=4, w=6, h=5)
|
||||
mask_dict = {
|
||||
"segment": mask,
|
||||
"area": int(mask.sum()),
|
||||
"bbox": [5, 4, 6, 5],
|
||||
"predicted_iou": 0.8,
|
||||
"stability_score": 0.9,
|
||||
}
|
||||
|
||||
features = compute_mask_features(mask_dict, image_shape=(20, 20))
|
||||
|
||||
assert features.area_ratio == 30 / 400
|
||||
assert features.fill_ratio == 1.0
|
||||
assert features.aspect_ratio == 6 / 5
|
||||
assert features.touch_top is False
|
||||
assert features.touch_left is False
|
||||
assert features.num_components == 1
|
||||
assert features.largest_component_ratio == 1.0
|
||||
assert features.num_holes == 0
|
||||
|
||||
|
||||
def test_rank_masks_rejects_extreme_small_and_fragmented_masks() -> None:
|
||||
cfg = MaskScoringConfig(min_area_ratio=0.02)
|
||||
good_mask = _rect_mask(height=30, width=30, x=6, y=6, w=10, h=10)
|
||||
|
||||
fragmented = np.zeros((30, 30), dtype=bool)
|
||||
fragmented[2, 2] = True
|
||||
fragmented[4, 7] = True
|
||||
fragmented[8, 12] = True
|
||||
fragmented[12, 16] = True
|
||||
fragmented[16, 20] = True
|
||||
fragmented[20, 24] = True
|
||||
fragmented[24, 26] = True
|
||||
|
||||
masks = [
|
||||
{"segment": np.zeros((30, 30), dtype=bool), "area": 1, "bbox": [0, 0, 1, 1]},
|
||||
{
|
||||
"segment": fragmented,
|
||||
"area": int(fragmented.sum()),
|
||||
"bbox": [2, 2, 25, 25],
|
||||
},
|
||||
{
|
||||
"segment": good_mask,
|
||||
"area": int(good_mask.sum()),
|
||||
"bbox": [6, 6, 10, 10],
|
||||
"predicted_iou": 0.9,
|
||||
"stability_score": 0.9,
|
||||
},
|
||||
]
|
||||
|
||||
ranked = rank_masks(masks=masks, image_shape=(30, 30), config=cfg, max_masks=3)
|
||||
|
||||
assert len(ranked) == 1
|
||||
assert ranked[0]["area"] == int(good_mask.sum())
|
||||
assert "mask_score" in ranked[0]
|
||||
|
||||
|
||||
def test_score_mask_prefers_stable_reasonable_object() -> None:
|
||||
cfg = MaskScoringConfig()
|
||||
|
||||
candidate = {
|
||||
"segment": _rect_mask(height=100, width=100, x=30, y=20, w=24, h=25),
|
||||
"area": 24 * 25,
|
||||
"bbox": [30, 20, 24, 25],
|
||||
"predicted_iou": 0.92,
|
||||
"stability_score": 0.91,
|
||||
}
|
||||
weak = {
|
||||
"segment": _rect_mask(height=100, width=100, x=0, y=0, w=4, h=60),
|
||||
"area": 4 * 60,
|
||||
"bbox": [0, 0, 4, 60],
|
||||
"predicted_iou": 0.4,
|
||||
"stability_score": 0.3,
|
||||
}
|
||||
|
||||
score_candidate = score_mask(candidate, image_shape=(100, 100), config=cfg)
|
||||
score_weak = score_mask(weak, image_shape=(100, 100), config=cfg)
|
||||
|
||||
assert score_candidate > score_weak
|
||||
|
||||
|
||||
def test_select_best_mask_falls_back_to_largest_area_when_all_rejected() -> None:
|
||||
cfg = MaskScoringConfig(min_area_ratio=0.2)
|
||||
tiny = _rect_mask(height=20, width=20, x=1, y=1, w=2, h=2)
|
||||
larger = _rect_mask(height=20, width=20, x=5, y=5, w=4, h=4)
|
||||
|
||||
masks = [
|
||||
{"segment": tiny, "area": int(tiny.sum()), "bbox": [1, 1, 2, 2]},
|
||||
{"segment": larger, "area": int(larger.sum()), "bbox": [5, 5, 4, 4]},
|
||||
]
|
||||
|
||||
best = select_best_mask(masks=masks, image_shape=(20, 20), config=cfg)
|
||||
|
||||
assert best is not None
|
||||
assert best["area"] == int(larger.sum())
|
||||
@@ -1,108 +0,0 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from simulator import collect_room_views_by_room
|
||||
|
||||
|
||||
class _FakeAgent:
|
||||
def __init__(self):
|
||||
self.positions = []
|
||||
|
||||
def set_state(self, state):
|
||||
self.positions.append(state.position)
|
||||
|
||||
|
||||
class _FakeSimulator:
|
||||
def __init__(self):
|
||||
self._frame_index = 0
|
||||
self.actions = []
|
||||
|
||||
def get_sensor_observations(self):
|
||||
observations = {
|
||||
"color_sensor": f"frame_{self._frame_index}",
|
||||
"depth_sensor": f"depth_{self._frame_index}",
|
||||
}
|
||||
self._frame_index += 1
|
||||
return observations
|
||||
|
||||
def step(self, action_name):
|
||||
self.actions.append(action_name)
|
||||
|
||||
|
||||
class _FakeAgentState:
|
||||
def __init__(self):
|
||||
self.position = None
|
||||
|
||||
|
||||
def test_collect_room_views_by_room_collects_grouped_frames_with_single_outer_progress():
|
||||
track_calls = []
|
||||
|
||||
def fake_track(iterable, description):
|
||||
track_calls.append(description)
|
||||
return iterable
|
||||
|
||||
agent = _FakeAgent()
|
||||
sim = _FakeSimulator()
|
||||
room_nodes = [
|
||||
SimpleNamespace(room_id="room_00", center=[1.0, 2.0, 3.0]),
|
||||
SimpleNamespace(room_id="room_01", center=[4.0, 5.0, 6.0]),
|
||||
]
|
||||
fake_habitat = SimpleNamespace(AgentState=_FakeAgentState)
|
||||
|
||||
room_views = collect_room_views_by_room(
|
||||
agent=agent,
|
||||
sim=sim,
|
||||
room_nodes=room_nodes,
|
||||
views_per_room=3,
|
||||
habitat_sim_module=fake_habitat,
|
||||
progress_track=fake_track,
|
||||
)
|
||||
|
||||
assert track_calls == ["Collecting room views"]
|
||||
assert room_views == {
|
||||
"room_00": ["frame_0", "frame_1", "frame_2"],
|
||||
"room_01": ["frame_3", "frame_4", "frame_5"],
|
||||
}
|
||||
assert sim.actions == [
|
||||
"turn_left",
|
||||
"turn_left",
|
||||
"turn_left",
|
||||
"turn_left",
|
||||
"turn_left",
|
||||
"turn_left",
|
||||
]
|
||||
assert agent.positions == [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
|
||||
|
||||
|
||||
def test_collect_room_views_by_room_uses_custom_sensor_and_turn_action():
|
||||
agent = _FakeAgent()
|
||||
sim = _FakeSimulator()
|
||||
room_nodes = [SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])]
|
||||
fake_habitat = SimpleNamespace(AgentState=_FakeAgentState)
|
||||
|
||||
room_views = collect_room_views_by_room(
|
||||
agent=agent,
|
||||
sim=sim,
|
||||
room_nodes=room_nodes,
|
||||
views_per_room=2,
|
||||
habitat_sim_module=fake_habitat,
|
||||
sensor_uuid="depth_sensor",
|
||||
turn_action="turn_right",
|
||||
progress_track=lambda iterable, description: iterable,
|
||||
)
|
||||
|
||||
assert room_views == {"room_00": ["depth_0", "depth_1"]}
|
||||
assert sim.actions == ["turn_right", "turn_right"]
|
||||
|
||||
|
||||
def test_collect_room_views_by_room_validates_views_per_room():
|
||||
with pytest.raises(ValueError, match="views_per_room"):
|
||||
collect_room_views_by_room(
|
||||
agent=_FakeAgent(),
|
||||
sim=_FakeSimulator(),
|
||||
room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])],
|
||||
views_per_room=0,
|
||||
habitat_sim_module=SimpleNamespace(AgentState=_FakeAgentState),
|
||||
progress_track=lambda iterable, description: iterable,
|
||||
)
|
||||
@@ -1,205 +0,0 @@
|
||||
"""Tests for SAM segmentation utilities.
|
||||
|
||||
Note: These tests mock the SAM model loading since SAM requires
|
||||
heavy model weights. The actual SAM integration should be tested
|
||||
separately in integration tests.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class TestSAMSegmentation:
|
||||
"""Test suite for SAM segmentation utilities."""
|
||||
|
||||
def test_segment_image_empty_masks(self):
|
||||
"""Test segment_image returns empty list when no masks generated."""
|
||||
from utils.sam import segment_image
|
||||
|
||||
# Create mock mask generator that returns empty list
|
||||
mock_generator = Mock()
|
||||
mock_generator.generate.return_value = []
|
||||
|
||||
result = segment_image(mock_generator, Image.new("RGB", (100, 100)))
|
||||
|
||||
assert result == []
|
||||
|
||||
def test_segment_image_filters_small_masks(self):
|
||||
"""Test segment_image filters masks below min_area threshold."""
|
||||
from utils.sam import segment_image
|
||||
|
||||
# Create mock masks with different areas
|
||||
small_mask = {
|
||||
"segment": np.zeros((10, 10), dtype=bool),
|
||||
"area": 50, # Below 32*32 = 1024
|
||||
"bbox": [0, 0, 10, 10],
|
||||
"predicted_iou": 0.9,
|
||||
"stability_score": 0.8,
|
||||
}
|
||||
large_mask = {
|
||||
"segment": np.ones((100, 100), dtype=bool),
|
||||
"area": 10000, # Above threshold
|
||||
"bbox": [0, 0, 100, 100],
|
||||
"predicted_iou": 0.95,
|
||||
"stability_score": 0.9,
|
||||
}
|
||||
|
||||
mock_generator = Mock()
|
||||
mock_generator.generate.return_value = [small_mask, large_mask]
|
||||
|
||||
result = segment_image(
|
||||
mock_generator,
|
||||
Image.new("RGB", (100, 100)),
|
||||
min_area=32 * 32,
|
||||
max_masks=5,
|
||||
)
|
||||
|
||||
# Should only return the large mask
|
||||
assert len(result) == 1
|
||||
assert result[0]["area"] == 10000
|
||||
|
||||
def test_segment_image_limits_max_masks(self):
|
||||
"""Test segment_image limits to max_masks largest masks."""
|
||||
from utils.sam import segment_image
|
||||
|
||||
# Create 10 masks with different areas
|
||||
masks = [
|
||||
{
|
||||
"segment": np.ones((i + 1, i + 1), dtype=bool),
|
||||
"area": (i + 1) * (i + 1),
|
||||
"bbox": [0, 0, i + 1, i + 1],
|
||||
"predicted_iou": 0.9,
|
||||
"stability_score": 0.8,
|
||||
}
|
||||
for i in range(10)
|
||||
]
|
||||
|
||||
mock_generator = Mock()
|
||||
mock_generator.generate.return_value = masks
|
||||
|
||||
result = segment_image(
|
||||
mock_generator,
|
||||
Image.new("RGB", (100, 100)),
|
||||
min_area=1,
|
||||
max_masks=3,
|
||||
)
|
||||
|
||||
# Should only return top 3 largest masks
|
||||
assert len(result) == 3
|
||||
# Check they are sorted by area (largest first)
|
||||
areas = [m["area"] for m in result]
|
||||
assert areas == sorted(areas, reverse=True)
|
||||
|
||||
def test_segment_image_sorted_by_area(self):
|
||||
"""Test segment_image returns masks sorted by area descending."""
|
||||
from utils.sam import segment_image
|
||||
|
||||
# Create masks with known areas (unordered)
|
||||
mask1 = {
|
||||
"segment": np.ones((5, 5), dtype=bool),
|
||||
"area": 25,
|
||||
"bbox": [0, 0, 5, 5],
|
||||
}
|
||||
mask2 = {
|
||||
"segment": np.ones((10, 10), dtype=bool),
|
||||
"area": 100,
|
||||
"bbox": [0, 0, 10, 10],
|
||||
}
|
||||
mask3 = {
|
||||
"segment": np.ones((3, 3), dtype=bool),
|
||||
"area": 9,
|
||||
"bbox": [0, 0, 3, 3],
|
||||
}
|
||||
|
||||
mock_generator = Mock()
|
||||
mock_generator.generate.return_value = [mask1, mask2, mask3]
|
||||
|
||||
result = segment_image(
|
||||
mock_generator,
|
||||
Image.new("RGB", (100, 100)),
|
||||
min_area=1,
|
||||
max_masks=10,
|
||||
)
|
||||
|
||||
# Should be sorted by area descending
|
||||
assert result[0]["area"] == 100
|
||||
assert result[1]["area"] == 25
|
||||
assert result[2]["area"] == 9
|
||||
|
||||
|
||||
class TestSAMLoading:
|
||||
@patch("utils.sam.pipeline")
|
||||
def test_load_sam_model_uses_transformers_pipeline(self, mock_pipeline):
|
||||
from utils.sam import Sam2MaskGenerator, load_sam_model
|
||||
|
||||
mock_pipe_obj = Mock()
|
||||
mock_pipe_obj.model = Mock()
|
||||
mock_pipeline.return_value = mock_pipe_obj
|
||||
|
||||
sam_model, mask_generator = load_sam_model(
|
||||
model_name="facebook/sam2.1-hiera-large",
|
||||
device="cpu",
|
||||
points_per_batch=16,
|
||||
)
|
||||
|
||||
assert sam_model is mock_pipe_obj.model
|
||||
assert isinstance(mask_generator, Sam2MaskGenerator)
|
||||
assert mask_generator.points_per_batch == 16
|
||||
|
||||
_, kwargs = mock_pipeline.call_args
|
||||
assert kwargs["task"] == "mask-generation"
|
||||
assert kwargs["model"] == "facebook/sam2.1-hiera-large"
|
||||
assert kwargs["device"] == -1
|
||||
|
||||
|
||||
class TestExtractMaskedRegion:
|
||||
"""Test suite for extracting masked regions from images."""
|
||||
|
||||
def test_extract_masked_region_binary(self):
|
||||
"""Test extracting masked region with binary mask."""
|
||||
from utils.sam import extract_masked_region
|
||||
|
||||
# Create a simple image
|
||||
image = Image.new("RGB", (10, 10), color=(255, 0, 0))
|
||||
|
||||
# Create a binary mask (half kept, half masked)
|
||||
mask = np.zeros((10, 10), dtype=bool)
|
||||
mask[:, :5] = True
|
||||
|
||||
result = extract_masked_region(image, mask)
|
||||
|
||||
# Check that left half is red, right half is black
|
||||
result_np = np.array(result)
|
||||
left_half = result_np[:, :5, :]
|
||||
right_half = result_np[:, 5:, :]
|
||||
|
||||
assert np.all(left_half == [255, 0, 0])
|
||||
assert np.all(right_half == [0, 0, 0])
|
||||
|
||||
def test_extract_masked_region_all_masked(self):
|
||||
"""Test extracting when entire image is masked."""
|
||||
from utils.sam import extract_masked_region
|
||||
|
||||
image = Image.new("RGB", (10, 10), color=(255, 0, 0))
|
||||
mask = np.ones((10, 10), dtype=bool)
|
||||
|
||||
result = extract_masked_region(image, mask)
|
||||
result_np = np.array(result)
|
||||
|
||||
# Entire image should be preserved
|
||||
assert np.all(result_np == [255, 0, 0])
|
||||
|
||||
def test_extract_masked_region_all_zero_mask(self):
|
||||
"""Test extracting when mask is all zeros."""
|
||||
from utils.sam import extract_masked_region
|
||||
|
||||
image = Image.new("RGB", (10, 10), color=(255, 0, 0))
|
||||
mask = np.zeros((10, 10), dtype=bool)
|
||||
|
||||
result = extract_masked_region(image, mask)
|
||||
result_np = np.array(result)
|
||||
|
||||
# Entire image should be black
|
||||
assert np.all(result_np == [0, 0, 0])
|
||||
@@ -1,121 +0,0 @@
|
||||
"""Tests for scene scoring algorithm in multi-object retrieval."""
|
||||
|
||||
import pytest
|
||||
from benchmarks.tasks.multi_object_retrieval import _compute_scene_score
|
||||
|
||||
|
||||
class TestSceneScoringAlgorithm:
|
||||
"""Test suite for scene scoring with co-occurrence penalty."""
|
||||
|
||||
def test_scene_score_basic(self):
|
||||
"""Test basic scene scoring with single match."""
|
||||
query_object_ids = ["obj_1", "obj_2", "obj_3"]
|
||||
|
||||
# Scene A has obj_1
|
||||
retrieved_results = {
|
||||
"scene_A": [("distance_1", "obj_1")],
|
||||
}
|
||||
|
||||
scores = _compute_scene_score(query_object_ids, retrieved_results, gamma=1.0)
|
||||
|
||||
# Hit rate = 1/3, similarity = 1/(1+distance_1)
|
||||
assert "scene_A" in scores
|
||||
assert scores["scene_A"] > 0
|
||||
|
||||
def test_scene_score_no_match(self):
|
||||
"""Test scene scoring when no objects match."""
|
||||
query_object_ids = ["obj_1", "obj_2", "obj_3"]
|
||||
|
||||
retrieved_results = {
|
||||
"scene_A": [("distance_1", "other_obj")],
|
||||
}
|
||||
|
||||
scores = _compute_scene_score(query_object_ids, retrieved_results, gamma=1.0)
|
||||
|
||||
assert scores["scene_A"] == 0.0
|
||||
|
||||
def test_scene_score_multiple_scenes(self):
|
||||
"""Test scoring across multiple scenes."""
|
||||
query_object_ids = ["obj_1", "obj_2"]
|
||||
|
||||
retrieved_results = {
|
||||
"scene_A": [("0.1", "obj_1")],
|
||||
"scene_B": [("0.1", "obj_2")],
|
||||
"scene_C": [("0.1", "other")],
|
||||
}
|
||||
|
||||
scores = _compute_scene_score(query_object_ids, retrieved_results, gamma=1.0)
|
||||
|
||||
# Scenes with matches should have positive scores
|
||||
assert scores["scene_A"] > 0
|
||||
assert scores["scene_B"] > 0
|
||||
# Scene C has no match, score should be 0
|
||||
assert scores["scene_C"] == 0.0
|
||||
|
||||
def test_scene_score_gamma_zero(self):
|
||||
"""Test scoring with gamma=0 (no penalty)."""
|
||||
query_object_ids = ["obj_1", "obj_2", "obj_3", "obj_4", "obj_5"]
|
||||
|
||||
retrieved_results = {
|
||||
"scene_A": [("0.1", "obj_1")],
|
||||
}
|
||||
|
||||
scores_gamma_0 = _compute_scene_score(query_object_ids, retrieved_results, gamma=0.0)
|
||||
scores_gamma_1 = _compute_scene_score(query_object_ids, retrieved_results, gamma=1.0)
|
||||
|
||||
# With gamma=0, hit_rate^0 = 1, so score = similarity
|
||||
# With gamma=1, hit_rate^1 = 1/5, so score = similarity * 1/5
|
||||
# scores_gamma_0 should be larger
|
||||
assert scores_gamma_0["scene_A"] > scores_gamma_1["scene_A"]
|
||||
|
||||
def test_scene_score_multiple_matches(self):
|
||||
"""Test scoring when scene has multiple matching objects."""
|
||||
query_object_ids = ["obj_1", "obj_2"]
|
||||
|
||||
retrieved_results = {
|
||||
"scene_A": [("0.1", "obj_1"), ("0.2", "obj_2")],
|
||||
}
|
||||
|
||||
scores = _compute_scene_score(query_object_ids, retrieved_results, gamma=1.0)
|
||||
|
||||
# Both objects match, hit_rate = 2/2 = 1.0
|
||||
# Score = (1/(1+0.1) + 1/(1+0.2)) * 1.0
|
||||
expected_similarity = 1.0 / 1.1 + 1.0 / 1.2
|
||||
assert abs(scores["scene_A"] - expected_similarity) < 0.01
|
||||
|
||||
def test_scene_score_distance_to_similarity(self):
|
||||
"""Test that smaller distance yields higher score."""
|
||||
query_object_ids = ["obj_1"]
|
||||
|
||||
retrieved_results = {
|
||||
"scene_close": [("0.01", "obj_1")],
|
||||
"scene_far": [("10.0", "obj_1")],
|
||||
}
|
||||
|
||||
scores = _compute_scene_score(query_object_ids, retrieved_results, gamma=1.0)
|
||||
|
||||
# Closer scene should have higher score
|
||||
assert scores["scene_close"] > scores["scene_far"]
|
||||
|
||||
def test_scene_score_empty_results(self):
|
||||
"""Test scoring with empty retrieved results."""
|
||||
query_object_ids = ["obj_1", "obj_2"]
|
||||
|
||||
retrieved_results = {}
|
||||
|
||||
scores = _compute_scene_score(query_object_ids, retrieved_results, gamma=1.0)
|
||||
|
||||
assert scores == {}
|
||||
|
||||
def test_scene_score_empty_query(self):
|
||||
"""Test scoring with empty query objects."""
|
||||
query_object_ids = []
|
||||
|
||||
retrieved_results = {
|
||||
"scene_A": [("0.1", "obj_1")],
|
||||
}
|
||||
|
||||
scores = _compute_scene_score(query_object_ids, retrieved_results, gamma=1.0)
|
||||
|
||||
# With empty query, no scenes should have positive score
|
||||
assert all(score == 0.0 for score in scores.values())
|
||||
@@ -1,126 +0,0 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from simulator import (
|
||||
TopDownRenderStyle,
|
||||
TopDownSceneElements,
|
||||
render_topdown_scene_map,
|
||||
)
|
||||
|
||||
|
||||
class _FakeMaps:
|
||||
def __init__(self):
|
||||
self.to_grid_calls: list[tuple[float, float]] = []
|
||||
|
||||
def get_topdown_map(self, pathfinder, height, meters_per_pixel):
|
||||
return [[0, 0], [0, 0]]
|
||||
|
||||
def to_grid(self, z, x, shape, pathfinder):
|
||||
self.to_grid_calls.append((z, x))
|
||||
return (int(z), int(x))
|
||||
|
||||
|
||||
class _FakePlt:
|
||||
def __init__(self):
|
||||
self.scatter_calls: list[tuple[int, int]] = []
|
||||
self.text_calls: list[str] = []
|
||||
self.shown = False
|
||||
|
||||
def figure(self, figsize):
|
||||
return None
|
||||
|
||||
def imshow(self, image, cmap):
|
||||
return None
|
||||
|
||||
def scatter(self, x, y, c, s):
|
||||
self.scatter_calls.append((x, y))
|
||||
|
||||
def text(self, x, y, text, color, fontsize):
|
||||
self.text_calls.append(text)
|
||||
|
||||
def title(self, title):
|
||||
return None
|
||||
|
||||
def axis(self, mode):
|
||||
return None
|
||||
|
||||
def show(self):
|
||||
self.shown = True
|
||||
|
||||
|
||||
def test_render_topdown_scene_map_renders_room_nodes_only():
|
||||
fake_maps = _FakeMaps()
|
||||
fake_plt = _FakePlt()
|
||||
room_nodes = [
|
||||
SimpleNamespace(room_id="room_00", center=[1.0, 2.0, 3.0]),
|
||||
SimpleNamespace(room_id="room_01", center=[4.0, 2.0, 5.0]),
|
||||
]
|
||||
elements = TopDownSceneElements(room_nodes=room_nodes)
|
||||
|
||||
top_down_map = render_topdown_scene_map(
|
||||
pathfinder=SimpleNamespace(),
|
||||
elements=elements,
|
||||
meters_per_pixel=0.05,
|
||||
style=TopDownRenderStyle(),
|
||||
maps_module=fake_maps,
|
||||
plt_module=fake_plt,
|
||||
)
|
||||
|
||||
assert top_down_map == [[0, 0], [0, 0]]
|
||||
assert fake_maps.to_grid_calls == [(3.0, 1.0), (5.0, 4.0)]
|
||||
assert fake_plt.scatter_calls == [(1, 3), (4, 5)]
|
||||
assert fake_plt.text_calls == ["room_00", "room_01"]
|
||||
assert fake_plt.shown is True
|
||||
|
||||
|
||||
def test_render_topdown_scene_map_validates_room_nodes():
|
||||
with pytest.raises(ValueError, match="room_nodes"):
|
||||
render_topdown_scene_map(
|
||||
pathfinder=SimpleNamespace(),
|
||||
elements=TopDownSceneElements(room_nodes=[]),
|
||||
meters_per_pixel=0.05,
|
||||
maps_module=_FakeMaps(),
|
||||
plt_module=_FakePlt(),
|
||||
)
|
||||
|
||||
|
||||
def test_render_topdown_scene_map_validates_meters_per_pixel():
|
||||
with pytest.raises(ValueError, match="meters_per_pixel"):
|
||||
render_topdown_scene_map(
|
||||
pathfinder=SimpleNamespace(),
|
||||
elements=TopDownSceneElements(
|
||||
room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])]
|
||||
),
|
||||
meters_per_pixel=0,
|
||||
maps_module=_FakeMaps(),
|
||||
plt_module=_FakePlt(),
|
||||
)
|
||||
|
||||
|
||||
def test_render_topdown_scene_map_rejects_object_nodes_before_implementation():
|
||||
with pytest.raises(NotImplementedError, match="object_nodes"):
|
||||
render_topdown_scene_map(
|
||||
pathfinder=SimpleNamespace(),
|
||||
elements=TopDownSceneElements(
|
||||
room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])],
|
||||
object_nodes=[SimpleNamespace(obj_id="obj_00")],
|
||||
),
|
||||
meters_per_pixel=0.05,
|
||||
maps_module=_FakeMaps(),
|
||||
plt_module=_FakePlt(),
|
||||
)
|
||||
|
||||
|
||||
def test_render_topdown_scene_map_rejects_edges_before_implementation():
|
||||
with pytest.raises(NotImplementedError, match="edge"):
|
||||
render_topdown_scene_map(
|
||||
pathfinder=SimpleNamespace(),
|
||||
elements=TopDownSceneElements(
|
||||
room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])],
|
||||
edges=[("room_00", "obj_00")],
|
||||
),
|
||||
meters_per_pixel=0.05,
|
||||
maps_module=_FakeMaps(),
|
||||
plt_module=_FakePlt(),
|
||||
)
|
||||
@@ -1,77 +0,0 @@
|
||||
"""Tests for visualizer app image upload similarity search."""
|
||||
|
||||
import base64
|
||||
import io
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from sklearn.metrics.pairwise import cosine_similarity
|
||||
|
||||
|
||||
class TestImageUploadSimilaritySearch:
|
||||
"""Test suite for image upload similarity search functionality."""
|
||||
|
||||
def test_base64_to_pil_image(self):
|
||||
"""Test conversion from base64 string to PIL Image."""
|
||||
# Create a test image
|
||||
img_array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
||||
img = Image.fromarray(img_array)
|
||||
|
||||
# Convert to base64
|
||||
buffer = io.BytesIO()
|
||||
img.save(buffer, format="PNG")
|
||||
img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
|
||||
# Add data URI prefix (as Dash provides)
|
||||
img_base64_with_prefix = f"data:image/png;base64,{img_base64}"
|
||||
|
||||
# Parse base64 to PIL Image
|
||||
# Remove prefix
|
||||
base64_str = img_base64_with_prefix.split(",")[1]
|
||||
img_bytes = base64.b64decode(base64_str)
|
||||
parsed_img = Image.open(io.BytesIO(img_bytes))
|
||||
|
||||
# Verify the image is valid
|
||||
assert parsed_img.size == (224, 224)
|
||||
assert parsed_img.mode == "RGB"
|
||||
|
||||
|
||||
class TestCosineSimilarity:
|
||||
"""Test suite for cosine similarity computation between feature vectors."""
|
||||
|
||||
def test_identical_vectors_return_one(self):
|
||||
"""Identical vectors should have cosine similarity of 1.0."""
|
||||
vec = np.random.randn(1024).tolist()
|
||||
similarity = cosine_similarity([vec], [vec])[0][0]
|
||||
assert np.isclose(similarity, 1.0)
|
||||
|
||||
def test_orthogonal_vectors_return_zero(self):
|
||||
"""Orthogonal vectors should have cosine similarity of 0.0."""
|
||||
vec_a = [1.0, 0.0]
|
||||
vec_b = [0.0, 1.0]
|
||||
similarity = cosine_similarity([vec_a], [vec_b])[0][0]
|
||||
assert np.isclose(similarity, 0.0)
|
||||
|
||||
def test_opposite_vectors_return_negative_one(self):
|
||||
"""Opposite vectors should have cosine similarity of -1.0."""
|
||||
vec_a = [1.0, 0.0, 0.0]
|
||||
vec_b = [-1.0, 0.0, 0.0]
|
||||
similarity = cosine_similarity([vec_a], [vec_b])[0][0]
|
||||
assert np.isclose(similarity, -1.0)
|
||||
|
||||
def test_similarity_range(self):
|
||||
"""Cosine similarity should always be within [-1, 1]."""
|
||||
# Random vectors
|
||||
for _ in range(10):
|
||||
vec_a = np.random.randn(1024).tolist()
|
||||
vec_b = np.random.randn(1024).tolist()
|
||||
similarity = cosine_similarity([vec_a], [vec_b])[0][0]
|
||||
assert -1.0 <= similarity <= 1.0
|
||||
|
||||
def test_similarity_with_list_input(self):
|
||||
"""Cosine similarity should work with Python list inputs (as stored in dcc.Store)."""
|
||||
# Simulate feature vectors stored as Python lists in dcc.Store
|
||||
vec_a = [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
vec_b = [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
similarity = cosine_similarity([vec_a], [vec_b])[0][0]
|
||||
assert np.isclose(similarity, 1.0)
|
||||
Reference in New Issue
Block a user