Skip to content

Commit 639f122

Browse files
authored
feat: Add mapping to colors into DocItemLabel (#123)
Signed-off-by: Christoph Auer <[email protected]>
1 parent 127dd2f commit 639f122

File tree

1 file changed

+27
-5
lines changed

1 file changed

+27
-5
lines changed

docling_core/types/doc/labels.py

+27-5
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
"""Models for the labels types."""
22

33
from enum import Enum
4+
from typing import Tuple
45

56

67
class 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

3759
class GroupLabel(str, Enum):

0 commit comments

Comments
 (0)