Files
breakpilot-lehrer/klausur-service/backend/tests/test_doclayout_detect.py
Benjamin Admin be7f5f1872 feat: Sprint 2 — TrOCR ONNX, PP-DocLayout, Model Management
D2: TrOCR ONNX export script (printed + handwritten, int8 quantization)
D3: PP-DocLayout ONNX export script (download or Docker-based conversion)
B3: Model Management admin page (PyTorch vs ONNX status, benchmarks, config)
A4: TrOCR ONNX service with runtime routing (auto/pytorch/onnx via TROCR_BACKEND)
A5: PP-DocLayout ONNX detection with OpenCV fallback (via GRAPHIC_DETECT_BACKEND)
B4: Structure Detection UI toggle (OpenCV vs PP-DocLayout) with class color coding
C3: TrOCR-ONNX.md documentation
C4: OCR-Pipeline.md ONNX section added
C5: mkdocs.yml nav updated, optimum added to requirements.txt

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-23 09:53:02 +01:00

395 lines
15 KiB
Python

"""
Tests for PP-DocLayout ONNX Document Layout Detection.
Uses mocking to avoid requiring the actual ONNX model file.
"""
import numpy as np
import pytest
from unittest.mock import patch, MagicMock
# We patch the module-level globals before importing to ensure clean state
# in tests that check "no model" behaviour.
import importlib
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _fresh_import():
"""Re-import cv_doclayout_detect with reset globals."""
import cv_doclayout_detect as mod
# Reset module-level caching so each test starts clean
mod._onnx_session = None
mod._model_path = None
mod._load_attempted = False
mod._load_error = None
return mod
# ---------------------------------------------------------------------------
# 1. is_doclayout_available — no model present
# ---------------------------------------------------------------------------
class TestIsDoclayoutAvailableNoModel:
def test_returns_false_when_no_onnx_file(self):
mod = _fresh_import()
with patch.object(mod, "_find_model_path", return_value=None):
assert mod.is_doclayout_available() is False
def test_returns_false_when_onnxruntime_missing(self):
mod = _fresh_import()
with patch.object(mod, "_find_model_path", return_value="/fake/model.onnx"):
with patch.dict("sys.modules", {"onnxruntime": None}):
# Force ImportError by making import fail
import builtins
real_import = builtins.__import__
def fake_import(name, *args, **kwargs):
if name == "onnxruntime":
raise ImportError("no onnxruntime")
return real_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=fake_import):
assert mod.is_doclayout_available() is False
# ---------------------------------------------------------------------------
# 2. LayoutRegion dataclass
# ---------------------------------------------------------------------------
class TestLayoutRegionDataclass:
def test_basic_creation(self):
from cv_doclayout_detect import LayoutRegion
region = LayoutRegion(
x=10, y=20, width=100, height=200,
label="figure", confidence=0.95, label_index=1,
)
assert region.x == 10
assert region.y == 20
assert region.width == 100
assert region.height == 200
assert region.label == "figure"
assert region.confidence == 0.95
assert region.label_index == 1
def test_all_fields_present(self):
from cv_doclayout_detect import LayoutRegion
import dataclasses
field_names = {f.name for f in dataclasses.fields(LayoutRegion)}
expected = {"x", "y", "width", "height", "label", "confidence", "label_index"}
assert field_names == expected
def test_different_labels(self):
from cv_doclayout_detect import LayoutRegion, DOCLAYOUT_CLASSES
for idx, label in enumerate(DOCLAYOUT_CLASSES):
region = LayoutRegion(
x=0, y=0, width=50, height=50,
label=label, confidence=0.8, label_index=idx,
)
assert region.label == label
assert region.label_index == idx
# ---------------------------------------------------------------------------
# 3. detect_layout_regions — no model available
# ---------------------------------------------------------------------------
class TestDetectLayoutRegionsNoModel:
def test_returns_empty_list_when_model_unavailable(self):
mod = _fresh_import()
with patch.object(mod, "_find_model_path", return_value=None):
img = np.zeros((480, 640, 3), dtype=np.uint8)
result = mod.detect_layout_regions(img)
assert result == []
def test_returns_empty_list_for_none_image(self):
mod = _fresh_import()
with patch.object(mod, "_find_model_path", return_value=None):
result = mod.detect_layout_regions(None)
assert result == []
def test_returns_empty_list_for_empty_image(self):
mod = _fresh_import()
with patch.object(mod, "_find_model_path", return_value=None):
img = np.array([], dtype=np.uint8)
result = mod.detect_layout_regions(img)
assert result == []
# ---------------------------------------------------------------------------
# 4. Preprocessing — tensor shape verification
# ---------------------------------------------------------------------------
class TestPreprocessingShapes:
def test_square_image(self):
from cv_doclayout_detect import preprocess_image
img = np.random.randint(0, 255, (800, 800, 3), dtype=np.uint8)
tensor, scale, pad_x, pad_y = preprocess_image(img)
assert tensor.shape == (1, 3, 800, 800)
assert tensor.dtype == np.float32
assert 0.0 <= tensor.min()
assert tensor.max() <= 1.0
def test_landscape_image(self):
from cv_doclayout_detect import preprocess_image
img = np.random.randint(0, 255, (600, 1200, 3), dtype=np.uint8)
tensor, scale, pad_x, pad_y = preprocess_image(img)
assert tensor.shape == (1, 3, 800, 800)
# Landscape: scale by width, should have vertical padding
expected_scale = 800 / 1200
assert abs(scale - expected_scale) < 1e-5
assert pad_y > 0 # vertical padding expected
def test_portrait_image(self):
from cv_doclayout_detect import preprocess_image
img = np.random.randint(0, 255, (1200, 600, 3), dtype=np.uint8)
tensor, scale, pad_x, pad_y = preprocess_image(img)
assert tensor.shape == (1, 3, 800, 800)
# Portrait: scale by height, should have horizontal padding
expected_scale = 800 / 1200
assert abs(scale - expected_scale) < 1e-5
assert pad_x > 0 # horizontal padding expected
def test_small_image(self):
from cv_doclayout_detect import preprocess_image
img = np.random.randint(0, 255, (100, 200, 3), dtype=np.uint8)
tensor, scale, pad_x, pad_y = preprocess_image(img)
assert tensor.shape == (1, 3, 800, 800)
def test_typical_scan_a4(self):
"""A4 scan at 300dpi: roughly 2480x3508 pixels."""
from cv_doclayout_detect import preprocess_image
img = np.random.randint(0, 255, (3508, 2480, 3), dtype=np.uint8)
tensor, scale, pad_x, pad_y = preprocess_image(img)
assert tensor.shape == (1, 3, 800, 800)
def test_values_normalized(self):
from cv_doclayout_detect import preprocess_image
# All white image
img = np.full((400, 400, 3), 255, dtype=np.uint8)
tensor, _, _, _ = preprocess_image(img)
# The padded region is 114/255 ≈ 0.447, the image region is 1.0
assert tensor.max() <= 1.0
assert tensor.min() >= 0.0
# ---------------------------------------------------------------------------
# 5. NMS logic
# ---------------------------------------------------------------------------
class TestNmsLogic:
def test_empty_input(self):
from cv_doclayout_detect import nms
boxes = np.array([]).reshape(0, 4)
scores = np.array([])
assert nms(boxes, scores) == []
def test_single_box(self):
from cv_doclayout_detect import nms
boxes = np.array([[10, 10, 100, 100]], dtype=np.float32)
scores = np.array([0.9])
kept = nms(boxes, scores, iou_threshold=0.5)
assert kept == [0]
def test_non_overlapping_boxes(self):
from cv_doclayout_detect import nms
boxes = np.array([
[0, 0, 50, 50],
[200, 200, 300, 300],
[400, 400, 500, 500],
], dtype=np.float32)
scores = np.array([0.9, 0.8, 0.7])
kept = nms(boxes, scores, iou_threshold=0.5)
assert len(kept) == 3
assert set(kept) == {0, 1, 2}
def test_overlapping_boxes_suppressed(self):
from cv_doclayout_detect import nms
# Two boxes that heavily overlap
boxes = np.array([
[10, 10, 110, 110], # 100x100
[15, 15, 115, 115], # 100x100, heavily overlapping with first
], dtype=np.float32)
scores = np.array([0.95, 0.80])
kept = nms(boxes, scores, iou_threshold=0.5)
# Only the higher-confidence box should survive
assert kept == [0]
def test_partially_overlapping_boxes_kept(self):
from cv_doclayout_detect import nms
# Two boxes that overlap ~25% (below 0.5 threshold)
boxes = np.array([
[0, 0, 100, 100], # 100x100
[75, 0, 175, 100], # 100x100, overlap 25x100 = 2500
], dtype=np.float32)
scores = np.array([0.9, 0.8])
# IoU = 2500 / (10000 + 10000 - 2500) = 2500/17500 ≈ 0.143
kept = nms(boxes, scores, iou_threshold=0.5)
assert len(kept) == 2
def test_nms_respects_score_ordering(self):
from cv_doclayout_detect import nms
# Three overlapping boxes — highest confidence should be kept first
boxes = np.array([
[10, 10, 110, 110],
[12, 12, 112, 112],
[14, 14, 114, 114],
], dtype=np.float32)
scores = np.array([0.5, 0.9, 0.7])
kept = nms(boxes, scores, iou_threshold=0.5)
# Index 1 has highest score → kept first, suppresses 0 and 2
assert kept[0] == 1
def test_iou_computation(self):
from cv_doclayout_detect import _compute_iou
box_a = np.array([0, 0, 100, 100], dtype=np.float32)
box_b = np.array([0, 0, 100, 100], dtype=np.float32)
assert abs(_compute_iou(box_a, box_b) - 1.0) < 1e-5
box_c = np.array([200, 200, 300, 300], dtype=np.float32)
assert _compute_iou(box_a, box_c) == 0.0
# ---------------------------------------------------------------------------
# 6. DOCLAYOUT_CLASSES verification
# ---------------------------------------------------------------------------
class TestDoclayoutClasses:
def test_correct_class_list(self):
from cv_doclayout_detect import DOCLAYOUT_CLASSES
expected = [
"table", "figure", "title", "text", "list",
"header", "footer", "equation", "reference", "abstract",
]
assert DOCLAYOUT_CLASSES == expected
def test_class_count(self):
from cv_doclayout_detect import DOCLAYOUT_CLASSES
assert len(DOCLAYOUT_CLASSES) == 10
def test_no_duplicates(self):
from cv_doclayout_detect import DOCLAYOUT_CLASSES
assert len(DOCLAYOUT_CLASSES) == len(set(DOCLAYOUT_CLASSES))
def test_all_lowercase(self):
from cv_doclayout_detect import DOCLAYOUT_CLASSES
for cls in DOCLAYOUT_CLASSES:
assert cls == cls.lower(), f"Class '{cls}' should be lowercase"
# ---------------------------------------------------------------------------
# 7. get_doclayout_status
# ---------------------------------------------------------------------------
class TestGetDoclayoutStatus:
def test_status_when_unavailable(self):
mod = _fresh_import()
with patch.object(mod, "_find_model_path", return_value=None):
status = mod.get_doclayout_status()
assert status["available"] is False
assert status["model_path"] is None
assert status["load_error"] is not None
assert status["classes"] == mod.DOCLAYOUT_CLASSES
assert status["class_count"] == 10
# ---------------------------------------------------------------------------
# 8. Post-processing with mocked ONNX outputs
# ---------------------------------------------------------------------------
class TestPostprocessing:
def test_single_tensor_format_6cols(self):
"""Test parsing of (1, N, 6) output format: x1,y1,x2,y2,score,class."""
from cv_doclayout_detect import _postprocess
# One detection: figure at (100,100)-(300,300) in 800x800 space
raw = np.array([[[100, 100, 300, 300, 0.92, 1]]], dtype=np.float32)
regions = _postprocess(
outputs=[raw],
scale=1.0, pad_x=0, pad_y=0,
orig_w=800, orig_h=800,
confidence_threshold=0.5,
max_regions=50,
)
assert len(regions) == 1
assert regions[0].label == "figure"
assert regions[0].confidence >= 0.9
def test_three_tensor_format(self):
"""Test parsing of 3-tensor output: boxes, scores, class_ids."""
from cv_doclayout_detect import _postprocess
boxes = np.array([[50, 50, 200, 150]], dtype=np.float32)
scores = np.array([0.88], dtype=np.float32)
class_ids = np.array([0], dtype=np.float32) # table
regions = _postprocess(
outputs=[boxes, scores, class_ids],
scale=1.0, pad_x=0, pad_y=0,
orig_w=800, orig_h=800,
confidence_threshold=0.5,
max_regions=50,
)
assert len(regions) == 1
assert regions[0].label == "table"
def test_confidence_filtering(self):
"""Detections below threshold should be excluded."""
from cv_doclayout_detect import _postprocess
raw = np.array([
[100, 100, 200, 200, 0.9, 1], # above threshold
[300, 300, 400, 400, 0.3, 2], # below threshold
], dtype=np.float32).reshape(1, 2, 6)
regions = _postprocess(
outputs=[raw],
scale=1.0, pad_x=0, pad_y=0,
orig_w=800, orig_h=800,
confidence_threshold=0.5,
max_regions=50,
)
assert len(regions) == 1
assert regions[0].label == "figure"
def test_coordinate_scaling(self):
"""Verify coordinates are correctly scaled back to original image."""
from cv_doclayout_detect import _postprocess
# Image was 1600x1200, scaled to fit 800x800 → scale=0.5, pad_y offset
scale = 800 / 1600 # 0.5
pad_x = 0
pad_y = (800 - int(1200 * scale)) // 2 # (800-600)//2 = 100
# Detection in 800x800 space at (100, 200) to (300, 400)
raw = np.array([[[100, 200, 300, 400, 0.95, 0]]], dtype=np.float32)
regions = _postprocess(
outputs=[raw],
scale=scale, pad_x=pad_x, pad_y=pad_y,
orig_w=1600, orig_h=1200,
confidence_threshold=0.5,
max_regions=50,
)
assert len(regions) == 1
r = regions[0]
# x1 = (100 - 0) / 0.5 = 200
assert r.x == 200
# y1 = (200 - 100) / 0.5 = 200
assert r.y == 200
def test_empty_output(self):
from cv_doclayout_detect import _postprocess
raw = np.array([]).reshape(1, 0, 6).astype(np.float32)
regions = _postprocess(
outputs=[raw],
scale=1.0, pad_x=0, pad_y=0,
orig_w=800, orig_h=800,
confidence_threshold=0.5,
max_regions=50,
)
assert regions == []