"""Tests for conversion modes, invert, expanded params, and mask regions.""" import numpy as np import cv2 import pytest from pipeline.preprocessing import preprocess, decode_image, apply_mask, VALID_MODES from pipeline.vectorize import potrace_trace, vtracer_trace def _make_test_image(width=100, height=100, color=True): """Create a test image with a colored rectangle on white background.""" if color: img = np.full((height, width, 3), 255, dtype=np.uint8) # Red rectangle top-left img[10:50, 10:50] = [0, 0, 200] # Blue rectangle bottom-right img[50:90, 50:90] = [200, 0, 0] # Green strip img[20:40, 60:90] = [0, 180, 0] else: img = np.full((height, width, 3), 255, dtype=np.uint8) img[10:50, 10:50] = [0, 0, 0] _, buf = cv2.imencode(".png", img) return buf.tobytes() def _make_grayscale_image(width=100, height=100): """Create a test image with distinct gray levels.""" img = np.full((height, width, 3), 255, dtype=np.uint8) # Dark gray img[10:40, 10:40] = [60, 60, 60] # Medium gray img[40:70, 40:70] = [128, 128, 128] # Light gray img[60:90, 60:90] = [200, 200, 200] _, buf = cv2.imencode(".png", img) return buf.tobytes() # ── Conversion mode tests ── class TestConversionModes: def test_bw_mode_returns_2d_binary(self): raw = _make_test_image(color=False) result = preprocess(raw, {"conversion_mode": "bw"}) assert result.ndim == 2 unique = set(np.unique(result)) assert unique <= {0, 255}, f"Expected binary, got values: {unique}" def test_grayscale_mode_returns_2d_non_binary(self): raw = _make_grayscale_image() result = preprocess(raw, {"conversion_mode": "grayscale"}) assert result.ndim == 2 # Should have more than 2 unique values (not just 0 and 255) unique_count = len(np.unique(result)) assert unique_count > 2, f"Expected multiple gray levels, got {unique_count}" def test_color_mode_returns_3d_bgr(self): raw = _make_test_image(color=True) result = preprocess(raw, {"conversion_mode": "color"}) assert result.ndim == 3 assert result.shape[2] == 3 # BGR channels def test_default_mode_is_bw(self): raw = _make_test_image(color=False) result = preprocess(raw, {}) assert result.ndim == 2 unique = set(np.unique(result)) assert unique <= {0, 255} def test_invalid_mode_raises(self): raw = _make_test_image() with pytest.raises(ValueError, match="Invalid conversion_mode"): preprocess(raw, {"conversion_mode": "invalid"}) def test_valid_modes_constant(self): assert VALID_MODES == {"bw", "grayscale", "color"} # ── Invert tests ── class TestInvert: def test_invert_flips_bw_output(self): raw = _make_test_image(color=False) normal = preprocess(raw, {"conversion_mode": "bw", "invert": False}) inverted = preprocess(raw, {"conversion_mode": "bw", "invert": True}) # Foreground and background should be swapped assert not np.array_equal(normal, inverted) # Pixels that were 0 should now be 255 and vice versa np.testing.assert_array_equal(normal, cv2.bitwise_not(inverted)) def test_invert_false_by_default(self): raw = _make_test_image(color=False) default_result = preprocess(raw, {"conversion_mode": "bw"}) explicit_false = preprocess(raw, {"conversion_mode": "bw", "invert": False}) np.testing.assert_array_equal(default_result, explicit_false) # ── Mask region tests ── class TestMaskRegions: def test_mask_sets_region_to_white_2d(self): img = np.zeros((100, 100), dtype=np.uint8) result = apply_mask(img, [{"x": 10, "y": 10, "width": 20, "height": 20}]) assert np.all(result[10:30, 10:30] == 255) # Outside mask should still be 0 assert result[0, 0] == 0 def test_mask_sets_region_to_white_3d(self): img = np.zeros((100, 100, 3), dtype=np.uint8) result = apply_mask(img, [{"x": 10, "y": 10, "width": 20, "height": 20}]) assert np.all(result[10:30, 10:30] == 255) assert np.all(result[0, 0] == 0) def test_multiple_masks(self): img = np.zeros((100, 100), dtype=np.uint8) masks = [ {"x": 0, "y": 0, "width": 10, "height": 10}, {"x": 50, "y": 50, "width": 10, "height": 10}, ] result = apply_mask(img, masks) assert np.all(result[0:10, 0:10] == 255) assert np.all(result[50:60, 50:60] == 255) assert result[30, 30] == 0 def test_mask_in_preprocessing_pipeline(self): raw = _make_test_image(color=False) # Mask the entire dark rectangle area result = preprocess(raw, { "conversion_mode": "bw", "mask_regions": [{"x": 5, "y": 5, "width": 50, "height": 50}], }) # The masked area should be white (255) in the output # since we masked the dark rectangle assert result.ndim == 2 def test_zero_size_mask_ignored(self): img = np.zeros((100, 100), dtype=np.uint8) result = apply_mask(img, [{"x": 10, "y": 10, "width": 0, "height": 20}]) assert np.all(result == 0) # ── Preprocessing parameter tests ── class TestPreprocessingParams: def test_threshold_manual_overrides_otsu(self): raw = _make_grayscale_image() # With THRESH_BINARY: pixels > threshold → 255 (white), else → 0 (black). # Disable morphological ops to isolate threshold effect. base = {"conversion_mode": "bw", "morph_dilate_iterations": 0, "morph_erode_iterations": 0} # Low threshold: most pixels are above 30 → lots of white low = preprocess(raw, {**base, "threshold_manual": 30}) # High threshold: only very bright pixels above 240 → less white high = preprocess(raw, {**base, "threshold_manual": 240}) # Low threshold produces more white pixels than high threshold assert np.sum(low == 255) > np.sum(high == 255) def test_morph_dilate_expands_shapes(self): raw = _make_test_image(color=False) no_dilate = preprocess(raw, { "conversion_mode": "bw", "morph_dilate_iterations": 0, "morph_erode_iterations": 0, }) heavy_dilate = preprocess(raw, { "conversion_mode": "bw", "morph_dilate_iterations": 5, "morph_erode_iterations": 0, }) # More dilation = more foreground pixels assert np.sum(heavy_dilate > 0) >= np.sum(no_dilate > 0) def test_morph_erode_shrinks_shapes(self): raw = _make_test_image(color=False) no_erode = preprocess(raw, { "conversion_mode": "bw", "morph_dilate_iterations": 0, "morph_erode_iterations": 0, }) heavy_erode = preprocess(raw, { "conversion_mode": "bw", "morph_dilate_iterations": 0, "morph_erode_iterations": 5, }) # More erosion = fewer foreground pixels assert np.sum(heavy_erode > 0) <= np.sum(no_erode > 0) def test_edge_detect_produces_edges(self): raw = _make_test_image(color=False) normal = preprocess(raw, {"conversion_mode": "bw", "edge_detect": False}) edges = preprocess(raw, {"conversion_mode": "bw", "edge_detect": True}) # Edge detection produces a visually different image assert not np.array_equal(normal, edges) # ── Vectorization with different modes ── class TestVectorizationModes: def test_bw_potrace_produces_svg(self): raw = _make_test_image(color=False) binary = preprocess(raw, {"conversion_mode": "bw"}) svg = potrace_trace(binary) assert "