"""Small checks for mistakes that would invalidate training or evaluation.""" import json import tempfile import unittest from pathlib import Path import torch from PIL import Image from data import ( FolderDataset, GridDataset, SupervisedGrids, prepare, prepare_jsonl, prepare_folders, split_records, target_name, transform, ) from models import GridModel, select_target, transfer_backbone from train import ( auxiliary_records, choose_threshold, grid_logits, metrics, partial_loss, thresholds_for, ) class TrainingTests(unittest.TestCase): def test_jsonl_dataset_preserves_splits_and_loads_images(self): with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) / "dataset" for grid in ("3x3", "4x4"): for split in ("train", "validation", "test", "review", "excluded"): folder = root / "data" / grid / split (folder / "images").mkdir(parents=True) Image.new("RGB", (30, 30), "red").save(folder / "images" / "grid.jpg") row = {"id": f"{grid}-{split}", "grid": grid, "target": "buses", "file_name": "images/grid.jpg", "correct_answers": [0], "uncertain_cells": [1], "machine_answers": [2]} (folder / "metadata.jsonl").write_text(json.dumps(row) + "\n") manifest, summary = prepare_jsonl(root, Path(temporary) / "output") self.assertEqual(len(manifest["records"]), 8) self.assertEqual(summary["splits"]["excluded"], 2) record = next(r for r in manifest["records"] if r["id"] == "3x3-validation") self.assertEqual((record["split"], record["target"], record["uncertain"]), ("val", "bus", [1])) sample = GridDataset(manifest["root"], [record], ["bus"], 16)[0] self.assertEqual(tuple(sample[0].shape), (9, 3, 16, 16)) review = next(r for r in manifest["records"] if r["id"] == "4x4-review") self.assertEqual(review["labeler_answers"], [2]) def test_enlarged_views_apply_only_to_requested_targets(self): class MeanModel(torch.nn.Module): tiles = 9 crop_target_ids = [0] def forward(self, images): mean = images.mean((1, 2, 3)) return torch.stack([mean, mean], dim=1) images = torch.zeros(2, 1, 3, 8, 8) images[:, :, :, :4, :4] = 1 output = grid_logits(MeanModel(), images, torch.tensor([0, 1])) self.assertTrue(torch.equal(output, torch.tensor([[1.0], [0.25]]))) def test_auxiliary_training_and_validation_stay_separate(self): records = [ {"id": str(i), "tiles": 16, "target": "motorcycle", "split": split} for i, split in enumerate(["train", "val", "test"]) ] self.assertEqual( [r["id"] for r in auxiliary_records({"records": records}, ["motorcycle"], "train")], ["0"], ) self.assertEqual(auxiliary_records({"records": records}, ["bus"], "train"), []) def test_auxiliary_crops_and_full_image_labels(self): with tempfile.TemporaryDirectory() as root: directory = Path(root) / "example" directory.mkdir() for name in ["full_images.png"] + [f"tile_{i}.png" for i in range(16)]: Image.new("RGB", (16, 16), "blue").save(directory / name) records = [{"id": "example", "tiles": 16, "target": "motorcycle", "answers": [0, 15]}] training = GridDataset( root, records, ["motorcycle"], 16, tiles=True, include_full_images=True ) self.assertEqual(len(training), 17) self.assertEqual([training[i][2].item() for i in [0, 1, 15, 16]], [1, 0, 1, 1]) validation = GridDataset(root, records, ["motorcycle"], 16, tile_grids=True) self.assertEqual(tuple(validation[0][0].shape), (16, 3, 16, 16)) self.assertEqual(validation[0][2].sum().item(), 2) def test_per_target_cutoffs_and_legacy_checkpoint(self): checkpoint = { "classes": ["bus", "motorcycle"], "threshold": 0.5, "thresholds": {"motorcycle": 0.3}, } cutoffs = thresholds_for(checkpoint, torch.tensor([0, 1])) result = metrics(torch.tensor([[0.4], [0.4]]), torch.tensor([[False], [True]]), cutoffs) self.assertEqual(result["exact_grid_accuracy"], 1.0) checkpoint.pop("thresholds") self.assertTrue( torch.equal( thresholds_for(checkpoint, torch.tensor([0, 1])), torch.tensor([[0.5], [0.5]]) ) ) def test_conflicting_duplicates_are_excluded_without_editing_labels(self): with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) / "dataset" for i in range(3): directory = root / str(i) directory.mkdir(parents=True) info = {"keyword": "buses", "correct_answers": [i]} (directory / "info.json").write_text(json.dumps(info)) Image.new("RGB", (30, 30), "red" if i < 2 else "blue").save( directory / "full_images.png" ) for tile in range(9): Image.new("RGB", (10, 10), (i * 70, tile * 20, 0)).save( directory / f"tile_{tile}.png" ) manifest, _ = prepare(root, Path(temporary) / "output") self.assertEqual([r["id"] for r in manifest["records"]], ["2"]) self.assertEqual(len(manifest["excluded_conflicting_labels"]), 2) self.assertEqual( json.loads((root / "0" / "info.json").read_text())["correct_answers"], [0] ) (root / "2" / "info.json").write_text( json.dumps({"keyword": "bus", "correct_answers": [9]}) ) with self.assertRaisesRegex(ValueError, "Invalid answers"): prepare(root, Path(temporary) / "output") def test_transitive_duplicate_groups_never_leak(self): records = [ {"id": str(i), "tiles": 9, "target": "bus", "hashes": [str(i)]} for i in range(30) ] records[0]["hashes"].append("shared-a") records[1]["hashes"].extend(["shared-a", "shared-b"]) records[2]["hashes"].append("shared-b") split_records(records, 42) self.assertEqual(len({r["split"] for r in records[:3]}), 1) self.assertEqual({r["split"] for r in records}, {"train", "val", "test"}) before = [r["split"] for r in records] split_records(records, 42) self.assertEqual(before, [r["split"] for r in records]) def test_singleton_is_training_only(self): records = [{"tiles": 16, "target": "tractor", "hashes": ["x"]}] self.assertEqual(split_records(records, 42)[0]["split"], "train") def test_aliases(self): self.assertEqual(target_name("buses"), target_name("bus")) self.assertEqual(target_name("a fire hydrant"), target_name("fire hydrants")) def test_unknown_targets_have_no_gradient(self): logits = torch.zeros(2, 3, requires_grad=True) loss = torch.nn.functional.binary_cross_entropy_with_logits( select_target(logits, torch.tensor([0, 2])), torch.tensor([1.0, 0.0]) ) loss.backward() self.assertTrue( torch.equal( logits.grad != 0, torch.tensor([[True, False, False], [False, False, True]]) ) ) def test_exact_grid_requires_every_cell_correct(self): probabilities = torch.tensor([[0.9, 0.1], [0.9, 0.8]]) truth = torch.tensor([[True, False], [True, False]]) result = metrics(probabilities, truth, 0.5) self.assertEqual(result["exact_grid_accuracy"], 0.5) self.assertEqual(result["cell_accuracy"], 0.75) self.assertGreater(choose_threshold(probabilities, truth), 0.8) def test_row_major_target_selection(self): logits = torch.arange(2 * 3 * 4 * 4).reshape(2, 3, 4, 4) selected = select_target(logits.flatten(2), torch.tensor([1, 2])) self.assertEqual(selected[0].tolist(), list(range(16, 32))) self.assertEqual(selected[1].tolist(), list(range(80, 96))) def test_models_have_expected_shapes(self): torch.set_num_threads(2) with torch.no_grad(): for tiles, size in [(9, 224), (16, 384)]: model = GridModel(tiles, 3, pretrained=False).eval() shape = (1, 9, 3, size, size) if tiles == 9 else (1, 3, size, size) self.assertEqual( tuple(grid_logits(model, torch.zeros(shape), torch.tensor([1])).shape), (1, tiles), ) if tiles == 9: self.assertEqual( tuple( grid_logits( model, torch.zeros(1, 16, 3, size, size), torch.tensor([1]) ).shape ), (1, 16), ) class FolderDataTests(unittest.TestCase): def test_folder_names_map_to_targets(self): self.assertEqual(target_name("Hydrant"), "fire hydrant") self.assertEqual(target_name("Traffic Light"), "traffic light") self.assertEqual(target_name("Stair"), "stairs") self.assertEqual(target_name("Palm"), target_name("palm trees")) def test_folders_merge_copies_group_near_copies_and_drop_old_heldout(self): with tempfile.TemporaryDirectory() as temporary: temporary = Path(temporary) old = temporary / "dataset" / "grid" old.mkdir(parents=True) held_out = Image.linear_gradient("L").convert("RGB").resize((64, 64)) for name in ["full_images.png"] + [f"tile_{i}.png" for i in range(9)]: held_out.save(old / name) (temporary / "split.json").write_text( json.dumps( { "root": str(temporary / "dataset"), "records": [{"id": "grid", "tiles": 9, "split": "test"}], } ) ) root = temporary / "folders" for split in ("train", "val"): for name in ("Bus", "Car", "Other"): (root / split / name).mkdir(parents=True) Image.new("RGB", (100, 100), "red").save(root / "train" / "Bus" / "a.png") Image.new("RGB", (100, 100), "red").save(root / "val" / "Bus" / "copy.png") Image.new("RGB", (100, 100), "red").save(root / "train" / "Car" / "also.png") held_out.resize((120, 120)).rotate(180).rotate(180).save( root / "train" / "Car" / "old.png" ) for i in range(20): image = Image.new("RGB", (100, 100), (0, 0, 0)) image.paste((255, 255, 255), (0, 0, 5 * (i + 1), 100)) image.save(root / "train" / "Car" / f"{i}.png") image.resize((101, 101)).save(root / "val" / "Car" / f"near{i}.png") Image.new("RGB", (450, 450), "blue").save(root / "train" / "Other" / "full.png") manifest, summary = prepare_folders( root, temporary / "split.json", temporary / "out", workers=2 ) self.assertEqual(summary["exact_duplicates_merged"], 2) self.assertEqual(manifest["dropped_matching_old_heldout"], ["train/Car/old.png"]) merged = [r for r in manifest["records"] if len(r["targets"]) > 1] self.assertEqual([r["targets"] for r in merged], [["bus", "car"]]) splits = {} for record in manifest["records"]: splits.setdefault(record["hashes"][1], set()).add(record["split"]) self.assertTrue(all(len(v) == 1 for v in splits.values())) self.assertEqual({r["split"] for r in manifest["records"]}, {"train", "val"}) self.assertEqual( [r["kind"] for r in manifest["records"] if r["targets"] == ["other"]], ["full"] ) def test_three_folder_layouts_and_glued_yolo_lines(self): with tempfile.TemporaryDirectory() as temporary: temporary = Path(temporary) (temporary / "split.json").write_text( json.dumps({"root": str(temporary), "records": []}) ) images = { "a/train/Bus/1.png": (10, 0, 0), "a/val/TLight/2.png": (0, 10, 0), "b/Stair/3.png": (0, 0, 10), "b/Bus/1.png": (10, 0, 0), "c/images/Chimney/4.png": (50, 0, 0), "c/images/Stair/5.png": (0, 50, 0), "c/images/Stair/6.png": (0, 0, 50), } for name, color in images.items(): (temporary / name).parent.mkdir(parents=True, exist_ok=True) Image.new("RGB", (100, 100), color).save(temporary / name) labels = temporary / "c" / "labels" (labels / "Chimney").mkdir(parents=True) (labels / "Stair").mkdir(parents=True) (labels / "Chimney" / "4.txt").write_text( "1 0.500000 0.500000 0.100000 0.1000002 0.300000 0.300000 0.200000 0.200000\n" ) (labels / "Stair" / "5.txt").write_text("2 0.500000 0.500000 0.100000 0.100000\n") (labels / "Stair" / "6.txt").write_text("2 0.400000 0.500000 0.100000 0.100000\n") manifest, summary = prepare_folders( [temporary / "a", temporary / "b", temporary / "c"], temporary / "split.json", temporary / "out", workers=2, ) targets = {r["path"]: r["targets"] for r in manifest["records"]} self.assertEqual(targets["a/val/TLight/2.png"], ["traffic light"]) self.assertEqual(targets["c/images/Chimney/4.png"], ["chimney", "stairs"]) self.assertEqual(targets["c/images/Stair/6.png"], ["stairs"]) self.assertEqual(summary["yolo_class_ids"], {"1": "chimney", "2": "stairs"}) self.assertEqual(summary["exact_duplicates_merged"], 1) self.assertEqual(summary["unique_images_by_first_source"], {"a": 2, "b": 1, "c": 3}) self.assertTrue(all((temporary / path).exists() for path in targets)) def test_folder_labels_are_weak_negatives_except_background(self): with tempfile.TemporaryDirectory() as root: Image.new("RGB", (8, 8)).save(Path(root) / "x.png") records = [ {"path": "x.png", "targets": ["bus"]}, {"path": "x.png", "targets": ["other"]}, ] dataset = FolderDataset(root, records, ["bus", "car"], 8, cells=16, negative_weight=0.2) _, _, cell_weights, labels, weights = dataset[0] self.assertEqual(labels.tolist(), [1, 0]) self.assertEqual([round(w, 3) for w in weights.tolist()], [1, 0.2]) self.assertEqual(cell_weights.sum().item(), 0) self.assertEqual(tuple(cell_weights.shape), (2, 16)) _, _, _, labels, weights = dataset[1] self.assertEqual((labels.tolist(), weights.tolist()), ([0, 0], [1, 1])) def test_answer_key_items_supervise_only_the_queried_target(self): with tempfile.TemporaryDirectory() as root: directory = Path(root) / "example" directory.mkdir() for name in ["full_images.png"] + [f"tile_{i}.png" for i in range(16)]: Image.new("RGB", (16, 16), "blue").save(directory / name) records = [{"id": "example", "tiles": 16, "target": "car", "answers": [0, 15]}] grids = SupervisedGrids(GridDataset(root, records, ["bus", "car"], 16), 16) _, labels, weights, _, any_weights = grids[0] self.assertEqual(weights.sum(1).tolist(), [0, 16]) self.assertEqual(labels[1].nonzero().flatten().tolist(), [0, 15]) self.assertEqual(any_weights.sum().item(), 0) tiles = SupervisedGrids(GridDataset(root, records, ["bus", "car"], 16, tiles=True), 1) self.assertEqual( (tiles[15][1].tolist(), tiles[15][2].tolist()), ([[0], [1]], [[0], [1]]) ) def test_partial_loss_matches_old_loss_and_ignores_unknowns(self): logits = torch.randn(2, 3, 16, requires_grad=True) labels, weights = torch.zeros(2, 3, 16), torch.zeros(2, 3, 16) labels[0, 1, :4], weights[0, 1], weights[1, 2] = 1, 1, 1 zero = torch.zeros(2, 3) loss, count = partial_loss(logits, labels, weights, zero, zero) expected = torch.nn.functional.binary_cross_entropy_with_logits( select_target(logits, torch.tensor([1, 2])), labels[[0, 1], [1, 2]] ) self.assertTrue(torch.allclose(loss, expected)) self.assertEqual(count.item(), 32) loss.backward() self.assertTrue( torch.equal( (logits.grad != 0).any(-1), torch.tensor([[False, True, False], [False, False, True]]), ) ) def test_image_level_label_trains_only_the_strongest_cell(self): logits = torch.zeros(1, 2, 16) logits[0, 0, 5] = 2 logits.requires_grad_() empty = torch.zeros(1, 2, 16) loss, _ = partial_loss( logits, empty, empty, torch.tensor([[1.0, 0.0]]), torch.tensor([[1.0, 0.0]]) ) loss.backward() self.assertEqual(logits.grad.nonzero().tolist(), [[0, 0, 5]]) def test_tile_backbone_initializes_grid_model(self): tile_model, grid_model = ( GridModel(9, 3, pretrained=False), GridModel(16, 5, pretrained=False), ) transfer_backbone(grid_model, tile_model.state_dict()) for name, value in grid_model.backbone.state_dict().items(): self.assertTrue(torch.equal(value, tile_model.backbone.state_dict()[name]), name) def test_torch_free_preprocessing_matches_training(self): try: from predict_onnx import preprocess except ImportError: self.skipTest("onnxruntime is not installed (uv sync --group export)") image = Image.effect_mandelbrot((97, 97), (-2, -1.5, 1, 1.5), 50).convert("RGB") expected = transform(64)(image).numpy() self.assertLess(abs(preprocess(image, 64) - expected).max(), 1e-5) class LabeledLayoutTests(unittest.TestCase): def test_invalid_metadata_and_sam_files_are_skipped(self): from data import scan with tempfile.TemporaryDirectory() as root: self.make(root, "valid", "3x3", [0]) self.make(root, "bad-target", "3x3", [0]) self.make(root, "bad-metadata", "3x3", [0]) self.make(root, "bad-sam", "4x4", [0]) (Path(root) / "bad-target" / "metadata.json").write_text( json.dumps({"target": None, "grid": "3x3", "correct_answers": [0]}) ) (Path(root) / "bad-metadata" / "metadata.json").write_text("[]") (Path(root) / "bad-sam" / "sam.json").write_text("{}") records, skipped = scan(root, workers=2, sam=True) self.assertEqual([record["id"] for record in records], ["valid"]) self.assertEqual( skipped, { "invalid target": 1, "invalid metadata.json": 1, "invalid sam.json": 1, }, ) def test_folder_preparation_excludes_tiles_from_labeled_holdouts(self): from data import grid_boxes, read_rgb, scan with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) data = root / "data" data.mkdir() self.make(data, "holdout", "3x3", [0]) records, _ = scan(data, workers=2) records[0]["split"] = "test" manifest = root / "split.json" manifest.write_text(json.dumps({"root": str(data), "records": records})) folder = root / "folders" / "bus" folder.mkdir(parents=True) image = read_rgb(data / "holdout" / "image.jpg") image.crop(grid_boxes(*image.size, 9)[0]).save(folder / "reused.png") Image.effect_mandelbrot((100, 100), (-2, -1.5, 1, 1.5), 50).save(folder / "new.png") prepared, summary = prepare_folders( root / "folders", manifest, root / "out", workers=2, ) self.assertEqual(summary["dropped_matching_old_heldout"], 1) self.assertEqual([record["path"] for record in prepared["records"]], ["bus/new.png"]) def make(self, root, name, grid, answers=None, uncertain=(), colors=None, size=300): folder = Path(root) / name folder.mkdir() n = 3 if grid == "3x3" else 4 image = Image.new("RGB", (size, size)) for i, color in enumerate(colors or [(i * 20, 0, 0) for i in range(n * n)]): r, c = divmod(i, n) image.paste( color, (c * size // n, r * size // n, (c + 1) * size // n, (r + 1) * size // n) ) image.save(folder / "image.jpg", quality=100) info = {"target": "buses", "grid": grid} if answers is not None: info.update(correct_answers=answers, uncertain_cells=list(uncertain)) (folder / "metadata.json").write_text(json.dumps(info)) def test_labeled_items_are_scanned_and_unlabeled_skipped(self): from data import scan with tempfile.TemporaryDirectory() as root: self.make(root, "a", "3x3", [1, 4], [4]) self.make(root, "b", "4x4", [5], size=450) self.make(root, "c", "3x3") records, skipped = scan(root, workers=2) self.assertEqual( [(r["id"], r["tiles"], r["target"], r["answers"], r["uncertain"]) for r in records], [("a", 9, "bus", [1, 4], [4]), ("b", 16, "bus", [5], [])], ) self.assertEqual(skipped, {"not labeled": 1}) self.assertEqual(len(records[0]["hashes"]), 10) def test_tiles_come_from_the_single_image_and_uncertain_cells_are_ignored(self): from data import scan with tempfile.TemporaryDirectory() as root: colors = [(250, 0, 0), (0, 250, 0), (0, 0, 250)] + [(0, 0, 0)] * 6 self.make(root, "a", "3x3", [1], [2], colors=colors) self.make(root, "b", "4x4", [5, 6], [6], size=450) records, _ = scan(root, workers=2) tiles = GridDataset(root, [records[0]], ["bus"], 8, tiles=True) self.assertEqual(len(tiles), 9) self.assertGreater( tiles[1][0][1].mean().item(), tiles[1][0][0].mean().item() ) # green tile items = SupervisedGrids(tiles, 1) self.assertEqual(items[1][1].tolist(), [[1.0]]) self.assertEqual(items[2][2].sum().item(), 0) # uncertain tile: no weight self.assertEqual(items[0][2].sum().item(), 1) grid = SupervisedGrids(GridDataset(root, [records[1]], ["bus"], 32), 16) _, labels, weights, _, _ = grid[0] self.assertEqual(labels[0].nonzero().flatten().tolist(), [5, 6]) self.assertEqual(weights[0].tolist(), [1.0] * 6 + [0.0] + [1.0] * 9) kept = SupervisedGrids( GridDataset(root, [records[1]], ["bus"], 32), 16, keep_uncertain=True ) self.assertEqual(kept[0][2].sum().item(), 16) def test_sam_labels_replace_4x4_answers_only(self): from data import scan with tempfile.TemporaryDirectory() as root: self.make(root, "a", "3x3", [1, 4], [4]) self.make(root, "b", "4x4", [5], [6], size=450) self.make(root, "c", "4x4", size=450) # Qwen never labeled it; SAM did self.make(root, "d", "4x4", [2], size=450) # no sam.json for name, cells in (("b", [5, 6, 9]), ("c", [0])): (Path(root) / name / "sam.json").write_text( json.dumps({"model": "facebook/sam3", "correct_answers": cells}) ) records, skipped = scan(root, workers=2, sam=True) self.assertEqual( [(r["id"], r["answers"], r["uncertain"], r["label_source"]) for r in records], [ ("a", [1, 4], [4], None), ("b", [5, 6, 9], [], "facebook/sam3:sam"), ("c", [0], [], "facebook/sam3:sam"), ], ) self.assertEqual(skipped, {"no sam.json": 1}) records, _ = scan(root, workers=2) self.assertEqual([r["answers"] for r in records], [[1, 4], [5], [2]]) def test_reviewed_grids_are_held_out_with_their_reviewed_answers(self): from data import prepare with tempfile.TemporaryDirectory() as root, tempfile.TemporaryDirectory() as out: for i in range(20): # g01 is a copy of the reviewed g00, so it must not reach training. j = 0 if i == 1 else i colors = [((j * 37 + k * 11) % 256, (k * 23) % 256, j % 256) for k in range(9)] self.make(root, f"g{i:02d}", "3x3", [0], colors=colors) key = Path(out) / "key.json" key.write_text(json.dumps({"items": {"g00": {"answers": [0, 8]}}})) manifest, counts = prepare(root, out, workers=2, review_keys=[key]) by_id = {r["id"]: r for r in manifest["records"]} self.assertEqual( (by_id["g00"]["split"], by_id["g00"]["answers"], by_id["g00"]["labeler_answers"]), ("review", [0, 8], [0]), ) self.assertNotIn("g01", by_id) self.assertEqual(counts["dropped_sharing_images_with_review"], 1) self.assertNotIn( "review", {r["split"] for r in manifest["records"] if r["id"] != "g00"} ) def test_flip_mirrors_4x4_cells_and_uncertain_cells(self): import random from data import scan with tempfile.TemporaryDirectory() as root: colors = [(250, 0, 0)] + [(0, 0, 0)] * 15 self.make(root, "b", "4x4", [0, 5], [1], colors=colors, size=448) records, _ = scan(root, workers=2) grids = GridDataset(root, records, ["bus"], 32) plain = SupervisedGrids(grids, 16)[0] seed = next( s for s in range(50) if random.Random(s).random() < 0.5 ) # a seed whose first draw flips random.seed(seed) flipped = SupervisedGrids(grids, 16, flip=True)[0] _, labels, weights, _, _ = flipped self.assertEqual(labels[0].nonzero().flatten().tolist(), [3, 6]) self.assertEqual( weights[0].tolist(), [1.0] * 2 + [0.0] + [1.0] * 13 ) # uncertain 1 -> 2 self.assertTrue(torch.equal(flipped[0], plain[0].flip(-1))) def test_reused_tiles_keep_grids_in_one_split(self): from data import scan with tempfile.TemporaryDirectory() as root: for i in range(40): colors = [((i * 37 + k * 11) % 256, (k * 23) % 256, i % 256) for k in range(9)] if i in (0, 1): colors[4] = (255, 255, 255) # the same photo reused in two grids self.make(root, f"g{i:02d}", "3x3", [0], colors=colors) records, _ = scan(root, workers=2) split_records(records, 42) self.assertEqual(records[0]["split"], records[1]["split"]) class LayoutTests(unittest.TestCase): def test_folder_metadata_takes_priority_over_image_size(self): from train import layout with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) Image.new("RGB", (300, 300)).save(root / "image.jpg") (root / "metadata.json").write_text(json.dumps({"grid": "4x4"})) self.assertEqual(layout(root), 16) def test_layout_from_size_or_seams(self): from train import layout with tempfile.TemporaryDirectory() as root: root = Path(root) Image.linear_gradient("L").convert("RGB").resize((300, 300)).save(root / "a.png") Image.linear_gradient("L").convert("RGB").resize((450, 450)).save(root / "b.png") tiles = Image.new("RGB", (600, 600)) for i in range(9): tiles.paste( (i * 28, 255 - i * 28, (i * 90) % 256), ((i % 3) * 200, (i // 3) * 200, (i % 3 + 1) * 200, (i // 3 + 1) * 200), ) tiles.save(root / "c.png") Image.linear_gradient("L").convert("RGB").resize((600, 600)).save(root / "d.png") self.assertEqual([layout(root / f"{n}.png") for n in "abcd"], [9, 16, 9, 16]) if __name__ == "__main__": unittest.main()