85 lines
2.2 KiB
Python
85 lines
2.2 KiB
Python
"""Tests for OpenCV preprocessing — runs without GPU."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
try:
|
|
import cv2
|
|
HAS_CV2 = True
|
|
except ImportError:
|
|
HAS_CV2 = False
|
|
|
|
requires_cv2 = pytest.mark.skipif(not HAS_CV2, reason="opencv-python-headless not installed")
|
|
|
|
# Add gpu/ to path so imports resolve (gpu modules use relative imports)
|
|
import sys
|
|
from pathlib import Path
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "gpu"))
|
|
|
|
|
|
def _make_image(w: int = 200, h: int = 60) -> np.ndarray:
|
|
"""White image with black text-like region."""
|
|
img = np.ones((h, w, 3), dtype=np.uint8) * 255
|
|
img[15:45, 20:180] = 30 # dark band simulating text
|
|
return img
|
|
|
|
|
|
@requires_cv2
|
|
def test_binarize():
|
|
from gpu.models.preprocess import binarize
|
|
|
|
img = _make_image()
|
|
result = binarize(img)
|
|
|
|
assert result.shape == img.shape
|
|
assert result.dtype == np.uint8
|
|
# Should be mostly black and white (no grays)
|
|
unique_values = np.unique(result)
|
|
assert len(unique_values) <= 3 # 0, 255, maybe one more from anti-aliasing
|
|
|
|
|
|
@requires_cv2
|
|
def test_enhance_contrast():
|
|
from gpu.models.preprocess import enhance_contrast
|
|
|
|
img = _make_image()
|
|
result = enhance_contrast(img)
|
|
|
|
assert result.shape == img.shape
|
|
assert result.dtype == np.uint8
|
|
|
|
|
|
@requires_cv2
|
|
def test_deskew_no_rotation():
|
|
from gpu.models.preprocess import deskew
|
|
|
|
img = _make_image()
|
|
result = deskew(img)
|
|
|
|
assert result.shape == img.shape
|
|
# Straight image should be unchanged (angle < 0.5 deg)
|
|
assert np.allclose(result, img, atol=5)
|
|
|
|
|
|
@requires_cv2
|
|
def test_preprocess_pipeline():
|
|
from gpu.models.preprocess import preprocess
|
|
|
|
img = _make_image()
|
|
|
|
result = preprocess(img, do_binarize=False, do_deskew=False, do_contrast=True)
|
|
assert result.shape == img.shape
|
|
|
|
result = preprocess(img, do_binarize=True, do_deskew=True, do_contrast=True)
|
|
assert result.shape[:2] == img.shape[:2] # h, w same; channels may differ then get converted back
|
|
|
|
|
|
@requires_cv2
|
|
def test_preprocess_all_disabled():
|
|
from gpu.models.preprocess import preprocess
|
|
|
|
img = _make_image()
|
|
result = preprocess(img, do_binarize=False, do_deskew=False, do_contrast=False)
|
|
|
|
assert np.array_equal(result, img)
|