File size: 610 Bytes
201b13c | 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 | from typing import Dict
# Forward mapping: model class ID → defect name
CLASS_MAP: Dict[int, str] = {
0: "crazing",
1: "patches",
2: "rolled_in_scale",
3: "pitted_surface",
}
# Reverse mapping: defect name → class ID
REVERSE_CLASS_MAP: Dict[str, int] = {v: k for k, v in CLASS_MAP.items()}
def get_class_name(class_id: int) -> str:
"""
Safely get class name from class ID.
"""
return CLASS_MAP.get(class_id, "unknown")
def get_class_id(class_name: str) -> int:
"""
Safely get class ID from class name.
"""
return REVERSE_CLASS_MAP.get(class_name, -1) |