227 lines
7.1 KiB
Python
227 lines
7.1 KiB
Python
"""
|
|
Tests for Training Export with uniform expand_bbox integration.
|
|
|
|
Tests the export endpoint's integration with uniform bbox expansion.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
from uuid import uuid4
|
|
|
|
from shared.bbox import expand_bbox, UNIFORM_PAD
|
|
from shared.fields import CLASS_NAMES, FIELD_CLASS_IDS
|
|
|
|
|
|
class TestExpandBboxForExport:
|
|
"""Tests for expand_bbox integration in export workflow."""
|
|
|
|
def test_expand_bbox_converts_normalized_to_pixel_and_back(self):
|
|
"""Verify expand_bbox works with pixel-to-normalized conversion."""
|
|
x_center_norm = 0.5
|
|
y_center_norm = 0.5
|
|
width_norm = 0.1
|
|
height_norm = 0.05
|
|
|
|
img_width = 2480
|
|
img_height = 3508
|
|
|
|
x_center_px = x_center_norm * img_width
|
|
y_center_px = y_center_norm * img_height
|
|
width_px = width_norm * img_width
|
|
height_px = height_norm * img_height
|
|
|
|
x0 = x_center_px - width_px / 2
|
|
y0 = y_center_px - height_px / 2
|
|
x1 = x_center_px + width_px / 2
|
|
y1 = y_center_px + height_px / 2
|
|
|
|
ex0, ey0, ex1, ey1 = expand_bbox(
|
|
bbox=(x0, y0, x1, y1),
|
|
image_width=img_width,
|
|
image_height=img_height,
|
|
)
|
|
|
|
assert ex0 < x0
|
|
assert ey0 < y0
|
|
assert ex1 > x1
|
|
assert ey1 > y1
|
|
|
|
new_x_center = (ex0 + ex1) / 2 / img_width
|
|
new_y_center = (ey0 + ey1) / 2 / img_height
|
|
new_width = (ex1 - ex0) / img_width
|
|
new_height = (ey1 - ey0) / img_height
|
|
|
|
assert 0 <= new_x_center <= 1
|
|
assert 0 <= new_y_center <= 1
|
|
assert 0 <= new_width <= 1
|
|
assert 0 <= new_height <= 1
|
|
|
|
def test_expand_bbox_uniform_for_all_sources(self):
|
|
"""Verify all annotation sources get the same uniform expansion."""
|
|
bbox = (100, 100, 200, 150)
|
|
img_width = 2480
|
|
img_height = 3508
|
|
|
|
# All sources now get the same uniform expansion
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=img_width,
|
|
image_height=img_height,
|
|
)
|
|
|
|
expected = (
|
|
100 - UNIFORM_PAD,
|
|
100 - UNIFORM_PAD,
|
|
200 + UNIFORM_PAD,
|
|
150 + UNIFORM_PAD,
|
|
)
|
|
assert result == expected
|
|
|
|
def test_expand_bbox_all_field_types_work(self):
|
|
"""Verify expand_bbox works for all field types (same result)."""
|
|
bbox = (100, 100, 200, 150)
|
|
img_width = 2480
|
|
img_height = 3508
|
|
|
|
# All fields should produce the same result with uniform padding
|
|
first_result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=img_width,
|
|
image_height=img_height,
|
|
)
|
|
|
|
assert len(first_result) == 4
|
|
x0, y0, x1, y1 = first_result
|
|
assert x0 >= 0
|
|
assert y0 >= 0
|
|
assert x1 <= img_width
|
|
assert y1 <= img_height
|
|
assert x1 > x0
|
|
assert y1 > y0
|
|
|
|
|
|
class TestExportAnnotationExpansion:
|
|
"""Tests for annotation expansion in export workflow."""
|
|
|
|
def test_annotation_bbox_conversion_workflow(self):
|
|
"""Test full annotation bbox conversion workflow."""
|
|
class MockAnnotation:
|
|
class_id = FIELD_CLASS_IDS["invoice_number"]
|
|
class_name = "invoice_number"
|
|
x_center = 0.3
|
|
y_center = 0.2
|
|
width = 0.15
|
|
height = 0.03
|
|
source = "auto"
|
|
|
|
ann = MockAnnotation()
|
|
img_width = 2480
|
|
img_height = 3508
|
|
|
|
half_w = (ann.width * img_width) / 2
|
|
half_h = (ann.height * img_height) / 2
|
|
x0 = ann.x_center * img_width - half_w
|
|
y0 = ann.y_center * img_height - half_h
|
|
x1 = ann.x_center * img_width + half_w
|
|
y1 = ann.y_center * img_height + half_h
|
|
|
|
ex0, ey0, ex1, ey1 = expand_bbox(
|
|
bbox=(x0, y0, x1, y1),
|
|
image_width=img_width,
|
|
image_height=img_height,
|
|
)
|
|
|
|
new_x_center = (ex0 + ex1) / 2 / img_width
|
|
new_y_center = (ey0 + ey1) / 2 / img_height
|
|
new_width = (ex1 - ex0) / img_width
|
|
new_height = (ey1 - ey0) / img_height
|
|
|
|
assert new_width > ann.width
|
|
assert new_height > ann.height
|
|
|
|
assert 0 <= new_x_center <= 1
|
|
assert 0 <= new_y_center <= 1
|
|
assert 0 < new_width <= 1
|
|
assert 0 < new_height <= 1
|
|
|
|
def test_export_applies_uniform_expansion_to_all_annotations(self):
|
|
"""Test that export applies uniform expansion to all annotations."""
|
|
annotations = [
|
|
{"class_name": "invoice_number", "source": "auto", "x_center": 0.3, "y_center": 0.2, "width": 0.05, "height": 0.02},
|
|
{"class_name": "ocr_number", "source": "manual", "x_center": 0.5, "y_center": 0.8, "width": 0.05, "height": 0.02},
|
|
{"class_name": "amount", "source": "imported", "x_center": 0.7, "y_center": 0.5, "width": 0.05, "height": 0.02},
|
|
]
|
|
|
|
img_width = 2480
|
|
img_height = 3508
|
|
|
|
expanded_annotations = []
|
|
for ann in annotations:
|
|
half_w = (ann["width"] * img_width) / 2
|
|
half_h = (ann["height"] * img_height) / 2
|
|
x0 = ann["x_center"] * img_width - half_w
|
|
y0 = ann["y_center"] * img_height - half_h
|
|
x1 = ann["x_center"] * img_width + half_w
|
|
y1 = ann["y_center"] * img_height + half_h
|
|
|
|
ex0, ey0, ex1, ey1 = expand_bbox(
|
|
bbox=(x0, y0, x1, y1),
|
|
image_width=img_width,
|
|
image_height=img_height,
|
|
)
|
|
|
|
expanded_annotations.append({
|
|
"class_name": ann["class_name"],
|
|
"source": ann["source"],
|
|
"x_center": (ex0 + ex1) / 2 / img_width,
|
|
"y_center": (ey0 + ey1) / 2 / img_height,
|
|
"width": (ex1 - ex0) / img_width,
|
|
"height": (ey1 - ey0) / img_height,
|
|
})
|
|
|
|
# All annotations get the same expansion
|
|
tolerance = 0.01
|
|
for orig, exp in zip(annotations, expanded_annotations):
|
|
assert exp["width"] >= orig["width"] * (1 - tolerance)
|
|
assert exp["height"] >= orig["height"] * (1 - tolerance)
|
|
|
|
|
|
class TestExpandBboxEdgeCases:
|
|
"""Tests for edge cases in export bbox expansion."""
|
|
|
|
def test_bbox_at_image_edge_left(self):
|
|
bbox = (0, 100, 50, 150)
|
|
|
|
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
|
|
|
assert result[0] >= 0
|
|
|
|
def test_bbox_at_image_edge_right(self):
|
|
bbox = (2400, 100, 2480, 150)
|
|
|
|
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
|
|
|
assert result[2] <= 2480
|
|
|
|
def test_bbox_at_image_edge_top(self):
|
|
bbox = (100, 0, 200, 50)
|
|
|
|
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
|
|
|
assert result[1] >= 0
|
|
|
|
def test_bbox_at_image_edge_bottom(self):
|
|
bbox = (100, 3400, 200, 3508)
|
|
|
|
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
|
|
|
assert result[3] <= 3508
|
|
|
|
def test_very_small_bbox(self):
|
|
bbox = (100, 100, 105, 105)
|
|
|
|
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
|
|
|
assert result[2] > result[0]
|
|
assert result[3] > result[1]
|