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>
395 lines
15 KiB
Python
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 == []
|