Spaces:
Running
Running
File size: 2,184 Bytes
23db765 cf388f7 23db765 cf388f7 23db765 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """Tests for the visualization module."""
import pytest
from PIL import Image
from src.parsing import BBox
from src.visualization import create_no_detection_overlay, draw_boxes
@pytest.fixture
def sample_image():
return Image.new("RGB", (640, 480), color=(30, 30, 60))
@pytest.fixture
def sample_boxes():
return [
BBox(x1=50, y1=50, x2=200, y2=150, confidence=0.9, label="debris"),
BBox(x1=300, y1=200, x2=450, y2=350, confidence=0.75, label="satellite"),
]
class TestDrawBoxes:
def test_returns_image(self, sample_image, sample_boxes):
result = draw_boxes(sample_image, sample_boxes)
assert isinstance(result, Image.Image)
assert result.size == sample_image.size
def test_empty_boxes_returns_copy(self, sample_image):
result = draw_boxes(sample_image, [])
assert isinstance(result, Image.Image)
assert result.size == sample_image.size
def test_custom_labels(self, sample_image, sample_boxes):
labels = ["fragment", "panel"]
result = draw_boxes(sample_image, sample_boxes, labels=labels)
assert isinstance(result, Image.Image)
def test_tiny_boxes_expanded(self, sample_image):
tiny_boxes = [BBox(x1=100, y1=100, x2=101, y2=101)]
result = draw_boxes(sample_image, tiny_boxes)
assert isinstance(result, Image.Image)
def test_out_of_bounds_boxes_clipped(self, sample_image):
boxes = [BBox(x1=-50, y1=-50, x2=800, y2=600)]
result = draw_boxes(sample_image, boxes)
assert isinstance(result, Image.Image)
def test_no_confidence_display(self, sample_image, sample_boxes):
result = draw_boxes(sample_image, sample_boxes, show_confidence=False)
assert isinstance(result, Image.Image)
class TestNoDetectionOverlay:
def test_returns_image(self, sample_image):
result = create_no_detection_overlay(sample_image)
assert isinstance(result, Image.Image)
assert result.size == sample_image.size
def test_custom_message(self, sample_image):
result = create_no_detection_overlay(sample_image, "Custom message")
assert isinstance(result, Image.Image)
|