File size: 3,034 Bytes
16760fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | from __future__ import annotations
import numpy as np
from toaster.core import LabelSchema
def test_colors_for_is_vectorized(schema):
colors = schema.colors_for(np.array([1, 2, 0]))
assert colors.tolist() == [[255, 0, 0], [0, 255, 0], [0, 0, 0]]
assert colors.dtype == np.uint8
def test_out_of_range_labels_clamp_to_unlabeled(schema):
colors = schema.colors_for(np.array([99, -1]))
assert colors.tolist() == [[0, 0, 0], [0, 0, 0]]
def test_from_config_parses_hex_and_names():
cfg = {
"ignore_index": 0,
"color_map": {0: "#000000", 1: "#ff0000", 2: [0, 0, 255]},
"semantic_map": {1: "car", 2: "road"},
}
s = LabelSchema.from_config(cfg)
assert s.unlabeled_id == 0
assert s.get(1).name == "car"
assert s.get(1).color == (255, 0, 0)
assert s.get(2).color == (0, 0, 255)
def test_set_color_updates_lut(schema):
assert schema.colors_for(np.array([1]))[0].tolist() == [255, 0, 0]
schema.set_color(1, (10, 20, 30))
assert schema.get(1).color == (10, 20, 30)
# The cached LUT is invalidated, so colours_for reflects the change.
assert schema.colors_for(np.array([1]))[0].tolist() == [10, 20, 30]
def test_yaml_round_trip(tmp_path):
import yaml
cfg = {"ignore_index": 0, "color_map": {0: "#010203", 1: "#0a0b0c"}}
path = tmp_path / "schema.yaml"
path.write_text(yaml.safe_dump(cfg))
s = LabelSchema.from_yaml(path)
assert s.get(1).color == (10, 11, 12)
def test_add_class_assigns_next_id_and_colours(schema):
cls = schema.add_class("tree", color=(1, 2, 3))
assert cls.id == 3 # one past the previous max (2)
assert schema.get(3).name == "tree"
# The LUT was invalidated, so the new class colours immediately.
assert schema.colors_for(np.array([3]))[0].tolist() == [1, 2, 3]
def test_add_class_auto_colour_is_unused(schema):
a = schema.add_class("x")
b = schema.add_class("y")
assert a.color != b.color # auto-assigned colours differ
assert a.color not in {(0, 0, 0), (255, 0, 0), (0, 255, 0)} # avoids existing ones
def test_rename_class(schema):
schema.rename(1, "car")
assert schema.get(1).name == "car"
assert schema.get(1).color == (255, 0, 0) # colour untouched
def test_remove_class(schema):
schema.remove(1)
assert 1 not in schema
assert [c.id for c in schema.classes] == [0, 2]
# A now-orphaned label clamps to unlabeled rather than crashing.
assert schema.colors_for(np.array([1]))[0].tolist() == [0, 0, 0]
def test_remove_unlabeled_is_refused(schema):
import pytest
with pytest.raises(ValueError, match="unlabeled"):
schema.remove(0)
def test_to_yaml_round_trip(schema, tmp_path):
schema.add_class("tree", color=(1, 2, 3))
out = schema.to_yaml(tmp_path / "out.yaml")
reloaded = LabelSchema.from_yaml(out)
assert reloaded.unlabeled_id == schema.unlabeled_id
assert {c.id: (c.name, c.color) for c in reloaded.classes} == {
c.id: (c.name, c.color) for c in schema.classes
}
|