11"""Models for the labels types."""
22
33from enum import Enum
4+ from typing import Tuple
45
56
67class DocItemLabel (str , Enum ):
78 """DocItemLabel."""
89
9- # DocLayNet v2
1010 CAPTION = "caption"
1111 FOOTNOTE = "footnote"
1212 FORMULA = "formula"
@@ -26,12 +26,34 @@ class DocItemLabel(str, Enum):
2626 KEY_VALUE_REGION = "key_value_region"
2727
2828 # Additional labels for markup-based formats (e.g. HTML, Word)
29- PARAGRAPH = "paragraph" # explicitly a paragraph and not arbitrary text
29+ PARAGRAPH = "paragraph"
3030 REFERENCE = "reference"
3131
32- def __str__ (self ):
33- """Get string value."""
34- return str (self .value )
32+ @staticmethod
33+ def get_color (label : "DocItemLabel" ) -> Tuple [int , int , int ]:
34+ """Return the RGB color associated with a given label."""
35+ color_map = {
36+ DocItemLabel .CAPTION : (255 , 204 , 153 ),
37+ DocItemLabel .FOOTNOTE : (200 , 200 , 255 ),
38+ DocItemLabel .FORMULA : (192 , 192 , 192 ),
39+ DocItemLabel .LIST_ITEM : (153 , 153 , 255 ),
40+ DocItemLabel .PAGE_FOOTER : (204 , 255 , 204 ),
41+ DocItemLabel .PAGE_HEADER : (204 , 255 , 204 ),
42+ DocItemLabel .PICTURE : (255 , 204 , 164 ),
43+ DocItemLabel .SECTION_HEADER : (255 , 153 , 153 ),
44+ DocItemLabel .TABLE : (255 , 204 , 204 ),
45+ DocItemLabel .TEXT : (255 , 255 , 153 ),
46+ DocItemLabel .TITLE : (255 , 153 , 153 ),
47+ DocItemLabel .DOCUMENT_INDEX : (220 , 220 , 220 ),
48+ DocItemLabel .CODE : (125 , 125 , 125 ),
49+ DocItemLabel .CHECKBOX_SELECTED : (255 , 182 , 193 ),
50+ DocItemLabel .CHECKBOX_UNSELECTED : (255 , 182 , 193 ),
51+ DocItemLabel .FORM : (200 , 255 , 255 ),
52+ DocItemLabel .KEY_VALUE_REGION : (183 , 65 , 14 ),
53+ DocItemLabel .PARAGRAPH : (255 , 255 , 153 ),
54+ DocItemLabel .REFERENCE : (176 , 224 , 230 ),
55+ }
56+ return color_map [label ]
3557
3658
3759class GroupLabel (str , Enum ):
0 commit comments