53 lines
2.1 KiB
Python
53 lines
2.1 KiB
Python
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)
|