feat: train fixed-size classifiers and export class probabilities

2026-10-14 17:08:14 +08:00
parent 63bab36147
commit f9a4032783
11 changed files with 478 additions and 208 deletions
@@ -1,83 +0,0 @@
-import pytest
-import torch
-from torch import nn
-
-from scripts.export_classifier_openvino import YoloClassifierAdapter, class_order, preprocessing
-
-
-class FixedClassifier(nn.Module):
-    def forward(self, x):
-        return torch.tensor([1.0, 3.0, -1.0]).expand(x.shape[0], -1)
-
-
-def test_yolo_contract_and_class_reordering():
-    adapter = YoloClassifierAdapter(FixedClassifier(), 128, 336, [2, 0, 1])
-    result = adapter(torch.rand(1, 3, 128, 128))
-    assert result.shape == (1, 7, 336)
-    torch.testing.assert_close(result[0, :4, 0], torch.tensor([64., 64., 128., 128.]))
-    expected = torch.tensor([1., 3., -1.]).softmax(0)[[2, 0, 1]]
-    torch.testing.assert_close(result[0, 4:, 0], expected)
-    assert result[0, 4:, 0].argmax().item() == 2
-    assert torch.count_nonzero(result[:, :, 1:]) == 0
-
-
-def test_normalization_matches_training_for_each_image():
-    adapter = YoloClassifierAdapter(FixedClassifier(), 128, 1, [0, 1, 2])
-    images = torch.stack((torch.rand(3, 128, 128), torch.ones(3, 128, 128)))
-    expected = torch.stack([(x - x.mean()) / (x.std() + 1e-5) for x in images])
-    torch.testing.assert_close(adapter.normalize(images), expected)
-    assert torch.isfinite(adapter(images)).all()
-
-
-def test_class_sets_must_match():
-    assert class_order(["normal", "mosquito"], ["mosquito", "normal"]) == [1, 0]
-    with pytest.raises(ValueError, match="Class sets differ"):
-        class_order(["normal", "mosquito"], ["normal", "blackdot"])
-    with pytest.raises(ValueError, match="Duplicate"):
-        class_order(["normal", "normal"], ["normal"])
-
-
-def test_background_suppression_uses_reordered_output_index():
-    # Source class 1 wins, but it is output class 2 after reordering.
-    adapter = YoloClassifierAdapter(FixedClassifier(), 128, 336, [2, 0, 1], background_index=2)
-    result = adapter(torch.rand(1, 3, 128, 128))
-    assert result.shape == (1, 7, 336)
-    assert torch.count_nonzero(result[:, 4:, :]) == 0
-    # The existing C++ decoder returns zero candidates at any positive threshold.
-    assert not (result[0, 4:, :].max(dim=0).values >= 0.25).any()
-
-
-def test_foreground_class_indices_and_scores_are_preserved():
-    foreground = YoloClassifierAdapter(FixedClassifier(), 128, 336, [0, 1, 2],
-                                      background_index=0)
-    legacy = YoloClassifierAdapter(FixedClassifier(), 128, 336, [0, 1, 2])
-    x = torch.rand(1, 3, 128, 128)
-    torch.testing.assert_close(foreground(x), legacy(x))
-    assert foreground(x)[0, 4:, 0].argmax() == 1
-
-
-def test_background_gate_is_per_image_in_batch():
-    class BatchClassifier(nn.Module):
-        def forward(self, x):
-            return x[:, :, 0, 0]
-
-    x = torch.zeros(2, 3, 4, 4)
-    x[0, 0, 0, 0] = 3  # normal
-    x[1, 1, 0, 0] = 3  # blackdot
-    adapter = YoloClassifierAdapter(BatchClassifier(), 4, 10, [0, 1, 2],
-                                    eps=None, background_index=0)
-    result = adapter(x)
-    assert not torch.count_nonzero(result[0, 4:, :])
-    assert result[1, 4:, 0].argmax() == 1
-    assert result[1, 5, 0] > 0.9
-
-
-def test_refuses_resampling_and_unknown_preprocessing():
-    config = {"data": {"val_transforms": [
-        {"name": "Resize", "size": [160, 160]}, {"name": "ToTensor"},
-        {"name": "PerImageStandardize"}]}}
-    with pytest.raises(ValueError, match="Training size differs"):
-        preprocessing(config, 128)
-    config["data"]["val_transforms"] = [{"name": "Normalize"}]
-    with pytest.raises(ValueError, match="Unsupported"):
-        preprocessing(config, 128)
@@ -0,0 +1,52 @@
+import pytest
+import torch
+from torch import nn
+
+from scripts.export_classifier_openvino import ClassifierProbabilities, ordered_names, preprocessing
+
+
+class FixedClassifier(nn.Module):
+    def forward(self, x):
+        return torch.tensor([1.0, 3.0, -1.0]).expand(x.shape[0], -1)
+
+
+def test_probabilities_keep_training_class_order_and_sum_to_one():
+    result = ClassifierProbabilities(FixedClassifier())(torch.rand(2, 3, 128, 128))
+    assert result.shape == (2, 3)
+    torch.testing.assert_close(result[0], torch.tensor([1., 3., -1.]).softmax(0))
+    torch.testing.assert_close(result.sum(1), torch.ones(2))
+    assert result.argmax(1).tolist() == [1, 1]
+
+
+def test_normal_probability_is_not_suppressed():
+    # Suppose class 1 is normal: its probability is returned like any other class.
+    result = ClassifierProbabilities(FixedClassifier())(torch.zeros(1, 3, 128, 128))
+    assert result[0, 1] > .8
+    assert (result > 0).all()
+
+
+def test_standardization_matches_training_per_image():
+    exported = ClassifierProbabilities(FixedClassifier(), eps=1e-5)
+    images = torch.stack((torch.rand(3, 128, 128), torch.ones(3, 128, 128)))
+    expected = torch.stack([(x-x.mean())/(x.std()+1e-5) for x in images])
+    torch.testing.assert_close(exported.normalize(images), expected)
+    assert torch.isfinite(exported(images)).all()
+
+
+def test_classes_validate_index_order_and_duplicates():
+    assert ordered_names({1: "blackdot", 0: "normal"}) == ["normal", "blackdot"]
+    with pytest.raises(ValueError, match="Duplicate"):
+        ordered_names(["normal", "normal"])
+    with pytest.raises(ValueError, match="contiguous"):
+        ordered_names({1: "normal"})
+
+
+def test_refuses_resampling_and_unknown_preprocessing():
+    config = {"data": {"val_transforms": [
+        {"name": "Resize", "size": [160, 160]}, {"name": "ToTensor"},
+        {"name": "PerImageStandardize"}]}}
+    with pytest.raises(ValueError, match="Training size differs"):
+        preprocessing(config, 128)
+    config["data"]["val_transforms"] = [{"name": "Normalize"}]
+    with pytest.raises(ValueError, match="Unsupported"):
+        preprocessing(config, 128)
@@ -0,0 +1,26 @@
+import numpy as np
+
+from scripts.compare_openvino_classes import matrix_report, output_class, probability_class
+
+
+def test_detector_output_channel_offset_and_background():
+    output = np.zeros((1, 13, 336), dtype=np.float32)
+    output[0, :4, 0] = [64, 64, 128, 128]
+    assert output_class(output, 0.25, 4) == 4
+    output[0, 5, 0] = 0.6
+    assert output_class(output, 0.25, 4) == 1  # blackdot, not black
+    output[0, 4, 50] = 0.9
+    assert output_class(output, 0.25, 4) == 0
+
+
+def test_report_excludes_background_from_foreground_agreement():
+    matrix = np.array([[1, 1], [0, 8]])
+    report = matrix_report(matrix, ["blackdot", "normal"], 1)
+    assert report["agreement"] == 0.9
+    assert report["foreground_agreement"] == 0.5
+
+
+def test_probability_output_has_no_coordinate_offset():
+    probabilities = np.array([[0.1, 0.8, 0.1]])
+    assert probability_class(probabilities, .25, 2) == 1
+    assert probability_class(np.array([[.3, .4, .3]]), .5, 2) == 2
@@ -0,0 +1,40 @@
+import pytest
+from PIL import Image
+
+from datasets.tile_classification import YoloTileClassificationDataset, tile_label
+from datasets.yolo import write_data_yaml
+
+
+def test_label_uses_human_class_ids_and_empty_background():
+    assert tile_label("", 9, 4) == 4
+    assert tile_label("1 0.5 0.5 0.2 0.2", 9, 4) == 1
+    assert tile_label("1 0.5 0.5 0.2 0.2\n3 0.5 0.5 0.4 0.4", 9, 4) == 3
+    with pytest.raises(ValueError):
+        tile_label("1 0.5 0.5 nan 0.2", 9, 4)
+
+
+def test_existing_split_is_preserved_and_no_resize(tmp_path):
+    for split in ("train", "val"):
+        (tmp_path / "images" / split).mkdir(parents=True)
+        (tmp_path / "labels" / split).mkdir(parents=True)
+        Image.new("RGB", (16, 16)).save(tmp_path / "images" / split / "sample.png")
+        (tmp_path / "labels" / split / "sample.txt").write_text(
+            "1 0.5 0.5 0.2 0.2" if split == "train" else "")
+    yaml = write_data_yaml(tmp_path / "data.yaml", ["normal", "blackdot"], root=tmp_path)
+    kwargs = {"yaml_path": yaml, "image_size": [16, 16], "transforms": [{"name": "ToTensor"}]}
+    train = YoloTileClassificationDataset(split="train", **kwargs)
+    val = YoloTileClassificationDataset(split="val", **kwargs)
+    assert train.class_counts() == [0, 1]
+    assert val.class_counts() == [1, 0]
+    assert train[0]["image"].shape == (3, 16, 16)
+    Image.new("RGB", (12, 16)).save(train.samples[0][0])
+    with pytest.raises(ValueError, match="resize is forbidden"):
+        train[0]
+
+
+def test_missing_labels_are_not_silent_background(tmp_path):
+    (tmp_path / "images" / "train").mkdir(parents=True)
+    Image.new("RGB", (16, 16)).save(tmp_path / "images" / "train" / "sample.png")
+    yaml = write_data_yaml(tmp_path / "data.yaml", ["normal", "blackdot"], root=tmp_path)
+    with pytest.raises(FileNotFoundError, match="Missing label"):
+        YoloTileClassificationDataset(yaml, "train")