| """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() |
| ) |
| items = SupervisedGrids(tiles, 1) |
| self.assertEqual(items[1][1].tolist(), [[1.0]]) |
| self.assertEqual(items[2][2].sum().item(), 0) |
| 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) |
| self.make(root, "d", "4x4", [2], size=450) |
| 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): |
| |
| 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 |
| ) |
| 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 |
| ) |
| 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) |
| 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() |
|
|
