10 dominant colors

anno_page/core/metadata.py

@@ -12,6 +12,18 @@
 logger = logging.getLogger(__name__)
 
 
+class DominantColorInfo:
+    def __init__(self, name, coverage):
+        self.name = name
+        self.coverage = coverage
+
+
+class ColorInfo:
+    def __init__(self, color_mode:str|None=None, dominant_colors:List[DominantColorInfo]|None=None):
+        self.color_mode = color_mode
+        self.dominant_colors = dominant_colors
+
+
 class BaseMetadata:
     def __init__(self, tag_id, mods_id, mods_uuid=None, record_identifier=None):
         self.tag_id = tag_id
@@ -231,7 +243,7 @@ def __init__(self,
                  description: Optional[str| Dict[Language, str]] = None,
                  caption: Optional[str | Dict[Language, str]] = None,
                  topics: Optional[str | Dict[Language, str] | Dict[Language, list[str]]] = None,
-                 color: Optional[str | Dict[Language, str]] = None,
+                 color: Optional[ColorInfo | Dict[Language, ColorInfo]] = None,
                  title: Optional[str | Dict[Language, str]] = None,
                  caption_lines_metadata: Optional[RelatedLinesMetadata] = None,
                  reference_lines_metadata: Optional[RelatedLinesMetadata] = None,
@@ -356,7 +368,7 @@ def _add_size_element(mods, mods_namespace, bounding_box):
         extent.attrib["unit"] = "pixels"
 
     def _add_color_elements(self, mods, mods_namespace):
-        if isinstance(self.color, str):
+        if isinstance(self.color, ColorInfo):
             self._add_color_element(mods, mods_namespace, "", self.color)
         elif isinstance(self.color, dict):
             for language, color in self.color.items():
@@ -366,13 +378,22 @@ def _add_color_elements(self, mods, mods_namespace):
             logger.warning(f"Color is not a string or dictionary in GraphicalObjectMetadata '{self.tag_id}'.")
 
     @staticmethod
-    def _add_color_element(mods, mods_namespace, language, color):
+    def _add_color_element(mods, mods_namespace, language, color_info: ColorInfo):
         physical_description = ET.SubElement(mods, f"{{{mods_namespace}}}physicalDescription")
         physical_description.attrib["altRepGroup"] = "color-1"
-        form = ET.SubElement(physical_description, f"{{{mods_namespace}}}form")
-        form.attrib["type"] = "color"
-        form.attrib["lang"] = language
-        form.text = color
+
+        form_color_mode = ET.SubElement(physical_description, f"{{{mods_namespace}}}form")
+        form_color_mode.attrib["type"] = "color"
+        form_color_mode.attrib["lang"] = language
+        form_color_mode.text = color_info.color_mode
+
+        if color_info.dominant_colors is not None:
+            for dominant_color in color_info.dominant_colors:
+                form_dominant_color = ET.SubElement(physical_description, f"{{{mods_namespace}}}form")
+                form_dominant_color.attrib["type"] = "dominant-color"
+                form_dominant_color.attrib["lang"] = language
+                form_dominant_color.attrib["coverage"] = f"{dominant_color.coverage:.2f}"
+                form_dominant_color.text = dominant_color.name
 
     def _add_caption_elements(self, mods, mods_namespace):
         if isinstance(self.caption, str):

anno_page/core/page_parser.py

@@ -4,7 +4,7 @@
 
 from anno_page.engines import (LayoutProcessingEngine, YoloDetectionEngine, HuggingfaceImageEmbeddingEngine,
                                OpenAICompletionsImageCaptioningEngine, CaptionYoloNearestEngine,
-                               CaptionYoloKeypointsEngine, CaptionYoloOrganizerEngine)
+                               CaptionYoloKeypointsEngine, CaptionYoloOrganizerEngine, DominantColorsEngine)
 
 
 def operation_factory(config, device, config_path) -> LayoutProcessingEngine | None:
@@ -24,6 +24,9 @@ def operation_factory(config, device, config_path) -> LayoutProcessingEngine | N
     elif config['METHOD'] == 'OPENAI_COMPLETIONS_IMAGE_CAPTIONING':
         logger.info("Creating OpenAICompletionsImageCaptioning engine")
         engine = OpenAICompletionsImageCaptioningEngine(config, device, config_path=config_path)
+    elif config['METHOD'] == 'DOMINANT_COLORS':
+        logger.info("Creating DominantColorsEngine engine")
+        engine = DominantColorsEngine(config, device, config_path=config_path)
     elif config['METHOD'] == 'CAPTION_YOLO_NEAREST':
         logger.info("Creating CaptionYoloNearestEngine engine")
         engine = CaptionYoloNearestEngine(config, device, config_path=config_path)

anno_page/engines/__init__.py

@@ -4,3 +4,4 @@
 from .detection import YoloDetectionEngine
 from .embedding import HuggingfaceTextEmbeddingEngine, HuggingfaceImageEmbeddingEngine
 from .translation import TranslationEngine
+from .color import DominantColorsEngine

anno_page/engines/captioning.py

@@ -14,8 +14,8 @@
 from urllib.parse import urljoin
 
 from anno_page.core.utils import compose_path, config_get_list
+from anno_page.core.metadata import GraphicalObjectMetadata, RelatedLinesMetadata, ColorInfo
 from anno_page.core.llm_api_aliases import get_llm_api_aliases
-from anno_page.core.metadata import GraphicalObjectMetadata, RelatedLinesMetadata
 from anno_page.engines import BaseEngine, LayoutProcessingEngine
 from anno_page.engines.detection import YoloDetector
 from anno_page.enums import Language, LineRelation
@@ -404,8 +404,8 @@ def process_image_captions(self, data: list[PromptData]):
             }
 
             metadata.color = {
-                Language.ENGLISH: item.result.color_en,
-                Language.CZECH: item.result.color_cz
+                Language.ENGLISH: ColorInfo(color_mode=item.result.color_en),
+                Language.CZECH: ColorInfo(color_mode=item.result.color_cz)
             }
 
             if metadata.prompts is None:

anno_page/engines/color.py

@@ -0,0 +1,153 @@
+import json
+
+import cv2
+import numpy as np
+
+from skimage.color import deltaE_ciede2000, rgb2lab
+from pydantic import BaseModel
+
+from anno_page.engines import LayoutProcessingEngine
+from anno_page.core.utils import compose_path, config_get_list
+from anno_page.core.metadata import GraphicalObjectMetadata, DominantColorInfo, ColorInfo
+from anno_page.enums import Language
+
+
+class ColorDefinition(BaseModel):
+    names: dict[str, str]
+    variants: dict[str, str]
+
+
+class ColorCoverage:
+    def __init__(self, names: dict[str, str], coverage: float):
+        self.names = names
+        self.coverage = coverage
+
+
+class NamedColors:
+    def __init__(self, colors, color_names, colors_mapping, colors_lab):
+        self.colors: np.ndarray = colors
+        self.color_names: list[dict[str, str]] = color_names
+        self.colors_mapping: np.ndarray = colors_mapping
+        self.colors_lab: np.ndarray = colors_lab
+
+
+class DominantColorsEngine(LayoutProcessingEngine):
+    def __init__(self, config, device, config_path):
+        super().__init__(config, device, config_path)
+
+        self.named_colors = self.load_colors(compose_path(self.config["colors"], self.config_path))
+        self.categories = config_get_list(self.config, key="categories", fallback=None, make_lowercase=True)
+        self.coverage_threshold = self.config.getfloat("coverage_threshold", 0.1)
+        self.max_size = self.config.getint("max_size", 256)
+        self.gaussian_blur_kernel_size = self.config.getint("gaussian_blur_kernel_size", fallback=0)
+
+    @staticmethod
+    def load_colors(path) -> NamedColors:
+        with open(path, "r", encoding="utf8") as fh:
+            data = json.load(fh)
+
+        color_definitions = [ColorDefinition(**item) for item in data]
+
+        colors = []
+        color_names = []
+        colors_mapping = []
+
+        for i, color_definition in enumerate(color_definitions):
+            color_names.append(color_definition.names)
+            for color_variant_name, color_variant in color_definition.variants.items():
+                colors.append(DominantColorsEngine.hex_to_rgb(color_variant))
+                colors_mapping.append(i)
+
+        colors = np.array(colors, dtype=np.uint8)
+        colors_mapping = np.array(colors_mapping)
+        colors_lab = rgb2lab(colors).reshape(-1, 3)
+
+        for i, color_name in enumerate(color_names):
+            color_variants = colors[colors_mapping == i]
+            color_name["color"] = np.clip(np.mean(color_variants, axis=0), 0, 255).astype(np.uint8)
+
+        named_colors = NamedColors(colors, color_names, colors_mapping, colors_lab)
+
+        return named_colors
+
+    def process_page(self, page_image, page_layout):
+        for i, region in enumerate(page_layout.regions):
+            if region.category is None or region.category.lower() == "text":
+                continue
+
+            if self.categories is None or region.category.lower() in self.categories:
+                x_min, y_min, x_max, y_max = region.get_polygon_bounding_box()
+                region_image = page_image[y_min:y_max, x_min:x_max]
+
+                if region_image.size == 0:
+                    continue
+
+                if self.gaussian_blur_kernel_size > 0:
+                    region_image = cv2.GaussianBlur(region_image, (self.gaussian_blur_kernel_size, self.gaussian_blur_kernel_size), 0)
+
+                region_image = cv2.cvtColor(region_image, cv2.COLOR_BGR2RGB)
+                region_image = self.resize_to_max_size(region_image)
+
+                region_colors = self.process_crop(region_image)
+                metadata: GraphicalObjectMetadata = region.graphical_metadata
+
+                for region_color in region_colors:
+                    if region_color.coverage >= self.coverage_threshold:
+                        if metadata.color is None:
+                            metadata.color = ColorInfo()
+
+                        if isinstance(metadata.color, dict):
+                            for language in metadata.color:
+                                dominant_color = DominantColorInfo(name=region_color.names[language.to_string()], coverage=region_color.coverage)
+
+                                if metadata.color[language].dominant_colors is None:
+                                    metadata.color[language].dominant_colors = [dominant_color]
+                                else:
+                                    metadata.color[language].dominant_colors.append(dominant_color)
+
+                        elif isinstance(metadata.color, ColorInfo):
+                            if Language.ENGLISH.to_string() in region_color.names:
+                                color_name = region_color.names[Language.ENGLISH.to_string()]
+                            elif Language.CZECH.to_string() in region_color.names:
+                                color_name = region_color.names[Language.CZECH.to_string()]
+                            else:
+                                color_name = list(region_color.names.values())[0]
+
+                            dominant_color = DominantColorInfo(name=color_name, coverage=region_color.coverage)
+
+                            if metadata.color.dominant_colors is None:
+                                metadata.color.dominant_colors = [dominant_color]
+                            else:
+                                metadata.color.dominant_colors.append(dominant_color)
+
+        return page_layout
+
+    def process_crop(self, image) -> list[ColorCoverage]:
+        pixels_lab = rgb2lab(image).reshape(-1, 3)
+        distances = deltaE_ciede2000(pixels_lab[:, None, :], self.named_colors.colors_lab[None, :, :])
+        assignments = np.argmin(distances, axis=1)
+        color_assignments = self.named_colors.colors_mapping[assignments]
+        counts = np.bincount(color_assignments, minlength=len(self.named_colors.color_names))
+        total_pixels = len(assignments)
+
+        coverages = [ColorCoverage(names=name, coverage=counts[index] / total_pixels) for index, name in enumerate(self.named_colors.color_names)]
+        return coverages
+
+    def resize_to_max_size(self, image):
+        if self.max_size <= 0:
+            return image
+
+        height, width = image.shape[:2]
+        current_size = max(width, height)
+        if current_size <= self.max_size:
+            return image
+
+        scale = self.max_size / current_size
+        new_size = (round(width * scale), round(height * scale))
+        resized = cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)
+        return resized
+
+    @staticmethod
+    def hex_to_rgb(hex_color):
+        hex_color = hex_color.lstrip("#")
+        return tuple(int(hex_color[i: i + 2], 16) for i in (0, 2, 4))