feat: train fixed-size classifiers and export class probabilities

tests/test_classifier_export.py

@@ -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)