mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
refactor(pipeline): integrate SAM segmentation and modularize model loading
This commit is contained in:
@@ -2,6 +2,7 @@
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from unittest.mock import Mock, patch
|
||||
from compressors import (
|
||||
BinarySign,
|
||||
HashCompressor,
|
||||
@@ -205,18 +206,54 @@ class TestVideoPositiveMask:
|
||||
class TestHashPipeline:
|
||||
"""Test suite for HashPipeline."""
|
||||
|
||||
def test_pipeline_init(self):
|
||||
@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()
|
||||
|
||||
def test_pipeline_hash_bits(self):
|
||||
@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
|
||||
|
||||
@@ -228,14 +265,33 @@ class TestHashPipeline:
|
||||
class TestConfigIntegration:
|
||||
"""Test suite for config integration with pipeline."""
|
||||
|
||||
def test_create_pipeline_from_config(self):
|
||||
@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."""
|
||||
|
||||
@@ -97,9 +97,21 @@ class TestSAMSegmentation:
|
||||
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]}
|
||||
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]
|
||||
@@ -117,6 +129,31 @@ class TestSAMSegmentation:
|
||||
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."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user