Team Ai
Apppublic

SahilCodevally/codevally-vision-language-action

sourceHugging Faceotherupdated 7mo agoView on Hugging Face
0likes
test_visualizer.py107 linesDownload Raw Back to tests
1"""2Tests for the Visualization Module — ActionVisualizer.3"""4 5import numpy as np6import pytest7 8from src.visualization.visualizer import ActionVisualizer9 10 11class TestActionVisualizer:12    """Tests for ActionVisualizer methods."""13 14    @pytest.fixture15    def visualizer(self) -> ActionVisualizer:16        return ActionVisualizer()17 18    @pytest.fixture19    def test_image(self) -> np.ndarray:20        """A 480x640 synthetic test image."""21        return np.zeros((480, 640, 3), dtype=np.uint8)22 23    @pytest.fixture24    def sample_detections(self) -> list[dict]:25        return [26            {"label": "bottle", "confidence": 0.92, "bbox": [100, 150, 200, 350]},27            {"label": "box", "confidence": 0.85, "bbox": [400, 300, 550, 500]},28        ]29 30    @pytest.fixture31    def sample_action_plan(self) -> dict:32        return {33            "action": "move",34            "steps": [35                "Step 1: Locate 'bottle'",36                "Step 2: Locate 'box'",37                "Step 3: Move 'bottle' → 'box'",38                "Step 4: Verify action completion",39            ],40            "source_bbox": [100, 150, 200, 350],41            "dest_bbox": [400, 300, 550, 500],42            "status": "success",43        }44 45    def test_draw_detections_returns_image(46        self, visualizer, test_image, sample_detections47    ):48        """draw_detections should return an image of the same shape."""49        result = visualizer.draw_detections(test_image, sample_detections)50 51        assert isinstance(result, np.ndarray)52        assert result.shape == test_image.shape53 54    def test_draw_detections_modifies_image(55        self, visualizer, test_image, sample_detections56    ):57        """draw_detections should modify the image (not return a blank copy)."""58        result = visualizer.draw_detections(test_image, sample_detections)59 60        # The annotated image should differ from the original61        assert not np.array_equal(result, test_image)62 63    def test_draw_detections_empty_list(self, visualizer, test_image):64        """draw_detections with empty list should still return an image."""65        result = visualizer.draw_detections(test_image, [])66 67        assert isinstance(result, np.ndarray)68        assert result.shape == test_image.shape69 70    def test_visualize_returns_image(71        self, visualizer, test_image, sample_action_plan72    ):73        """visualize should return an image of the same shape."""74        result = visualizer.visualize(test_image, sample_action_plan)75 76        assert isinstance(result, np.ndarray)77        assert result.shape == test_image.shape78 79    def test_visualize_modifies_image(80        self, visualizer, test_image, sample_action_plan81    ):82        """visualize should modify the image."""83        result = visualizer.visualize(test_image, sample_action_plan)84 85        assert not np.array_equal(result, test_image)86 87    def test_visualize_with_no_bboxes(self, visualizer, test_image):88        """visualize should handle missing bounding boxes gracefully."""89        plan = {90            "action": "inspect",91            "steps": ["Step 1: Inspect object"],92            "source_bbox": None,93            "dest_bbox": None,94            "status": "failed",95        }96 97        result = visualizer.visualize(test_image, plan)98 99        assert isinstance(result, np.ndarray)100        assert result.shape == test_image.shape101 102    def test_draw_detections_invalid_image(self, visualizer):103        """draw_detections with None should return a fallback image."""104        result = visualizer.draw_detections(None, [])105 106        assert isinstance(result, np.ndarray)107