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
This commit is contained in:
Emil
+2
-1
@@ -17,7 +17,8 @@ output/
|
||||
feedback.db
|
||||
feedback_dataset.jsonl
|
||||
|
||||
# IDE
|
||||
# pytest
|
||||
.pytest_cache/
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
|
||||
+1
-1
@@ -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,
|
||||
|
||||
+112
-72
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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("_ - _") == []
|
||||
+166
@@ -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
|
||||
Reference in New Issue
Block a user