From 68cba874d5a19753bb12c064207ac1db4139c6e7 Mon Sep 17 00:00:00 2001 From: Emil Date: Sun, 28 Jun 2026 16:34:47 +0300 Subject: [PATCH] Add tests (64 passing) + fix bg removal algorithm - test_feedback.py: 46 tests for FeedbackDB (add, rate, search, stats, export, delete) - test_server.py: 18 tests for post-processing (pixelate, bg removal, prompt building) - Fix _pixelate: handle pixel_size=0 without division by zero - Fix _remove_background: remove global magenta normalization that ate interior highlights Now uses single-pass flood-fill from edges, preserving interior bright pixels - Add pytest to .gitignore --- .gitignore | 3 +- feedback.py | 2 +- review_gui.py | 184 +++++++++++++++---------- server.py | 35 ++--- test_feedback.py | 342 +++++++++++++++++++++++++++++++++++++++++++++++ test_server.py | 166 +++++++++++++++++++++++ 6 files changed, 634 insertions(+), 98 deletions(-) create mode 100644 test_feedback.py create mode 100644 test_server.py diff --git a/.gitignore b/.gitignore index 7688996..a14046c 100644 --- a/.gitignore +++ b/.gitignore @@ -17,7 +17,8 @@ output/ feedback.db feedback_dataset.jsonl -# IDE +# pytest +.pytest_cache/ .idea/ .vscode/ *.swp diff --git a/feedback.py b/feedback.py index a62c57e..f0f0f9e 100644 --- a/feedback.py +++ b/feedback.py @@ -41,7 +41,7 @@ class FeedbackDB: @classmethod def open(cls, path: str) -> "FeedbackDB": - conn = sqlite3.connect(path) + conn = sqlite3.connect(path, check_same_thread=False) conn.execute(""" CREATE TABLE IF NOT EXISTS feedback ( id TEXT PRIMARY KEY, diff --git a/review_gui.py b/review_gui.py index ac2415d..8c9c706 100644 --- a/review_gui.py +++ b/review_gui.py @@ -22,6 +22,7 @@ _db = FeedbackDB.open(DB_PATH) _entries = [] _thumbnails = {} _filter = "all" +_texture_registry = None def _load_thumbnail(path, tag): @@ -34,9 +35,17 @@ def _load_thumbnail(path, tag): bg = Image.new("RGBA", (THUMB_SIZE, THUMB_SIZE), (40, 40, 40, 255)) offset = ((THUMB_SIZE - img.width) // 2, (THUMB_SIZE - img.height) // 2) bg.paste(img, offset, img if img.mode == "RGBA" else None) - rgb = bg.convert("RGB") - dpg.add_static_texture(THUMB_SIZE, THUMB_SIZE, list(rgb.getdata()), tag=tag) + pixels = list(bg.getdata()) + flat = [c / 255.0 for pixel in pixels for c in pixel] + + dpg.add_static_texture( + THUMB_SIZE, + THUMB_SIZE, + flat, + tag=tag, + parent=_texture_registry, + ) _thumbnails[tag] = True return True except Exception: @@ -59,74 +68,105 @@ def _refresh(): f"{stats.total} sprites | {stats.rated} rated | {stats.unrated} unrated | avg {stats.avg_rating:.1f}", ) - dpg.delete_item("sprite_list", children=True) + dpg.delete_item("sprite_list", children_only=True) - for i, entry in enumerate(_entries): - with dpg.group(parent="sprite_list", tag=f"row_{entry.id}"): - with dpg.group(horizontal=True): - thumb_tag = f"thumb_{entry.id}" - if _load_thumbnail(entry.image_path, thumb_tag): - dpg.add_image(thumb_tag) - else: - dpg.add_text("[no img]", default_value="[no img]") + if not _entries: + dpg.add_text( + "No sprites in database. Generate some first.", parent="sprite_list" + ) + return - with dpg.group(): - prompt_display = ( - entry.prompt[:90] + "..." - if len(entry.prompt) > 90 - else entry.prompt + for entry in _entries: + _build_entry_row(entry) + + +def _build_entry_row(entry): + with dpg.group(parent="sprite_list"): + with dpg.group(horizontal=True): + thumb_tag = f"thumb_{entry.id}" + if _load_thumbnail(entry.image_path, thumb_tag): + dpg.add_image(thumb_tag) + else: + dpg.add_text("[no img]") + + dpg.add_spacer(width=8) + + with dpg.group(width=700): + prompt_display = ( + entry.prompt[:90] + "..." + if len(entry.prompt) > 90 + else entry.prompt + ) + dpg.add_text(prompt_display, color=(140, 180, 255), wrap=650) + + dpg.add_spacer(height=2) + + if entry.image_path: + p = entry.image_path + if len(p) > 80: + p = "..." + p[-77:] + dpg.add_text(p, color=(100, 100, 100)) + dpg.add_spacer(height=4) + + with dpg.group(horizontal=True): + for star in range(1, 6): + filled = star <= entry.rating + label = "\u2605" if filled else "\u2606" + color = (255, 200, 80) if filled else (80, 80, 90) + btn_tag = f"star_{entry.id}_{star}" + dpg.add_button( + label=label, + tag=btn_tag, + callback=lambda s, a, e=entry: _rate(e, s), + width=28, + height=24, + ) + dpg.bind_item_theme(btn_tag, _star_theme(color)) + + dpg.add_spacer(width=8) + + rating_text = ( + f"{entry.rating}\u2605" if entry.rating > 0 else "unrated" ) - dpg.add_text(prompt_display, color=(140, 180, 255)) + dpg.add_text(rating_text, tag=f"rating_label_{entry.id}") - if entry.image_path: - p = entry.image_path - if len(p) > 80: - p = "..." + p[-77:] - dpg.add_text(p, color=(100, 100, 100)) + dpg.add_spacer(width=12) - with dpg.group(horizontal=True): - for star in range(1, 6): - filled = star <= entry.rating - label = "★" if filled else "☆" - color = (255, 200, 80) if filled else (80, 80, 90) - dpg.add_button( - label=label, - tag=f"star_{entry.id}_{star}", - callback=lambda s, a, e=entry: _rate(e, s), - width=24, - ) - dpg.bind_item_theme( - f"star_{entry.id}_{star}", _star_theme(color) - ) + dpg.add_input_text( + hint="feedback...", + default_value=entry.feedback or "", + tag=f"feedback_{entry.id}", + width=200, + height=24, + ) - rating_text = ( - f"{entry.rating}★" if entry.rating > 0 else "unrated" - ) - dpg.add_text(rating_text, tag=f"rating_label_{entry.id}") + dpg.add_spacer(width=4) - dpg.add_input_text( - hint="feedback...", - default_value=entry.feedback or "", - tag=f"feedback_{entry.id}", - width=180, - ) + dpg.add_button( + label="Save", + callback=lambda s, a, e=entry: _save_feedback(e), + width=55, + height=24, + ) - dpg.add_button( - label="Save", - callback=lambda s, a, e=entry: _save_feedback(e), - width=50, - ) + dpg.add_spacer(width=4) - dpg.add_button( - label="Del", - callback=lambda s, a, e=entry: _delete(e), - width=40, - ) + dpg.add_button( + label="Del", + callback=lambda s, a, e=entry: _delete(e), + width=40, + height=24, + ) - if entry.feedback: - dpg.add_text(f'saved: "{entry.feedback}"', color=(80, 180, 100)) + if entry.feedback: + dpg.add_spacer(height=2) + dpg.add_text( + f'saved: "{entry.feedback}"', color=(80, 180, 100), wrap=650 + ) - dpg.add_separator() + dpg.add_spacer(height=4) + dpg.add_separator() + dpg.add_spacer(height=4) def _rate(entry, star_tag): @@ -169,8 +209,6 @@ def _on_filter(sender, app_data): def _export_dataset(): - import dearpygui.dearpygui as dpg_filedialog - def _do_export(sender, app_data): path = ( app_data["file_path_name"] @@ -198,20 +236,23 @@ def _star_theme(color): dpg.add_theme_color(dpg.mvThemeCol_Button, (0, 0, 0, 0)) dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, (60, 60, 60, 50)) dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, (40, 40, 40, 50)) - with dpg.theme_component(dpg.mvButton, enabled_state=False): - dpg.add_theme_color(dpg.mvThemeCol_Text, color) return theme def main(): dpg.create_context() + global _texture_registry + _texture_registry = dpg.add_texture_registry() + with dpg.window(tag="main_window"): with dpg.group(horizontal=True): - dpg.add_button(label="Refresh", callback=_refresh) - dpg.add_button(label="Export Dataset", callback=_export_dataset) - - dpg.add_text("Filter:", indent=20) + dpg.add_button(label="Refresh", callback=_refresh, width=80) + dpg.add_spacer(width=4) + dpg.add_button(label="Export Dataset", callback=_export_dataset, width=120) + dpg.add_spacer(width=20) + dpg.add_text("Filter:") + dpg.add_spacer(width=4) dpg.add_radio_button( ["all", "unrated", "top"], tag="filter_radio", @@ -220,18 +261,17 @@ def main(): default_value="all", ) + dpg.add_spacer(height=4) dpg.add_text("", tag="stats_text", color=(180, 180, 200)) + dpg.add_spacer(height=4) dpg.add_separator() + dpg.add_spacer(height=4) with dpg.child_window(tag="sprite_list", autosize_x=True, autosize_y=True): pass - dpg.create_viewport(title="Imagen — Sprite Review", width=1000, height=750) + dpg.create_viewport(title="Imagen - Sprite Review", width=1000, height=750) dpg.set_viewport_resizable(True) - dpg.configure_app(init_file="") - - texture_registry = dpg.add_texture_registry() - dpg.bind_texture_registry(texture_registry) dpg.setup_dearpygui() dpg.show_viewport() diff --git a/server.py b/server.py index 85edb22..1e67ab7 100644 --- a/server.py +++ b/server.py @@ -145,12 +145,11 @@ def _generate( def _remove_background(image: Image.Image, threshold: int = 30) -> Image.Image: """Remove background using flood-fill from edges. - Two-pass approach: - 1. Detect border color, replace all near-border pixels with a flat fill color - 2. Flood-fill from edges to remove the flat color cleanly - - This normalizes gradient/noisy backgrounds into one solid color, - making flood-fill removal much cleaner. + Detects the border color, then flood-fills from all edge pixels, + removing any pixel that is within threshold of the border color AND + connected to the border. Interior pixels with similar colors (e.g. + highlights on armor) are preserved because they are not connected + to the border through similar-colored pixels. """ from collections import deque @@ -170,20 +169,10 @@ def _remove_background(image: Image.Image, threshold: int = 30) -> Image.Image: border_colors = np.array(border_colors) bg_color = np.median(border_colors, axis=0).astype(int) - # Pass 1: normalize background — replace all pixels within threshold - # of border color with a flat fill color (pure magenta, unlikely in sprites) - fill_color = np.array([255, 0, 255], dtype=int) - dist_to_bg = np.abs(arr - bg_color).sum(axis=2) - bg_mask = dist_to_bg < threshold * 3 - arr[bg_mask] = fill_color - - # Pass 2: flood-fill from edges to remove connected fill_color regions alpha = np.full((h, w), 255, dtype=np.uint8) visited = np.zeros((h, w), dtype=bool) queue = deque() - fill_dist_threshold = 30 # tolerance for near-fill pixels - # Seed from all border pixels for x in range(w): for y in [0, h - 1]: @@ -196,11 +185,11 @@ def _remove_background(image: Image.Image, threshold: int = 30) -> Image.Image: queue.append((y, x)) visited[y, x] = True - # BFS flood-fill + # BFS flood-fill: remove pixels close to bg_color that are connected to border while queue: y, x = queue.popleft() - dist = np.abs(arr[y, x] - fill_color).sum() - if dist > fill_dist_threshold: + dist = np.abs(arr[y, x] - bg_color).sum() + if dist > threshold * 3: continue alpha[y, x] = 0 @@ -210,11 +199,6 @@ def _remove_background(image: Image.Image, threshold: int = 30) -> Image.Image: visited[ny, nx] = True queue.append((ny, nx)) - # Clean up: any remaining near-magenta pixels that weren't flood-filled - # (small isolated background pockets) get removed too - remaining_bg = np.abs(arr - fill_color).sum(axis=2) < fill_dist_threshold - alpha[remaining_bg] = 0 - rgba = np.dstack([arr.astype(np.uint8), alpha]) return Image.fromarray(rgba, mode="RGBA") @@ -223,7 +207,10 @@ def _pixelate(image: Image.Image, pixel_size: int = 8) -> Image.Image: """Downscale then upscale with NEAREST to create chunky pixel-art effect. pixel_size=8 means each "pixel" in the result is an 8x8 block. + pixel_size=0 returns the original image unchanged. """ + if pixel_size <= 0: + return image w, h = image.size small = image.resize((w // pixel_size, h // pixel_size), Image.LANCZOS) return small.resize((w, h), Image.NEAREST) diff --git a/test_feedback.py b/test_feedback.py new file mode 100644 index 0000000..8d07809 --- /dev/null +++ b/test_feedback.py @@ -0,0 +1,342 @@ +""" +Tests for FeedbackDB — SQLite-backed feedback database. +""" + +import json +import os +import sys +import tempfile + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from feedback import FeedbackDB, FeedbackEntry, DBStats, _tokenize + + +@pytest.fixture +def db(): + """Create a temporary in-file database.""" + with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: + path = f.name + database = FeedbackDB.open(path) + yield database + os.unlink(path) + + +@pytest.fixture +def db_mem(): + """Create an in-memory database via temp file (SQLite needs a path).""" + with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: + path = f.name + database = FeedbackDB.open(path) + yield database + os.unlink(path) + + +def _add_sample( + db, prompt="a brave knight", rating=0, feedback=None, image_path="/tmp/test.png" +): + entry_id = db.add( + prompt=prompt, + params={"seed": 42, "width": 512, "height": 512, "steps": 8}, + image_path=image_path, + ) + if rating > 0: + db.update_rating(entry_id, rating, feedback) + return entry_id + + +class TestOpenAndInit: + def test_open_creates_tables(self, db): + stats = db.stats() + assert stats.total == 0 + + def test_open_creates_indexes(self, db): + cursor = db.conn.execute("SELECT name FROM sqlite_master WHERE type='index'") + names = [row[0] for row in cursor] + assert "idx_prompt" in names + assert "idx_rating" in names + + +class TestAdd: + def test_add_returns_id(self, db): + entry_id = db.add("knight", {"seed": 1}, "/tmp/k.png") + assert isinstance(entry_id, str) + assert len(entry_id) > 0 + + def test_add_increments_count(self, db): + db.add("knight", {"seed": 1}) + db.add("archer", {"seed": 2}) + assert db.stats().total == 2 + + def test_add_default_rating_is_zero(self, db): + entry_id = db.add("knight", {"seed": 1}) + entries = db.get_all() + assert entries[0].rating == 0 + + def test_add_stores_params(self, db): + params = {"seed": 42, "width": 512, "height": 512, "steps": 8} + db.add("knight", params) + entries = db.get_all() + assert entries[0].params == params + + def test_add_stores_image_path(self, db): + db.add("knight", {"seed": 1}, "/tmp/knight.png") + entries = db.get_all() + assert entries[0].image_path == "/tmp/knight.png" + + def test_add_without_image_path(self, db): + db.add("knight", {"seed": 1}, None) + entries = db.get_all() + assert entries[0].image_path is None + + +class TestUpdateRating: + def test_update_rating(self, db): + entry_id = _add_sample(db) + db.update_rating(entry_id, 5) + entries = db.get_all() + assert entries[0].rating == 5 + + def test_update_rating_with_feedback(self, db): + entry_id = _add_sample(db) + db.update_rating(entry_id, 4, "great colors") + entries = db.get_all() + assert entries[0].rating == 4 + assert entries[0].feedback == "great colors" + + def test_update_rating_clamps_high(self, db): + entry_id = _add_sample(db) + db.update_rating(entry_id, 10) + entries = db.get_all() + assert entries[0].rating == 5 + + def test_update_rating_clamps_negative(self, db): + entry_id = _add_sample(db) + db.update_rating(entry_id, -3) + entries = db.get_all() + assert entries[0].rating == 0 + + def test_update_rating_to_zero(self, db): + entry_id = _add_sample(db, rating=5) + db.update_rating(entry_id, 0) + entries = db.get_all() + assert entries[0].rating == 0 + + +class TestGetAll: + def test_get_all_empty(self, db): + assert db.get_all() == [] + + def test_get_all_returns_entries(self, db): + _add_sample(db, "knight") + _add_sample(db, "archer") + entries = db.get_all() + assert len(entries) == 2 + + def test_get_all_ordered_newest_first(self, db): + id1 = _add_sample(db, "first") + # Force different timestamp + import time as _time + + _time.sleep(1.1) + id2 = _add_sample(db, "second") + entries = db.get_all() + assert entries[0].prompt == "second" + assert entries[1].prompt == "first" + + def test_get_all_returns_feedback_entry(self, db): + _add_sample(db, "knight", rating=5, feedback="perfect") + entry = db.get_all()[0] + assert isinstance(entry, FeedbackEntry) + assert entry.prompt == "knight" + assert entry.rating == 5 + assert entry.feedback == "perfect" + + +class TestGetUnrated: + def test_get_unrated_empty(self, db): + assert db.get_unrated() == [] + + def test_get_unrated_only_unrated(self, db): + id1 = _add_sample(db, "rated", rating=5) + id2 = _add_sample(db, "unrated") + unrated = db.get_unrated() + assert len(unrated) == 1 + assert unrated[0].prompt == "unrated" + + def test_get_unrated_all_unrated(self, db): + _add_sample(db, "sprite1") + _add_sample(db, "sprite2") + assert len(db.get_unrated()) == 2 + + +class TestTopRated: + def test_top_rated_empty(self, db): + assert db.top_rated(10, 1) == [] + + def test_top_rated_filters_min_rating(self, db): + _add_sample(db, "low", rating=1) + _add_sample(db, "high", rating=5) + top = db.top_rated(10, 4) + assert len(top) == 1 + assert top[0].prompt == "high" + + def test_top_rated_orders_by_rating(self, db): + _add_sample(db, "three", rating=3) + _add_sample(db, "five", rating=5) + _add_sample(db, "four", rating=4) + top = db.top_rated(3, 1) + assert top[0].rating == 5 + assert top[1].rating == 4 + assert top[2].rating == 3 + + def test_top_rated_respects_limit(self, db): + for i in range(10): + _add_sample(db, f"sprite_{i}", rating=5) + top = db.top_rated(3, 1) + assert len(top) == 3 + + +class TestSearchSimilar: + def test_search_similar_empty_db(self, db): + assert db.search_similar("knight", 5) == [] + + def test_search_similar_exact_match(self, db): + _add_sample(db, "brave knight", rating=5) + _add_sample(db, "fire dragon", rating=4) + results = db.search_similar("knight", 5) + assert len(results) == 1 + assert "knight" in results[0].prompt + + def test_search_similar_multiple_keywords(self, db): + _add_sample(db, "brave knight", rating=5) + _add_sample(db, "brave warrior", rating=4) + _add_sample(db, "fire dragon", rating=3) + results = db.search_similar("brave knight", 5) + assert len(results) == 2 + assert "knight" in results[0].prompt + assert "warrior" in results[1].prompt + + def test_search_similar_no_match(self, db): + _add_sample(db, "fire dragon", rating=5) + results = db.search_similar("knight", 5) + assert results == [] + + def test_search_similar_empty_query_returns_top(self, db): + _add_sample(db, "knight", rating=5) + _add_sample(db, "dragon", rating=3) + results = db.search_similar("", 5) + assert len(results) == 2 + + def test_search_similar_respects_limit(self, db): + for i in range(10): + _add_sample(db, f"knight variant {i}", rating=5) + results = db.search_similar("knight", 3) + assert len(results) == 3 + + +class TestStats: + def test_stats_empty(self, db): + stats = db.stats() + assert stats.total == 0 + assert stats.rated == 0 + assert stats.unrated == 0 + assert stats.avg_rating == 0.0 + + def test_stats_with_entries(self, db): + _add_sample(db, "s1", rating=4) + _add_sample(db, "s2", rating=2) + _add_sample(db, "s3") + stats = db.stats() + assert stats.total == 3 + assert stats.rated == 2 + assert stats.unrated == 1 + assert abs(stats.avg_rating - 3.0) < 0.1 + + def test_stats_all_unrated(self, db): + _add_sample(db, "s1") + _add_sample(db, "s2") + stats = db.stats() + assert stats.rated == 0 + assert stats.unrated == 2 + assert stats.avg_rating == 0.0 + + +class TestDelete: + def test_delete_entry(self, db): + entry_id = _add_sample(db, "knight") + assert db.stats().total == 1 + db.delete(entry_id) + assert db.stats().total == 0 + + def test_delete_nonexistent_id(self, db): + db.delete("nonexistent-id") + assert db.stats().total == 0 + + def test_delete_specific_entry(self, db): + id1 = _add_sample(db, "knight") + id2 = _add_sample(db, "archer") + db.delete(id1) + entries = db.get_all() + assert len(entries) == 1 + assert entries[0].prompt == "archer" + + +class TestExportJsonl: + def test_export_jsonl(self, db, tmp_path): + _add_sample(db, "knight", rating=5, feedback="great") + _add_sample(db, "archer", rating=4) + _add_sample(db, "goblin", rating=1) + + path = str(tmp_path / "export.jsonl") + count = db.export_jsonl(path, min_rating=4) + + assert count == 2 + assert os.path.exists(path) + + with open(path) as f: + lines = f.readlines() + + assert len(lines) == 2 + data = json.loads(lines[0]) + assert "instruction" in data + assert "response" in data + assert "rating" in data + + def test_export_jsonl_empty(self, db, tmp_path): + path = str(tmp_path / "empty.jsonl") + count = db.export_jsonl(path, min_rating=4) + assert count == 0 + + def test_export_jsonl_min_rating_filter(self, db, tmp_path): + _add_sample(db, "high", rating=5) + _add_sample(db, "mid", rating=3) + _add_sample(db, "low", rating=1) + + path = str(tmp_path / "filter.jsonl") + count = db.export_jsonl(path, min_rating=3) + assert count == 2 + + +class TestTokenize: + def test_simple_words(self): + assert _tokenize("knight armor") == ["knight", "armor"] + + def test_underscores(self): + assert _tokenize("missile_launch") == ["missile", "launch"] + + def test_hyphens(self): + assert _tokenize("fire-ball") == ["fire", "ball"] + + def test_single_char_filtered(self): + assert _tokenize("a b c") == [] + + def test_mixed_case(self): + assert _tokenize("Brave Knight") == ["brave", "knight"] + + def test_empty_string(self): + assert _tokenize("") == [] + + def test_only_separators(self): + assert _tokenize("_ - _") == [] diff --git a/test_server.py b/test_server.py new file mode 100644 index 0000000..eb81bf2 --- /dev/null +++ b/test_server.py @@ -0,0 +1,166 @@ +""" +Tests for server.py — post-processing functions (pixelate, background removal, prompt building). + +Does NOT test model loading or generation (requires GPU). +""" + +import os +import sys + +import numpy as np +from PIL import Image + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import server + + +class TestBuildPrompt: + def test_basic_prompt(self): + result = server._build_prompt("a brave knight") + assert "pixel art" in result + assert "a brave knight" in result + assert "game asset" in result + + def test_empty_prompt(self): + result = server._build_prompt("") + assert "pixel art" in result + + +class TestPixelate: + def _make_image(self, size=512): + arr = np.random.randint(0, 255, (size, size, 3), dtype=np.uint8) + return Image.fromarray(arr, mode="RGB") + + def test_pixelate_preserves_size(self): + img = self._make_image(512) + result = server._pixelate(img, pixel_size=4) + assert result.size == (512, 512) + + def test_pixelate_zero_returns_original(self): + img = self._make_image(64) + result = server._pixelate(img, pixel_size=0) + assert result.size == (64, 64) + + def test_pixelate_reduces_unique_colors(self): + arr = np.zeros((64, 64, 3), dtype=np.uint8) + for y in range(64): + for x in range(64): + arr[y, x] = [x * 4, y * 4, 128] + img = Image.fromarray(arr, mode="RGB") + + result = server._pixelate(img, pixel_size=8) + + original_colors = len(np.unique(np.array(img).reshape(-1, 3), axis=0)) + result_colors = len(np.unique(np.array(result).reshape(-1, 3), axis=0)) + assert result_colors <= original_colors + + def test_pixelate_creates_blocks(self): + arr = np.zeros((8, 8, 3), dtype=np.uint8) + arr[:4, :4] = [255, 0, 0] + arr[4:, 4:] = [0, 255, 0] + img = Image.fromarray(arr, mode="RGB") + + result = server._pixelate(img, pixel_size=4) + result_arr = np.array(result) + + block_tl = result_arr[:4, :4] + assert np.all(block_tl == block_tl[0, 0]) + + block_br = result_arr[4:, 4:] + assert np.all(block_br == block_br[0, 0]) + + +class TestRemoveBackground: + def _make_sprite_on_white(self, size=64): + arr = np.full((size, size, 3), 255, dtype=np.uint8) + arr[16:48, 16:48] = [200, 50, 50] + return Image.fromarray(arr, mode="RGB") + + def test_returns_rgba(self): + img = self._make_sprite_on_white() + result = server._remove_background(img) + assert result.mode == "RGBA" + + def test_white_background_becomes_transparent(self): + img = self._make_sprite_on_white() + result = server._remove_background(img) + arr = np.array(result) + alpha = arr[:, :, 3] + + corner_alpha = alpha[0, 0] + assert corner_alpha == 0 + + def test_sprite_pixels_remain_opaque(self): + img = self._make_sprite_on_white() + result = server._remove_background(img) + arr = np.array(result) + alpha = arr[:, :, 3] + + center_alpha = alpha[32, 32] + assert center_alpha == 255 + + def test_preserves_sprite_colors(self): + img = self._make_sprite_on_white() + result = server._remove_background(img) + arr = np.array(result) + center_pixel = arr[32, 32] + assert center_pixel[0] == 200 + assert center_pixel[1] == 50 + assert center_pixel[2] == 50 + + def test_interior_bright_pixel_not_removed(self): + # Sprite is dark red on white background, with a white highlight inside + arr = np.full((64, 64, 3), 255, dtype=np.uint8) + arr[8:56, 8:56] = [180, 30, 30] + # White highlight inside the sprite (like a shine on armor) + arr[30:34, 30:34] = [255, 255, 255] + img = Image.fromarray(arr, mode="RGB") + + result = server._remove_background(img) + result_arr = np.array(result) + alpha = result_arr[:, :, 3] + + # Background corners should be transparent + assert alpha[0, 0] == 0 + # Interior highlight should be opaque (not removed as background) + assert alpha[32, 32] == 255 + # Sprite body should be opaque + assert alpha[20, 20] == 255 + + def test_black_background_removed(self): + arr = np.zeros((64, 64, 3), dtype=np.uint8) + arr[16:48, 16:48] = [200, 50, 50] + img = Image.fromarray(arr, mode="RGB") + + result = server._remove_background(img) + result_arr = np.array(result) + alpha = result_arr[:, :, 3] + + assert alpha[0, 0] == 0 + assert alpha[32, 32] == 255 + + +class TestEnsureDir: + def test_ensure_dir_creates_nested(self, tmp_path): + path = str(tmp_path / "a" / "b" / "c" / "test.png") + server._ensure_dir(path) + assert os.path.isdir(str(tmp_path / "a" / "b" / "c")) + + def test_ensure_dir_empty_dir(self): + server._ensure_dir("test.png") + + +class TestEnvPaths: + def test_default_model_dir(self): + assert "models" in server.MODEL_DIR + + def test_default_lora_dir(self): + assert "models" in server.LORA_DIR + + def test_default_output_dir(self): + assert "output" in server.OUTPUT_DIR + + def test_default_db_path(self): + assert "feedback.db" in server.DB_PATH