0990239e9c
Implement center-point based bbox scaling with directional compensation to capture field labels that typically appear above or to the left of field values. This improves YOLO training data quality by including contextual information around field values. Key changes: - Add shared.bbox module with ScaleStrategy dataclass and expand_bbox function - Define field-specific strategies (ocr_number, bankgiro, invoice_date, etc.) - Support manual_mode for minimal padding (no scaling) - Integrate expand_bbox into AnnotationGenerator - Add FIELD_TO_CLASS mapping for field_name to class_name lookup - Comprehensive tests with 100% coverage (45 tests) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
557 lines
17 KiB
Python
557 lines
17 KiB
Python
"""
|
|
Tests for expand_bbox function.
|
|
|
|
Tests verify that bbox expansion works correctly with center-point scaling,
|
|
directional compensation, max padding clamping, and image boundary handling.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from shared.bbox import (
|
|
expand_bbox,
|
|
ScaleStrategy,
|
|
FIELD_SCALE_STRATEGIES,
|
|
DEFAULT_STRATEGY,
|
|
)
|
|
|
|
|
|
class TestExpandBboxCenterScaling:
|
|
"""Tests for center-point based scaling."""
|
|
|
|
def test_center_scaling_expands_symmetrically(self):
|
|
"""Verify bbox expands symmetrically around center when no extra ratios."""
|
|
# 100x50 bbox at (100, 200)
|
|
bbox = (100, 200, 200, 250)
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.2, # 20% wider
|
|
scale_y=1.4, # 40% taller
|
|
max_pad_x=1000, # Large to avoid clamping
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# Original: width=100, height=50
|
|
# New: width=120, height=70
|
|
# Center: (150, 225)
|
|
# Expected: x0=150-60=90, x1=150+60=210, y0=225-35=190, y1=225+35=260
|
|
assert result[0] == 90 # x0
|
|
assert result[1] == 190 # y0
|
|
assert result[2] == 210 # x1
|
|
assert result[3] == 260 # y1
|
|
|
|
def test_no_scaling_returns_original(self):
|
|
"""Verify scale=1.0 with no extras returns original bbox."""
|
|
bbox = (100, 200, 200, 250)
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
assert result == (100, 200, 200, 250)
|
|
|
|
|
|
class TestExpandBboxDirectionalCompensation:
|
|
"""Tests for directional compensation (extra ratios)."""
|
|
|
|
def test_extra_top_expands_upward(self):
|
|
"""Verify extra_top_ratio adds expansion toward top."""
|
|
bbox = (100, 200, 200, 250) # width=100, height=50
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_top_ratio=0.5, # Add 50% of height to top
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# extra_top = 50 * 0.5 = 25
|
|
assert result[0] == 100 # x0 unchanged
|
|
assert result[1] == 175 # y0 = 200 - 25
|
|
assert result[2] == 200 # x1 unchanged
|
|
assert result[3] == 250 # y1 unchanged
|
|
|
|
def test_extra_left_expands_leftward(self):
|
|
"""Verify extra_left_ratio adds expansion toward left."""
|
|
bbox = (100, 200, 200, 250) # width=100
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_left_ratio=0.8, # Add 80% of width to left
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# extra_left = 100 * 0.8 = 80
|
|
assert result[0] == 20 # x0 = 100 - 80
|
|
assert result[1] == 200 # y0 unchanged
|
|
assert result[2] == 200 # x1 unchanged
|
|
assert result[3] == 250 # y1 unchanged
|
|
|
|
def test_extra_right_expands_rightward(self):
|
|
"""Verify extra_right_ratio adds expansion toward right."""
|
|
bbox = (100, 200, 200, 250) # width=100
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_right_ratio=0.3, # Add 30% of width to right
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# extra_right = 100 * 0.3 = 30
|
|
assert result[0] == 100 # x0 unchanged
|
|
assert result[1] == 200 # y0 unchanged
|
|
assert result[2] == 230 # x1 = 200 + 30
|
|
assert result[3] == 250 # y1 unchanged
|
|
|
|
def test_extra_bottom_expands_downward(self):
|
|
"""Verify extra_bottom_ratio adds expansion toward bottom."""
|
|
bbox = (100, 200, 200, 250) # height=50
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_bottom_ratio=0.4, # Add 40% of height to bottom
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# extra_bottom = 50 * 0.4 = 20
|
|
assert result[0] == 100 # x0 unchanged
|
|
assert result[1] == 200 # y0 unchanged
|
|
assert result[2] == 200 # x1 unchanged
|
|
assert result[3] == 270 # y1 = 250 + 20
|
|
|
|
def test_combined_scaling_and_directional(self):
|
|
"""Verify scale + directional compensation work together."""
|
|
bbox = (100, 200, 200, 250) # width=100, height=50
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.2, # 20% wider -> 120 width
|
|
scale_y=1.0, # no height change
|
|
extra_left_ratio=0.5, # Add 50% of width to left
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# Center: x=150
|
|
# After scale: width=120 -> x0=150-60=90, x1=150+60=210
|
|
# After extra_left: x0 = 90 - (100 * 0.5) = 40
|
|
assert result[0] == 40 # x0
|
|
assert result[2] == 210 # x1
|
|
|
|
|
|
class TestExpandBboxMaxPadClamping:
|
|
"""Tests for max padding clamping."""
|
|
|
|
def test_max_pad_x_limits_horizontal_expansion(self):
|
|
"""Verify max_pad_x limits expansion on left and right."""
|
|
bbox = (100, 200, 200, 250) # width=100
|
|
strategy = ScaleStrategy(
|
|
scale_x=2.0, # Double width (would add 50 each side)
|
|
scale_y=1.0,
|
|
max_pad_x=30, # Limit to 30 pixels each side
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# Scale would make: x0=100, x1=200 -> x0=50, x1=250 (50px each side)
|
|
# But max_pad_x=30 limits to: x0=70, x1=230
|
|
assert result[0] == 70 # x0 = 100 - 30
|
|
assert result[2] == 230 # x1 = 200 + 30
|
|
|
|
def test_max_pad_y_limits_vertical_expansion(self):
|
|
"""Verify max_pad_y limits expansion on top and bottom."""
|
|
bbox = (100, 200, 200, 250) # height=50
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=3.0, # Triple height (would add 50 each side)
|
|
max_pad_x=1000,
|
|
max_pad_y=20, # Limit to 20 pixels each side
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# Scale would make: y0=175, y1=275 (50px each side)
|
|
# But max_pad_y=20 limits to: y0=180, y1=270
|
|
assert result[1] == 180 # y0 = 200 - 20
|
|
assert result[3] == 270 # y1 = 250 + 20
|
|
|
|
def test_max_pad_preserves_asymmetry(self):
|
|
"""Verify max_pad clamping preserves asymmetric expansion."""
|
|
bbox = (100, 200, 200, 250) # width=100
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_left_ratio=1.0, # 100px left expansion
|
|
extra_right_ratio=0.0, # No right expansion
|
|
max_pad_x=50, # Limit to 50 pixels
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
# Left would expand 100, clamped to 50
|
|
# Right stays at 0
|
|
assert result[0] == 50 # x0 = 100 - 50
|
|
assert result[2] == 200 # x1 unchanged
|
|
|
|
|
|
class TestExpandBboxImageBoundaryClamping:
|
|
"""Tests for image boundary clamping."""
|
|
|
|
def test_clamps_to_left_boundary(self):
|
|
"""Verify x0 is clamped to 0."""
|
|
bbox = (10, 200, 110, 250) # Close to left edge
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_left_ratio=0.5, # Would push x0 below 0
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
assert result[0] == 0 # Clamped to 0
|
|
|
|
def test_clamps_to_top_boundary(self):
|
|
"""Verify y0 is clamped to 0."""
|
|
bbox = (100, 10, 200, 60) # Close to top edge
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_top_ratio=0.5, # Would push y0 below 0
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
assert result[1] == 0 # Clamped to 0
|
|
|
|
def test_clamps_to_right_boundary(self):
|
|
"""Verify x1 is clamped to image_width."""
|
|
bbox = (900, 200, 990, 250) # Close to right edge
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_right_ratio=0.5, # Would push x1 beyond image_width
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
assert result[2] == 1000 # Clamped to image_width
|
|
|
|
def test_clamps_to_bottom_boundary(self):
|
|
"""Verify y1 is clamped to image_height."""
|
|
bbox = (100, 940, 200, 990) # Close to bottom edge
|
|
strategy = ScaleStrategy(
|
|
scale_x=1.0,
|
|
scale_y=1.0,
|
|
extra_bottom_ratio=0.5, # Would push y1 beyond image_height
|
|
max_pad_x=1000,
|
|
max_pad_y=1000,
|
|
)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test_field",
|
|
strategies={"test_field": strategy},
|
|
)
|
|
|
|
assert result[3] == 1000 # Clamped to image_height
|
|
|
|
|
|
class TestExpandBboxUnknownField:
|
|
"""Tests for unknown field handling."""
|
|
|
|
def test_unknown_field_uses_default_strategy(self):
|
|
"""Verify unknown field types use DEFAULT_STRATEGY."""
|
|
bbox = (100, 200, 200, 250)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="unknown_field_xyz",
|
|
)
|
|
|
|
# DEFAULT_STRATEGY: scale_x=1.15, scale_y=1.15
|
|
# Original: width=100, height=50
|
|
# New: width=115, height=57.5
|
|
# Center: (150, 225)
|
|
# x0 = 150 - 57.5 = 92.5 -> 92
|
|
# x1 = 150 + 57.5 = 207.5 -> 207
|
|
# y0 = 225 - 28.75 = 196.25 -> 196
|
|
# y1 = 225 + 28.75 = 253.75 -> 253
|
|
# But max_pad_x=50 may clamp...
|
|
# Left pad = 100 - 92.5 = 7.5 (< 50, ok)
|
|
# Right pad = 207.5 - 200 = 7.5 (< 50, ok)
|
|
assert result[0] == 92
|
|
assert result[2] == 207
|
|
|
|
|
|
class TestExpandBboxWithRealStrategies:
|
|
"""Tests using actual FIELD_SCALE_STRATEGIES."""
|
|
|
|
def test_ocr_number_expands_significantly_upward(self):
|
|
"""Verify ocr_number field gets significant upward expansion."""
|
|
bbox = (100, 200, 200, 230) # Small height=30
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="ocr_number",
|
|
)
|
|
|
|
# extra_top_ratio=0.60 -> 30 * 0.6 = 18 extra top
|
|
# y0 should decrease significantly
|
|
assert result[1] < 200 - 10 # At least 10px upward expansion
|
|
|
|
def test_bankgiro_expands_significantly_leftward(self):
|
|
"""Verify bankgiro field gets significant leftward expansion."""
|
|
bbox = (200, 200, 300, 230) # width=100
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="bankgiro",
|
|
)
|
|
|
|
# extra_left_ratio=0.80 -> 100 * 0.8 = 80 extra left
|
|
# x0 should decrease significantly
|
|
assert result[0] < 200 - 30 # At least 30px leftward expansion
|
|
|
|
def test_amount_expands_rightward(self):
|
|
"""Verify amount field gets rightward expansion for currency."""
|
|
bbox = (100, 200, 200, 230) # width=100
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="amount",
|
|
)
|
|
|
|
# extra_right_ratio=0.30 -> 100 * 0.3 = 30 extra right
|
|
# x1 should increase
|
|
assert result[2] > 200 + 10 # At least 10px rightward expansion
|
|
|
|
|
|
class TestExpandBboxReturnType:
|
|
"""Tests for return type and value format."""
|
|
|
|
def test_returns_tuple_of_four_ints(self):
|
|
"""Verify return type is tuple of 4 integers."""
|
|
bbox = (100.5, 200.3, 200.7, 250.9)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="invoice_number",
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 4
|
|
assert all(isinstance(v, int) for v in result)
|
|
|
|
def test_returns_valid_bbox_format(self):
|
|
"""Verify returned bbox has x0 < x1 and y0 < y1."""
|
|
bbox = (100, 200, 200, 250)
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="invoice_number",
|
|
)
|
|
|
|
x0, y0, x1, y1 = result
|
|
assert x0 < x1, "x0 should be less than x1"
|
|
assert y0 < y1, "y0 should be less than y1"
|
|
|
|
|
|
class TestManualLabelMode:
|
|
"""Tests for manual_mode parameter."""
|
|
|
|
def test_manual_mode_uses_minimal_padding(self):
|
|
"""Verify manual_mode uses MANUAL_LABEL_STRATEGY with minimal padding."""
|
|
bbox = (100, 200, 200, 250) # width=100, height=50
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="bankgiro", # Would normally expand left significantly
|
|
manual_mode=True,
|
|
)
|
|
|
|
# MANUAL_LABEL_STRATEGY: scale=1.0, max_pad=10
|
|
# Should only add 10px padding each side (but scale=1.0 means no scaling)
|
|
# Actually with scale=1.0, no extra ratios, we get 0 expansion from scaling
|
|
# Only max_pad=10 applies as a limit, but there's no expansion to limit
|
|
# So result should be same as original
|
|
assert result == (100, 200, 200, 250)
|
|
|
|
def test_manual_mode_ignores_field_type(self):
|
|
"""Verify manual_mode ignores field-specific strategies."""
|
|
bbox = (100, 200, 200, 250)
|
|
|
|
# Different fields should give same result in manual_mode
|
|
result_bankgiro = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="bankgiro",
|
|
manual_mode=True,
|
|
)
|
|
|
|
result_ocr = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="ocr_number",
|
|
manual_mode=True,
|
|
)
|
|
|
|
assert result_bankgiro == result_ocr
|
|
|
|
def test_manual_mode_vs_auto_mode_different(self):
|
|
"""Verify manual_mode produces different results than auto mode."""
|
|
bbox = (100, 200, 200, 250)
|
|
|
|
auto_result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="bankgiro", # Has extra_left_ratio=0.80
|
|
manual_mode=False,
|
|
)
|
|
|
|
manual_result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="bankgiro",
|
|
manual_mode=True,
|
|
)
|
|
|
|
# Auto mode should expand more (especially to the left for bankgiro)
|
|
assert auto_result[0] < manual_result[0] # Auto x0 is more left
|
|
|
|
def test_manual_mode_clamps_to_image_bounds(self):
|
|
"""Verify manual_mode still respects image boundaries."""
|
|
bbox = (5, 5, 50, 50) # Close to top-left corner
|
|
|
|
result = expand_bbox(
|
|
bbox=bbox,
|
|
image_width=1000,
|
|
image_height=1000,
|
|
field_type="test",
|
|
manual_mode=True,
|
|
)
|
|
|
|
# Should clamp to 0
|
|
assert result[0] >= 0
|
|
assert result[1] >= 0
|