| """Tests for visualization helpers.""" |
|
|
| from shelfie.models import BookItem, BoundingBox, Catalog, ShelfRegion, SpineRegion |
| from shelfie.visualize import to_gradio_annotations |
|
|
|
|
| def test_to_gradio_annotations_format(): |
| spine = SpineRegion(bbox=BoundingBox(x=10, y=20, w=30, h=200), shelf_index=0) |
| items = [BookItem(spine=spine, title="Test Book", author="Author", ocr_confidence=0.8)] |
| cat = Catalog( |
| items=items, |
| shelves=[ShelfRegion(bbox=BoundingBox(0, 0, 100, 300), index=0)], |
| source_image_path="test.jpg", |
| ) |
|
|
| annotations = to_gradio_annotations(cat) |
| assert len(annotations) == 1 |
| bbox, label = annotations[0] |
| assert bbox == (10, 20, 40, 220) |
| assert "Test Book" in label |
| assert "Author" in label |
|
|
|
|
| def test_to_gradio_annotations_highlight(): |
| spine = SpineRegion(bbox=BoundingBox(x=0, y=0, w=30, h=200), shelf_index=0) |
| items = [ |
| BookItem(spine=spine, title="A", ocr_confidence=0.5), |
| BookItem(spine=spine, title="B", ocr_confidence=0.5), |
| ] |
| cat = Catalog( |
| items=items, |
| shelves=[ShelfRegion(bbox=BoundingBox(0, 0, 100, 300), index=0)], |
| source_image_path="test.jpg", |
| ) |
|
|
| annotations = to_gradio_annotations(cat, highlight_indices=[1]) |
| assert " *" not in annotations[0][1] |
| assert " *" in annotations[1][1] |
|
|
|
|
| def test_to_gradio_annotations_unknown_item(): |
| spine = SpineRegion(bbox=BoundingBox(x=0, y=0, w=30, h=200), shelf_index=0) |
| items = [BookItem(spine=spine, title=None, ocr_confidence=0.0)] |
| cat = Catalog( |
| items=items, |
| shelves=[ShelfRegion(bbox=BoundingBox(0, 0, 100, 300), index=0)], |
| source_image_path="test.jpg", |
| ) |
|
|
| annotations = to_gradio_annotations(cat) |
| assert "[Unknown #1]" in annotations[0][1] |
|
|
