from __future__ import annotations from pathlib import Path import pytest from PIL import Image from kneiff.datasets.export.image_grid import ( build_existing_training_export_result, TRAINING_IMAGE_GRID_FILENAME, write_training_image_grid, ) from kneiff.datasets.export.types import ( TrainingExportResult, TrainingSamplingResult, TrainingSubsetExportResult, TrainingSubsetSamplingResult, ) def _subset_result(subset: str) -> TrainingSubsetExportResult: """Return a tiny export subset summary for image-grid tests.""" return TrainingSubsetExportResult( subset=subset, source_image_count=1, output_image_count=1, caption_file_count=1, min_outputs_per_source=1, max_outputs_per_source=1, average_outputs_per_source=1.0, ) def _export_result(export_root: Path, *, dry_run: bool = False) -> TrainingExportResult: """Return a tiny export result with one training subset and one SFW subset.""" return TrainingExportResult( export_root=export_root, source_image_count=2, output_image_count=2, caption_file_count=2, min_outputs_per_source=1, max_outputs_per_source=1, average_outputs_per_source=1.0, subset_results=(_subset_result("fullbody"), _subset_result("sfw")), dry_run=dry_run, ) def _write_subset_image( export_root: Path, subset: str, color: str, *, filename: str = "scene__orig.png", ) -> None: """Write one exported image below a subset directory.""" subset_dir = export_root / subset subset_dir.mkdir(parents=True, exist_ok=True) Image.new("RGB", (16, 16), color=color).save(subset_dir / filename) def test_training_image_grid_excludes_mirrored_variants(tmp_path: Path) -> None: export_root = tmp_path / "export" output_path = tmp_path / "grid.png" _write_subset_image(export_root, "fullbody", "red") _write_subset_image( export_root, "fullbody", "green", filename="scene__orig__tag.png", ) _write_subset_image( export_root, "fullbody", "blue", filename="scene__aug-mirror.png", ) _write_subset_image( export_root, "fullbody", "yellow", filename="scene__aug-mirror__tag.png", ) grid_path = write_training_image_grid( _export_result(export_root), output_path=output_path, thumbnail_size=16, ) assert grid_path == output_path with Image.open(output_path) as rendered: colors = rendered.convert("RGB").getcolors(maxcolors=100000) assert colors is not None color_counts = {color: count for count, color in colors} assert color_counts.get((255, 0, 0), 0) > 0 assert color_counts.get((0, 128, 0), 0) > 0 assert color_counts.get((0, 0, 255), 0) == 0 assert color_counts.get((255, 255, 0), 0) == 0 def test_training_image_grid_excludes_sfw_when_training_sampling_omits_it( tmp_path: Path, ) -> None: export_root = tmp_path / "export" _write_subset_image(export_root, "fullbody", "red") _write_subset_image(export_root, "sfw", "green") sampling = TrainingSamplingResult( sampling_method="auto-weighting", subset_results=( TrainingSubsetSamplingResult( subset="fullbody", probability=1.0, relative_percent=100.0, ), ), ) sampled_grid = write_training_image_grid( _export_result(export_root), training_sampling=sampling, thumbnail_size=16, ) if sampled_grid is None: sampled_height = 0 else: with Image.open(sampled_grid) as image: sampled_height = image.height (export_root / TRAINING_IMAGE_GRID_FILENAME).unlink() full_grid = write_training_image_grid( _export_result(export_root), thumbnail_size=16 ) if full_grid is None: full_height = 0 else: with Image.open(full_grid) as image: full_height = image.height assert sampled_grid == export_root / TRAINING_IMAGE_GRID_FILENAME assert sampled_height > 0 assert sampled_height < full_height def test_training_image_grid_skips_dry_run(tmp_path: Path) -> None: export_root = tmp_path / "export" _write_subset_image(export_root, "fullbody", "red") grid_path = write_training_image_grid(_export_result(export_root, dry_run=True)) assert grid_path is None assert not (export_root / TRAINING_IMAGE_GRID_FILENAME).exists() @pytest.mark.parametrize( "subset_name", ["../../outside", "/absolute", "nested/subset", " padded ", ""], ) def test_existing_export_result_rejects_non_direct_subset_names( tmp_path: Path, subset_name: str, ) -> None: export_root = tmp_path / "export" export_root.mkdir() with pytest.raises(ValueError, match="direct directory names"): build_existing_training_export_result( export_root, subset_names=[subset_name], ) def test_existing_export_result_rejects_duplicate_subset_names(tmp_path: Path) -> None: export_root = tmp_path / "export" _write_subset_image(export_root, "fullbody", "red") with pytest.raises(ValueError, match="Duplicate training subset"): build_existing_training_export_result( export_root, subset_names=["fullbody", "fullbody"], ) def test_existing_export_result_rejects_symlinked_subset(tmp_path: Path) -> None: export_root = tmp_path / "export" export_root.mkdir() outside = tmp_path / "outside" _write_subset_image(outside, "images", "red") (export_root / "linked").symlink_to(outside / "images", target_is_directory=True) with pytest.raises(ValueError, match="must not be symlinks"): build_existing_training_export_result(export_root) def test_existing_export_result_rejects_symlinked_image(tmp_path: Path) -> None: export_root = tmp_path / "export" subset_dir = export_root / "fullbody" subset_dir.mkdir(parents=True) outside_image = tmp_path / "outside.png" Image.new("RGB", (16, 16), color="red").save(outside_image) (subset_dir / "linked.png").symlink_to(outside_image) with pytest.raises(ValueError, match="Exported images must not be symlinks"): build_existing_training_export_result(export_root) def test_training_image_grid_rejects_symlinked_output(tmp_path: Path) -> None: export_root = tmp_path / "export" _write_subset_image(export_root, "fullbody", "red") outside = tmp_path / "outside-grid.jpg" outside.write_text("keep", encoding="utf-8") output_path = export_root / TRAINING_IMAGE_GRID_FILENAME output_path.symlink_to(outside) with pytest.raises(ValueError, match="Contact sheet output must not cross"): write_training_image_grid(_export_result(export_root), thumbnail_size=16) assert outside.read_text(encoding="utf-8") == "keep"