""" Tests for Training Export with expand_bbox integration. Tests the export endpoint's integration with field-specific bbox expansion. """ import pytest from unittest.mock import MagicMock, patch from uuid import uuid4 from shared.bbox import expand_bbox 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.""" # 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_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 # 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 bbox = (100, 100, 200, 150) img_width = 2480 img_height = 3508 # Auto mode (field-specific expansion) auto_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, ) # 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 def test_expand_bbox_all_field_types_work(self): """Verify expand_bbox works for all field types.""" 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, ) # 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 class TestExportAnnotationExpansion: """Tests for annotation expansion in export workflow.""" 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" x_center = 0.3 y_center = 0.2 width = 0.15 height = 0.03 source = "auto" ann = MockAnnotation() 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 y0 = ann.y_center * img_height - half_h 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 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: # 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 y0 = ann["y_center"] * img_height - half_h 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"], "x_center": (ex0 + ex1) / 2 / img_width, "y_center": (ey0 + ey1) / 2 / img_height, "width": (ex1 - ex0) / img_width, "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']}" 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", ) # 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", ) # Right edge should be clamped to image width assert result[2] <= img_width 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", ) # 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", ) # Bottom edge should be clamped to image height assert result[3] <= img_height 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 result = expand_bbox( bbox=bbox, image_width=img_width, image_height=img_height, field_type="invoice_number", ) # Should still produce a valid expanded bbox assert result[2] > result[0] assert result[3] > result[1]