Spaces:
Paused
Paused
| """ | |
| Unit tests for library/mask_generator.py. | |
| These cover the procedural mask functions that drive --train_inpainting: | |
| shape/size/mode invariants, seed determinism, binary-value invariant, coverage | |
| bounds, and combine_masks logic. visualize_masks.py covers the qualitative side | |
| (eyeballing variety); these tests cover the contractual side. | |
| """ | |
| import os | |
| import sys | |
| import numpy as np | |
| import pytest | |
| from PIL import Image | |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) | |
| from library.mask_generator import ( | |
| cloud_mask, | |
| polygon_mask, | |
| shape_mask, | |
| wobbly_ellipse_mask, | |
| combine_masks, | |
| random_mask, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Helpers | |
| # --------------------------------------------------------------------------- | |
| def _coverage(mask: Image.Image) -> float: | |
| arr = np.array(mask, dtype=np.uint8) | |
| return float((arr >= 128).sum()) / arr.size | |
| def _is_binary(mask: Image.Image) -> bool: | |
| arr = np.array(mask, dtype=np.uint8) | |
| unique = set(np.unique(arr).tolist()) | |
| return unique.issubset({0, 255}) | |
| # --------------------------------------------------------------------------- | |
| # Shape / mode / size invariants | |
| # --------------------------------------------------------------------------- | |
| def test_mask_shape_and_mode(size, fn): | |
| w, h = size | |
| m = fn(w, h, 0) | |
| assert isinstance(m, Image.Image) | |
| assert m.mode == "L" | |
| assert m.size == (w, h) | |
| def test_mask_is_binary(fn): | |
| # Each generator must produce strictly {0, 255} values. | |
| assert _is_binary(fn(42)) | |
| # --------------------------------------------------------------------------- | |
| # Seed determinism | |
| # --------------------------------------------------------------------------- | |
| def test_seed_reproducible(fn): | |
| a = np.array(fn(123)) | |
| b = np.array(fn(123)) | |
| assert np.array_equal(a, b) | |
| def test_seed_diversity(fn): | |
| # Different seeds should very likely give different masks. We sample many | |
| # seeds so a stuck generator can't pass by coincidence. | |
| samples = {np.array(fn(s)).tobytes() for s in range(8)} | |
| assert len(samples) > 1 | |
| # --------------------------------------------------------------------------- | |
| # Coverage bounds | |
| # --------------------------------------------------------------------------- | |
| def test_cloud_threshold_controls_coverage(): | |
| # threshold is approximate but should track monotonically. | |
| low = _coverage(cloud_mask(128, 128, threshold=0.1, seed=0)) | |
| high = _coverage(cloud_mask(128, 128, threshold=0.3, seed=0)) | |
| assert low < high | |
| # threshold ~ target fraction of masked area; allow generous slack. | |
| assert 0.05 <= low <= 0.20 | |
| assert 0.20 <= high <= 0.45 | |
| def test_polygon_coverage_within_bounds(): | |
| # Single polygon, fixed seed; coverage should fall roughly within bounds. | |
| # The bounds describe each polygon's area target; the rasterised result | |
| # can over-/under-shoot a little, so we use loose envelopes. | |
| samples = [ | |
| _coverage(polygon_mask(128, 128, n_points=6, irregularity=0.0, | |
| min_coverage=0.15, max_coverage=0.35, | |
| n_polygons=1, seed=s)) | |
| for s in range(16) | |
| ] | |
| mean = float(np.mean(samples)) | |
| assert 0.05 < mean < 0.5 | |
| def test_wobbly_ellipse_coverage_tracks_target(): | |
| # coverage parameter should approximately control the masked fraction. | |
| samples_low = [_coverage(wobbly_ellipse_mask(128, 128, coverage=0.15, | |
| wobble_scale=0.1, seed=s)) | |
| for s in range(8)] | |
| samples_high = [_coverage(wobbly_ellipse_mask(128, 128, coverage=0.4, | |
| wobble_scale=0.1, seed=s)) | |
| for s in range(8)] | |
| assert np.mean(samples_low) < np.mean(samples_high) | |
| # --------------------------------------------------------------------------- | |
| # combine_masks | |
| # --------------------------------------------------------------------------- | |
| def test_combine_masks_union_is_or(): | |
| a = Image.fromarray(np.array([[0, 255], [0, 0]], dtype=np.uint8), mode="L") | |
| b = Image.fromarray(np.array([[0, 0], [255, 0]], dtype=np.uint8), mode="L") | |
| out = np.array(combine_masks(a, b, mode="union")) | |
| assert np.array_equal(out, np.array([[0, 255], [255, 0]], dtype=np.uint8)) | |
| def test_combine_masks_intersection_is_and(): | |
| a = Image.fromarray(np.array([[255, 255], [0, 0]], dtype=np.uint8), mode="L") | |
| b = Image.fromarray(np.array([[255, 0], [0, 0]], dtype=np.uint8), mode="L") | |
| out = np.array(combine_masks(a, b, mode="intersection")) | |
| assert np.array_equal(out, np.array([[255, 0], [0, 0]], dtype=np.uint8)) | |
| def test_combine_masks_three_inputs_union(): | |
| a = Image.fromarray(np.array([[255, 0]], dtype=np.uint8), mode="L") | |
| b = Image.fromarray(np.array([[0, 0]], dtype=np.uint8), mode="L") | |
| c = Image.fromarray(np.array([[0, 255]], dtype=np.uint8), mode="L") | |
| out = np.array(combine_masks(a, b, c, mode="union")) | |
| assert np.array_equal(out, np.array([[255, 255]], dtype=np.uint8)) | |
| # --------------------------------------------------------------------------- | |
| # random_mask sanity | |
| # --------------------------------------------------------------------------- | |
| def test_random_mask_nonempty(): | |
| # Across a handful of seeds, at least one produced mask must have non-zero | |
| # coverage (a generator that always returns all-zero would silently break | |
| # training without a crash). | |
| coverages = [_coverage(random_mask(128, 128, seed=s)) for s in range(8)] | |
| assert max(coverages) > 0.01 | |