"""Tests for the OpenCV preprocessing pipeline.""" import cv2 import numpy as np import pytest from pipeline.preprocessing import ( decode_image, denoise, edge_detect, enhance_contrast, morphological_ops, preprocess, threshold, to_grayscale, ) def _make_test_image(width: int = 100, height: int = 80, color: bool = True) -> np.ndarray: """Create a synthetic test image with some structure.""" if color: img = np.zeros((height, width, 3), dtype=np.uint8) # White rectangle on black background cv2.rectangle(img, (20, 15), (80, 65), (255, 255, 255), -1) # Gray circle for tonal variation cv2.circle(img, (50, 40), 15, (128, 128, 128), -1) else: img = np.zeros((height, width), dtype=np.uint8) cv2.rectangle(img, (20, 15), (80, 65), 255, -1) cv2.circle(img, (50, 40), 15, 128, -1) return img def _encode_png(img: np.ndarray) -> bytes: """Encode numpy array to PNG bytes.""" ok, buf = cv2.imencode(".png", img) assert ok return buf.tobytes() # --- decode_image --- class TestDecodeImage: def test_decodes_valid_png(self): img = _make_test_image() raw = _encode_png(img) result = decode_image(raw) assert result.shape == img.shape assert result.dtype == np.uint8 def test_rejects_invalid_bytes(self): with pytest.raises(ValueError, match="Failed to decode"): decode_image(b"not an image") def test_decodes_jpeg(self): img = _make_test_image() ok, buf = cv2.imencode(".jpg", img) assert ok result = decode_image(buf.tobytes()) assert result.shape == img.shape # --- to_grayscale --- class TestToGrayscale: def test_converts_color_to_gray(self): img = _make_test_image(color=True) result = to_grayscale(img) assert len(result.shape) == 2 assert result.shape == (80, 100) def test_passthrough_already_gray(self): img = _make_test_image(color=False) result = to_grayscale(img) assert len(result.shape) == 2 np.testing.assert_array_equal(result, img) # --- denoise --- class TestDenoise: def test_returns_same_shape(self): img = _make_test_image(color=False) result = denoise(img) assert result.shape == img.shape assert result.dtype == np.uint8 def test_custom_params(self): img = _make_test_image(color=False) result = denoise(img, d=5, sigma_color=50.0, sigma_space=50.0) assert result.shape == img.shape def test_reduces_noise(self): img = _make_test_image(color=False).astype(np.float64) rng = np.random.default_rng(42) noisy = np.clip(img + rng.normal(0, 25, img.shape), 0, 255).astype(np.uint8) result = denoise(noisy, d=9, sigma_color=75, sigma_space=75) # Denoised image should be closer to original than noisy version orig = _make_test_image(color=False) noise_diff = np.mean(np.abs(noisy.astype(float) - orig.astype(float))) clean_diff = np.mean(np.abs(result.astype(float) - orig.astype(float))) assert clean_diff < noise_diff # --- enhance_contrast --- class TestEnhanceContrast: def test_returns_same_shape(self): img = _make_test_image(color=False) result = enhance_contrast(img) assert result.shape == img.shape def test_custom_clip_limit(self): img = _make_test_image(color=False) result = enhance_contrast(img, clip_limit=4.0, tile_grid_size=(4, 4)) assert result.shape == img.shape def test_increases_dynamic_range(self): # Low-contrast gray image img = np.full((80, 100), 120, dtype=np.uint8) img[20:60, 20:80] = 130 result = enhance_contrast(img, clip_limit=2.0) # CLAHE should increase the spread assert result.std() >= img.std() # --- threshold --- class TestThreshold: def test_otsu_produces_binary(self): img = _make_test_image(color=False) result = threshold(img) unique = set(np.unique(result)) assert unique <= {0, 255} def test_manual_threshold(self): img = _make_test_image(color=False) result = threshold(img, manual_thresh=100) unique = set(np.unique(result)) assert unique <= {0, 255} def test_manual_threshold_value(self): img = np.array([[50, 150], [100, 200]], dtype=np.uint8) result = threshold(img, manual_thresh=120) expected = np.array([[0, 255], [0, 255]], dtype=np.uint8) np.testing.assert_array_equal(result, expected) # --- edge_detect --- class TestEdgeDetect: def test_returns_same_shape(self): img = _make_test_image(color=False) result = edge_detect(img) assert result.shape == img.shape def test_detects_edges(self): img = np.zeros((100, 100), dtype=np.uint8) img[25:75, 25:75] = 255 result = edge_detect(img, low=50, high=150) # Should have non-zero pixels along the rectangle edges assert np.count_nonzero(result) > 0 # Interior and exterior should be zero assert result[50, 50] == 0 assert result[0, 0] == 0 def test_custom_thresholds(self): img = _make_test_image(color=False) result = edge_detect(img, low=100, high=200) assert result.dtype == np.uint8 # --- morphological_ops --- class TestMorphologicalOps: def test_returns_same_shape(self): img = threshold(_make_test_image(color=False)) result = morphological_ops(img) assert result.shape == img.shape def test_fills_small_gaps(self): # Create binary image with a small hole img = np.zeros((100, 100), dtype=np.uint8) img[20:80, 20:80] = 255 img[49:51, 49:51] = 0 # small gap result = morphological_ops(img, kernel_size=3, dilate_iterations=1, erode_iterations=1) # Gap should be filled after dilation assert result[49, 49] == 255 def test_custom_kernel_and_iterations(self): img = threshold(_make_test_image(color=False)) result = morphological_ops(img, kernel_size=5, dilate_iterations=2, erode_iterations=2) assert result.shape == img.shape # --- full pipeline --- class TestPreprocess: def test_end_to_end_default_params(self): img = _make_test_image() raw = _encode_png(img) result = preprocess(raw) assert len(result.shape) == 2 assert result.dtype == np.uint8 unique = set(np.unique(result)) # After threshold + morph, should be mostly binary assert unique <= {0, 255} def test_with_edge_detection(self): img = _make_test_image() raw = _encode_png(img) result = preprocess(raw, params={"edge_detect": True}) assert len(result.shape) == 2 def test_custom_params(self): img = _make_test_image() raw = _encode_png(img) result = preprocess(raw, params={ "denoise_d": 5, "denoise_sigma_color": 50.0, "clahe_clip_limit": 4.0, "threshold_manual": 128, "morph_kernel_size": 5, "morph_dilate_iterations": 2, "morph_erode_iterations": 2, }) assert result.dtype == np.uint8 assert len(result.shape) == 2 def test_rejects_bad_input(self): with pytest.raises(ValueError): preprocess(b"garbage data")