def test_exception():
serial = SegmentationLabels()
with pytest.raises(Exception):
sample = torch.zeros(1, 5, 2, 3)
serial.serialize(sample)
with pytest.raises(Exception):
sample = torch.zeros(2, 3)
serial.serialize(sample)
def test_serialize():
serial = SegmentationLabels()
sample = torch.zeros(5, 2, 3)
sample[1, 1, 2] = 1 # add peak in class 2
sample[3, 0, 1] = 1 # add peak in class 4
classes = serial.serialize({DefaultDataKeys.PREDS: sample})
assert torch.tensor(classes)[1, 2] == 1
assert torch.tensor(classes)[0, 1] == 3