WIP
This commit is contained in:
@@ -1,14 +1,14 @@
|
||||
"""
|
||||
Tests for Training Export with expand_bbox integration.
|
||||
Tests for Training Export with uniform expand_bbox integration.
|
||||
|
||||
Tests the export endpoint's integration with field-specific bbox expansion.
|
||||
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
|
||||
from shared.bbox import expand_bbox, UNIFORM_PAD
|
||||
from shared.fields import CLASS_NAMES, FIELD_CLASS_IDS
|
||||
|
||||
|
||||
@@ -17,149 +17,87 @@ class TestExpandBboxForExport:
|
||||
|
||||
def test_expand_bbox_converts_normalized_to_pixel_and_back(self):
|
||||
"""Verify expand_bbox works with pixel-to-normalized conversion."""
|
||||
# Annotation stored as normalized coords
|
||||
x_center_norm = 0.5
|
||||
y_center_norm = 0.5
|
||||
width_norm = 0.1
|
||||
height_norm = 0.05
|
||||
|
||||
# Image dimensions
|
||||
img_width = 2480 # A4 at 300 DPI
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
|
||||
# Convert to pixel coords
|
||||
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
|
||||
|
||||
# Convert to corner coords
|
||||
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
|
||||
|
||||
# Apply expansion
|
||||
class_name = "invoice_number"
|
||||
ex0, ey0, ex1, ey1 = expand_bbox(
|
||||
bbox=(x0, y0, x1, y1),
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type=class_name,
|
||||
)
|
||||
|
||||
# Verify expanded bbox is larger
|
||||
assert ex0 < x0 # Left expanded
|
||||
assert ey0 < y0 # Top expanded
|
||||
assert ex1 > x1 # Right expanded
|
||||
assert ey1 > y1 # Bottom expanded
|
||||
assert ex0 < x0
|
||||
assert ey0 < y0
|
||||
assert ex1 > x1
|
||||
assert ey1 > y1
|
||||
|
||||
# Convert back to normalized
|
||||
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
|
||||
|
||||
# Verify valid normalized coords
|
||||
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_manual_mode_minimal_expansion(self):
|
||||
"""Verify manual annotations use minimal expansion."""
|
||||
# Small bbox
|
||||
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
|
||||
|
||||
# Auto mode (field-specific expansion)
|
||||
auto_result = expand_bbox(
|
||||
# All sources now get the same uniform expansion
|
||||
result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type="invoice_number",
|
||||
manual_mode=False,
|
||||
)
|
||||
|
||||
# Manual mode (minimal expansion)
|
||||
manual_result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type="invoice_number",
|
||||
manual_mode=True,
|
||||
expected = (
|
||||
100 - UNIFORM_PAD,
|
||||
100 - UNIFORM_PAD,
|
||||
200 + UNIFORM_PAD,
|
||||
150 + UNIFORM_PAD,
|
||||
)
|
||||
|
||||
# Auto expansion should be larger than manual
|
||||
auto_width = auto_result[2] - auto_result[0]
|
||||
manual_width = manual_result[2] - manual_result[0]
|
||||
assert auto_width > manual_width
|
||||
|
||||
auto_height = auto_result[3] - auto_result[1]
|
||||
manual_height = manual_result[3] - manual_result[1]
|
||||
assert auto_height > manual_height
|
||||
|
||||
def test_expand_bbox_different_sources_use_correct_mode(self):
|
||||
"""Verify different annotation sources use correct expansion mode."""
|
||||
bbox = (100, 100, 200, 150)
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
|
||||
# Define source to manual_mode mapping
|
||||
source_mode_mapping = {
|
||||
"manual": True, # Manual annotations -> minimal expansion
|
||||
"auto": False, # Auto-labeled -> field-specific expansion
|
||||
"imported": True, # Imported (from CSV) -> minimal expansion
|
||||
}
|
||||
|
||||
results = {}
|
||||
for source, manual_mode in source_mode_mapping.items():
|
||||
result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type="ocr_number",
|
||||
manual_mode=manual_mode,
|
||||
)
|
||||
results[source] = result
|
||||
|
||||
# Auto should have largest expansion
|
||||
auto_area = (results["auto"][2] - results["auto"][0]) * \
|
||||
(results["auto"][3] - results["auto"][1])
|
||||
manual_area = (results["manual"][2] - results["manual"][0]) * \
|
||||
(results["manual"][3] - results["manual"][1])
|
||||
imported_area = (results["imported"][2] - results["imported"][0]) * \
|
||||
(results["imported"][3] - results["imported"][1])
|
||||
|
||||
assert auto_area > manual_area
|
||||
assert auto_area > imported_area
|
||||
# Manual and imported should be the same (both use minimal mode)
|
||||
assert manual_area == imported_area
|
||||
assert result == expected
|
||||
|
||||
def test_expand_bbox_all_field_types_work(self):
|
||||
"""Verify expand_bbox works for all field types."""
|
||||
"""Verify expand_bbox works for all field types (same result)."""
|
||||
bbox = (100, 100, 200, 150)
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
|
||||
for class_name in CLASS_NAMES:
|
||||
result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type=class_name,
|
||||
)
|
||||
# All fields should produce the same result with uniform padding
|
||||
first_result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
)
|
||||
|
||||
# Verify result is a valid bbox
|
||||
assert len(result) == 4
|
||||
x0, y0, x1, y1 = result
|
||||
assert x0 >= 0
|
||||
assert y0 >= 0
|
||||
assert x1 <= img_width
|
||||
assert y1 <= img_height
|
||||
assert x1 > x0
|
||||
assert y1 > y0
|
||||
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:
|
||||
@@ -167,7 +105,6 @@ class TestExportAnnotationExpansion:
|
||||
|
||||
def test_annotation_bbox_conversion_workflow(self):
|
||||
"""Test full annotation bbox conversion workflow."""
|
||||
# Simulate stored annotation (normalized coords)
|
||||
class MockAnnotation:
|
||||
class_id = FIELD_CLASS_IDS["invoice_number"]
|
||||
class_name = "invoice_number"
|
||||
@@ -181,7 +118,6 @@ class TestExportAnnotationExpansion:
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
|
||||
# Step 1: Convert normalized to pixel corner coords
|
||||
half_w = (ann.width * img_width) / 2
|
||||
half_h = (ann.height * img_height) / 2
|
||||
x0 = ann.x_center * img_width - half_w
|
||||
@@ -189,38 +125,27 @@ class TestExportAnnotationExpansion:
|
||||
x1 = ann.x_center * img_width + half_w
|
||||
y1 = ann.y_center * img_height + half_h
|
||||
|
||||
# Step 2: Determine manual_mode based on source
|
||||
manual_mode = ann.source in ("manual", "imported")
|
||||
|
||||
# Step 3: Apply expand_bbox
|
||||
ex0, ey0, ex1, ey1 = expand_bbox(
|
||||
bbox=(x0, y0, x1, y1),
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type=ann.class_name,
|
||||
manual_mode=manual_mode,
|
||||
)
|
||||
|
||||
# Step 4: Convert back to normalized
|
||||
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
|
||||
|
||||
# Verify expansion happened (auto mode)
|
||||
assert new_width > ann.width
|
||||
assert new_height > ann.height
|
||||
|
||||
# Verify valid YOLO format
|
||||
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_expansion_to_each_annotation(self):
|
||||
"""Test that export applies expansion to each annotation."""
|
||||
# Simulate multiple annotations with different sources
|
||||
# Use smaller bboxes so manual mode padding has visible effect
|
||||
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},
|
||||
@@ -232,7 +157,6 @@ class TestExportAnnotationExpansion:
|
||||
|
||||
expanded_annotations = []
|
||||
for ann in annotations:
|
||||
# Convert to pixel coords
|
||||
half_w = (ann["width"] * img_width) / 2
|
||||
half_h = (ann["height"] * img_height) / 2
|
||||
x0 = ann["x_center"] * img_width - half_w
|
||||
@@ -240,19 +164,12 @@ class TestExportAnnotationExpansion:
|
||||
x1 = ann["x_center"] * img_width + half_w
|
||||
y1 = ann["y_center"] * img_height + half_h
|
||||
|
||||
# Determine manual_mode
|
||||
manual_mode = ann["source"] in ("manual", "imported")
|
||||
|
||||
# Apply expansion
|
||||
ex0, ey0, ex1, ey1 = expand_bbox(
|
||||
bbox=(x0, y0, x1, y1),
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type=ann["class_name"],
|
||||
manual_mode=manual_mode,
|
||||
)
|
||||
|
||||
# Convert back to normalized
|
||||
expanded_annotations.append({
|
||||
"class_name": ann["class_name"],
|
||||
"source": ann["source"],
|
||||
@@ -262,106 +179,48 @@ class TestExportAnnotationExpansion:
|
||||
"height": (ey1 - ey0) / img_height,
|
||||
})
|
||||
|
||||
# Verify auto-labeled annotation expanded more than manual/imported
|
||||
auto_ann = next(a for a in expanded_annotations if a["source"] == "auto")
|
||||
manual_ann = next(a for a in expanded_annotations if a["source"] == "manual")
|
||||
|
||||
# Auto mode should expand more than manual mode
|
||||
# (auto has larger scale factors and max_pad)
|
||||
assert auto_ann["width"] > manual_ann["width"]
|
||||
assert auto_ann["height"] > manual_ann["height"]
|
||||
|
||||
# All annotations should be expanded (at least slightly for manual mode)
|
||||
# Allow small precision loss (< 1%) due to integer conversion in expand_bbox
|
||||
for i, (orig, exp) in enumerate(zip(annotations, expanded_annotations)):
|
||||
# Width and height should be >= original (expansion or equal, with small tolerance)
|
||||
tolerance = 0.01 # 1% tolerance for integer rounding
|
||||
assert exp["width"] >= orig["width"] * (1 - tolerance), \
|
||||
f"Annotation {i} width unexpectedly smaller: {exp['width']} < {orig['width']}"
|
||||
assert exp["height"] >= orig["height"] * (1 - tolerance), \
|
||||
f"Annotation {i} height unexpectedly smaller: {exp['height']} < {orig['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):
|
||||
"""Test bbox at left edge of image."""
|
||||
bbox = (0, 100, 50, 150)
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
|
||||
result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type="invoice_number",
|
||||
)
|
||||
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
||||
|
||||
# Left edge should be clamped to 0
|
||||
assert result[0] >= 0
|
||||
|
||||
def test_bbox_at_image_edge_right(self):
|
||||
"""Test bbox at right edge of image."""
|
||||
bbox = (2400, 100, 2480, 150)
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
|
||||
result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type="invoice_number",
|
||||
)
|
||||
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
||||
|
||||
# Right edge should be clamped to image width
|
||||
assert result[2] <= img_width
|
||||
assert result[2] <= 2480
|
||||
|
||||
def test_bbox_at_image_edge_top(self):
|
||||
"""Test bbox at top edge of image."""
|
||||
bbox = (100, 0, 200, 50)
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
|
||||
result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type="invoice_number",
|
||||
)
|
||||
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
||||
|
||||
# Top edge should be clamped to 0
|
||||
assert result[1] >= 0
|
||||
|
||||
def test_bbox_at_image_edge_bottom(self):
|
||||
"""Test bbox at bottom edge of image."""
|
||||
bbox = (100, 3400, 200, 3508)
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
|
||||
result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type="invoice_number",
|
||||
)
|
||||
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
||||
|
||||
# Bottom edge should be clamped to image height
|
||||
assert result[3] <= img_height
|
||||
assert result[3] <= 3508
|
||||
|
||||
def test_very_small_bbox(self):
|
||||
"""Test very small bbox gets expanded."""
|
||||
bbox = (100, 100, 105, 105) # 5x5 pixel bbox
|
||||
img_width = 2480
|
||||
img_height = 3508
|
||||
bbox = (100, 100, 105, 105)
|
||||
|
||||
result = expand_bbox(
|
||||
bbox=bbox,
|
||||
image_width=img_width,
|
||||
image_height=img_height,
|
||||
field_type="invoice_number",
|
||||
)
|
||||
result = expand_bbox(bbox=bbox, image_width=2480, image_height=3508)
|
||||
|
||||
# Should still produce a valid expanded bbox
|
||||
assert result[2] > result[0]
|
||||
assert result[3] > result[1]
|
||||
|
||||
Reference in New Issue
Block a user