Spaces:
Runtime error
Runtime error
| """ | |
| tests/test_extraction.py | |
| -------------------------- | |
| Tests for pixel-exact sprite extraction: cropping, padding, and the | |
| NEAREST-only preview rule. This is the module responsible for the | |
| project's most important guarantee: ORIGINAL PIXELS = OUTPUT PIXELS. | |
| Run with: pytest tests/test_extraction.py -v | |
| """ | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| from PIL import Image | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) | |
| from background import detect_background | |
| from image_utils import crop_exact, nearest_preview, pad_nearest, to_pil | |
| from segmentation import DetectedSprite | |
| from image_utils import BBox | |
| from sprite_extractor import extract_sprite | |
| def test_crop_exact_matches_source_pixels_exactly(): | |
| """The crop must contain the EXACT same pixel values as the source region.""" | |
| sheet = np.zeros((30, 30, 4), dtype=np.uint8) | |
| sheet[:, :, 3] = 255 | |
| rng = np.random.default_rng(1) | |
| sheet[5:15, 5:15] = np.dstack([ | |
| rng.integers(0, 255, (10, 10)), | |
| rng.integers(0, 255, (10, 10)), | |
| rng.integers(0, 255, (10, 10)), | |
| np.full((10, 10), 255), | |
| ]).astype(np.uint8) | |
| crop = crop_exact(sheet, 5, 5, 10, 10) | |
| assert crop.shape == (10, 10, 4) | |
| assert np.array_equal(crop, sheet[5:15, 5:15]) | |
| def test_crop_clamps_out_of_bounds_without_crashing(): | |
| sheet = np.zeros((10, 10, 4), dtype=np.uint8) | |
| crop = crop_exact(sheet, 5, 5, 100, 100) # way beyond sheet bounds | |
| assert crop.shape[0] <= 5 | |
| assert crop.shape[1] <= 5 | |
| def test_pad_nearest_zero_padding_is_noop(): | |
| arr = np.ones((5, 5, 4), dtype=np.uint8) * 7 | |
| out = pad_nearest(arr, 0) | |
| assert np.array_equal(out, arr) | |
| def test_pad_nearest_replicates_edge_pixels(): | |
| arr = np.zeros((3, 3, 4), dtype=np.uint8) | |
| arr[1, 1] = [255, 0, 0, 255] # center pixel distinct | |
| arr[0, 0] = [10, 20, 30, 255] # corner pixel | |
| out = pad_nearest(arr, 2) | |
| assert out.shape == (7, 7, 4) # 3 + 2*2 on each axis | |
| # The padded corner region should replicate the original corner pixel value. | |
| assert np.array_equal(out[0, 0], arr[0, 0]) | |
| assert np.array_equal(out[0, 1], arr[0, 0]) | |
| def test_pad_nearest_does_not_alter_original_pixels(): | |
| arr = np.zeros((4, 4, 4), dtype=np.uint8) | |
| rng = np.random.default_rng(2) | |
| arr[:, :, :3] = rng.integers(0, 255, (4, 4, 3)) | |
| arr[:, :, 3] = 255 | |
| out = pad_nearest(arr, 3) | |
| # Original pixels must appear unchanged, exactly offset by the padding. | |
| assert np.array_equal(out[3:7, 3:7], arr) | |
| def test_nearest_preview_uses_nearest_resampling_only(): | |
| """ | |
| A tiny checkerboard, when upscaled via nearest_preview, must produce | |
| HARD pixel boundaries (no blended/intermediate colors) -- this is the | |
| signature of NEAREST resampling vs. any smooth interpolation method. | |
| """ | |
| tiny = np.zeros((2, 2, 4), dtype=np.uint8) | |
| tiny[0, 0] = [255, 0, 0, 255] | |
| tiny[0, 1] = [0, 255, 0, 255] | |
| tiny[1, 0] = [0, 0, 255, 255] | |
| tiny[1, 1] = [255, 255, 0, 255] | |
| img = to_pil(tiny) | |
| preview = nearest_preview(img, target_long_side=8) | |
| arr = np.array(preview) | |
| unique_colors = np.unique(arr.reshape(-1, 4), axis=0) | |
| # NEAREST upscaling of a 4-color 2x2 image must produce EXACTLY those | |
| # 4 colors, never a blended/interpolated fifth color. | |
| assert len(unique_colors) == 4 | |
| expected = {(255, 0, 0, 255), (0, 255, 0, 255), (0, 0, 255, 255), (255, 255, 0, 255)} | |
| actual = {tuple(c) for c in unique_colors} | |
| assert actual == expected | |
| def test_extract_sprite_produces_exact_pixel_crop_no_resize(): | |
| sheet = np.zeros((40, 40, 4), dtype=np.uint8) | |
| sheet[:, :, 3] = 255 | |
| sheet[10:20, 10:25] = [90, 200, 90, 255] # 15x10 sprite region | |
| bg_info = detect_background(sheet) | |
| detection = DetectedSprite( | |
| sprite_index=1, bbox=BBox(10, 10, 15, 10), confidence=100.0, | |
| confidence_reasons=[], component_count=1, pixel_count=150, | |
| ) | |
| extracted = extract_sprite(sheet, bg_info, detection, source_file="test.png") | |
| assert extracted.rgba.shape == (10, 15, 4) # height, width, channels -- NOT resized | |
| assert np.array_equal(extracted.rgba, sheet[10:20, 10:25]) | |
| def test_extract_sprite_with_padding_preserves_original_pixels_at_offset(): | |
| sheet = np.zeros((40, 40, 4), dtype=np.uint8) | |
| sheet[:, :, 3] = 255 | |
| sheet[10:20, 10:20] = [90, 200, 90, 255] | |
| bg_info = detect_background(sheet) | |
| detection = DetectedSprite( | |
| sprite_index=1, bbox=BBox(10, 10, 10, 10), confidence=100.0, | |
| confidence_reasons=[], component_count=1, pixel_count=100, | |
| ) | |
| extracted = extract_sprite(sheet, bg_info, detection, source_file="test.png", padding=2) | |
| assert extracted.rgba.shape == (14, 14, 4) # 10 + 2*2 | |
| assert np.array_equal(extracted.rgba[2:12, 2:12], sheet[10:20, 10:20]) | |
| def test_extract_sprite_transparent_variant_only_clears_background(): | |
| """ | |
| The background-transparent export variant must only clear pixels that | |
| were actually part of the border-connected background -- a sprite's | |
| own same-colored interior pixel must remain opaque. | |
| """ | |
| sheet = np.zeros((30, 30, 4), dtype=np.uint8) | |
| sheet[:, :, 3] = 255 # black background | |
| sheet[5:15, 5:15] = [200, 50, 50, 255] | |
| sheet[9, 9] = [0, 0, 0, 255] # black pixel inside the sprite | |
| bg_info = detect_background(sheet) | |
| detection = DetectedSprite( | |
| sprite_index=1, bbox=BBox(5, 5, 10, 10), confidence=100.0, | |
| confidence_reasons=[], component_count=1, pixel_count=100, | |
| ) | |
| extracted = extract_sprite(sheet, bg_info, detection, source_file="test.png") | |
| # relative position of (9,9) within the crop starting at (5,5) | |
| rel_y, rel_x = 9 - 5, 9 - 5 | |
| assert extracted.rgba_transparent_bg[rel_y, rel_x, 3] == 255 # untouched | |
| assert extracted.rgba[rel_y, rel_x, 3] == 255 | |
| def test_png_roundtrip_is_lossless(): | |
| """Saving to PNG and reloading must reproduce identical pixel values (PNG is lossless).""" | |
| import io | |
| arr = np.zeros((12, 12, 4), dtype=np.uint8) | |
| rng = np.random.default_rng(3) | |
| arr[:, :, :3] = rng.integers(0, 255, (12, 12, 3)) | |
| arr[:, :, 3] = 255 | |
| img = to_pil(arr) | |
| buf = io.BytesIO() | |
| img.save(buf, format="PNG") | |
| buf.seek(0) | |
| reloaded = np.array(Image.open(buf).convert("RGBA")) | |
| assert np.array_equal(arr, reloaded) | |
| if __name__ == "__main__": | |
| import traceback | |
| tests = [(name, fn) for name, fn in list(globals().items()) if name.startswith("test_") and callable(fn)] | |
| passed, failed = 0, 0 | |
| for name, fn in tests: | |
| try: | |
| fn() | |
| print(f"PASS {name}") | |
| passed += 1 | |
| except Exception: | |
| print(f"FAIL {name}") | |
| traceback.print_exc() | |
| failed += 1 | |
| print(f"\n{passed} passed, {failed} failed") | |
| sys.exit(1 if failed else 0) | |