def test_segmentation_classes():
seg1 = Segmentation(shape=(3008, 3008), dtype="uint8", num_classes=5)
seg2 = Segmentation(
shape=(3008, 3008), dtype="uint8", names=["apple", "orange", "banana"]
)
assert seg1.get_segmentation_classes() == ["0", "1", "2", "3", "4"]
assert seg2.get_segmentation_classes() == ["apple", "orange", "banana"]
def test_segmentation_repr():
seg1 = Segmentation(shape=(3008, 3008), dtype="uint8", num_classes=5)
seg2 = Segmentation(
shape=(3008, 3008), dtype="uint8", names=["apple", "orange", "banana"]
)
text1 = "Segmentation(shape=(3008, 3008), dtype='uint8', num_classes=5)"
text2 = "Segmentation(shape=(3008, 3008), dtype='uint8', names=['apple', 'orange', 'banana'], num_classes=3)"
assert seg1.__repr__() == text1
assert seg2.__repr__() == text2