1
1
"""Models for the labels types."""
2
2
3
3
from enum import Enum
4
+ from typing import Tuple
4
5
5
6
6
7
class DocItemLabel (str , Enum ):
7
8
"""DocItemLabel."""
8
9
9
- # DocLayNet v2
10
10
CAPTION = "caption"
11
11
FOOTNOTE = "footnote"
12
12
FORMULA = "formula"
@@ -26,12 +26,34 @@ class DocItemLabel(str, Enum):
26
26
KEY_VALUE_REGION = "key_value_region"
27
27
28
28
# Additional labels for markup-based formats (e.g. HTML, Word)
29
- PARAGRAPH = "paragraph" # explicitly a paragraph and not arbitrary text
29
+ PARAGRAPH = "paragraph"
30
30
REFERENCE = "reference"
31
31
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 ]
35
57
36
58
37
59
class GroupLabel (str , Enum ):
0 commit comments