docling-eval/docling_eval/evaluators/ocr/processing_utils.py

import json
import logging
import traceback
from typing import Any, Dict, List, Optional, Tuple

from docling_core.types.doc import BoundingBox, CoordOrigin
from docling_core.types.doc.page import (
    BoundingRectangle,
    SegmentedPage,
    TextCell,
    TextDirection,
)

from docling_eval.evaluators.ocr.evaluation_models import Word, _CalculationConstants
from docling_eval.evaluators.ocr.geometry_utils import create_polygon_from_bbox

_log = logging.getLogger(__name__)


def extract_word_from_text_cell(text_cell: TextCell, page_height: float) -> Word:
    rect_to_process = text_cell.rect
    if rect_to_process.coord_origin != CoordOrigin.TOPLEFT:
        rect_to_process = rect_to_process.to_top_left_origin(page_height=page_height)

    polygon_points: List[List[float]] = [
        [float(rect_to_process.r_x0), float(rect_to_process.r_y0)],
        [float(rect_to_process.r_x1), float(rect_to_process.r_y1)],
        [float(rect_to_process.r_x2), float(rect_to_process.r_y2)],
        [float(rect_to_process.r_x3), float(rect_to_process.r_y3)],
    ]

    bbox = rect_to_process.to_bounding_box()

    width_val: float = bbox.width
    height_val: float = bbox.height

    is_vertical_flag: bool = False
    if (
        width_val > _CalculationConstants.EPS
        and height_val > (2 * width_val)
        and len(text_cell.text) > 1
    ):
        is_vertical_flag = True

    return Word(
        rect=rect_to_process,
        text=text_cell.text,
        orig=text_cell.orig,
        text_direction=text_cell.text_direction,
        confidence=text_cell.confidence,
        from_ocr=text_cell.from_ocr,
        vertical=is_vertical_flag,
        polygon=polygon_points,
    )


def convert_word_to_text_cell(word_obj: Word) -> TextCell:
    source_bbox = word_obj.bbox

    bbox_for_conversion = BoundingBox(
        l=source_bbox.l,
        t=source_bbox.t,
        r=source_bbox.r,
        b=source_bbox.b,
        coord_origin=CoordOrigin.TOPLEFT,
    )
    if source_bbox.coord_origin != CoordOrigin.TOPLEFT:
        pass

    text_cell_rect = BoundingRectangle.from_bounding_box(bbox_for_conversion)

    return TextCell(
        rect=text_cell_rect,
        text=word_obj.text,
        orig=word_obj.orig,
        confidence=word_obj.confidence,
        text_direction=word_obj.text_direction,
        from_ocr=word_obj.from_ocr,
    )


def merge_words_into_one(
    words: List[Word], add_space_between_words: bool = True
) -> Word:
    if not words:
        default_bbox = BoundingBox(l=0, t=0, r=0, b=0, coord_origin=CoordOrigin.TOPLEFT)
        default_rect = BoundingRectangle.from_bounding_box(default_bbox)
        return Word(
            text="",
            rect=default_rect,
            orig="",
            confidence=1.0,
            from_ocr=False,
            text_direction=TextDirection.LEFT_TO_RIGHT,
            vertical=False,
            polygon=create_polygon_from_bbox(default_bbox),
        )

    separator: str = " " if add_space_between_words else ""
    merged_text_parts = []

    min_left: float = float("inf")
    min_top: float = float("inf")
    max_right: float = -float("inf")
    max_bottom: float = -float("inf")

    sorted_words: List[Word] = sorted(words, key=lambda k: k.bbox.l)

    first_word_for_metadata = sorted_words[0]

    for word_item in sorted_words:
        merged_text_parts.append(word_item.text)
        current_bbox = word_item.bbox
        min_left = min(min_left, current_bbox.l)
        min_top = min(min_top, current_bbox.t)
        max_right = max(max_right, current_bbox.r)
        max_bottom = max(max_bottom, current_bbox.b)

    merged_text: str = separator.join(merged_text_parts)

    merged_bbox = BoundingBox(
        l=min_left,
        t=min_top,
        r=max_right,
        b=max_bottom,
        coord_origin=CoordOrigin.TOPLEFT,
    )

    merged_polygon: List[List[float]] = create_polygon_from_bbox(merged_bbox)
    merged_rect = BoundingRectangle.from_bounding_box(merged_bbox)

    return Word(
        text=merged_text,
        rect=merged_rect,
        orig=merged_text,
        confidence=first_word_for_metadata.confidence,
        from_ocr=any(w.from_ocr for w in sorted_words),
        text_direction=first_word_for_metadata.text_direction,
        vertical=first_word_for_metadata.vertical,
        polygon=merged_polygon,
    )


class _IgnoreZoneFilter:
    def __init__(self) -> None:
        pass

    def filter_words_in_ignore_zones(
        self, prediction_words: List[Word], ground_truth_words: List[Word]
    ) -> Tuple[List[Word], List[Word], List[Word]]:
        ignore_zones: List[Word] = []

        temp_ground_truth_words: List[Word] = list(ground_truth_words)
        for gt_word in temp_ground_truth_words:
            if gt_word.ignore_zone is True:
                ignore_zones.append(gt_word)
                gt_word.to_remove = True

        for ignore_zone_word in ignore_zones:
            self._mark_intersecting_words_for_removal(
                ignore_zone_word.bbox, ground_truth_words
            )
            self._mark_intersecting_words_for_removal(
                ignore_zone_word.bbox, prediction_words
            )

        filtered_ground_truth_words: List[Word] = [
            word for word in ground_truth_words if not word.to_remove
        ]
        filtered_prediction_words: List[Word] = [
            word for word in prediction_words if not word.to_remove
        ]

        return filtered_ground_truth_words, filtered_prediction_words, ignore_zones

    def _mark_intersecting_words_for_removal(
        self, ignore_zone_bbox: BoundingBox, words_list: List[Word]
    ) -> None:
        for word_item in words_list:
            if self._check_intersection(word_item.bbox, ignore_zone_bbox):
                word_item.to_remove = True

    def _check_intersection(self, bbox1: BoundingBox, bbox2: BoundingBox) -> bool:
        bbox1_width: float = bbox1.width
        bbox1_height: float = bbox1.height

        x_overlap: float = bbox1.x_overlap_with(bbox2)
        y_overlap: float = bbox1.y_overlap_with(bbox2)

        x_overlap_ratio: float = 0.0 if bbox1_width == 0 else x_overlap / bbox1_width
        y_overlap_ratio: float = 0.0 if bbox1_height == 0 else y_overlap / bbox1_height

        if y_overlap_ratio < 0.1 or x_overlap_ratio < 0.1:
            return False
        else:
            return True


def parse_segmented_pages(
    segmented_pages_raw_data: Any, document_id: str
) -> Optional[Dict[int, SegmentedPage]]:
    segmented_pages_map: Dict[int, SegmentedPage] = {}
    if isinstance(segmented_pages_raw_data, (bytes, str)):
        try:
            segmented_pages_payload: Any = json.loads(segmented_pages_raw_data)
        except json.JSONDecodeError as e:
            _log.warning(
                f"JSONDecodeError for doc {document_id}: {e}. Data: {str(segmented_pages_raw_data)[:200]}"
            )
            return None
    elif isinstance(segmented_pages_raw_data, dict):
        segmented_pages_payload = segmented_pages_raw_data
    else:
        _log.warning(
            f"Unrecognized segmented_pages data format for doc {document_id}: {type(segmented_pages_raw_data)}"
        )
        return None

    if not isinstance(segmented_pages_payload, dict):
        _log.warning(
            f"Expected dict payload for segmented_pages for doc {document_id}, got {type(segmented_pages_payload)}"
        )
        return None

    for page_index_str, page_data in segmented_pages_payload.items():
        try:
            page_index: int = int(page_index_str)
        except ValueError:
            _log.warning(
                f"Invalid page index string '{page_index_str}' for doc {document_id}. Skipping page."
            )
            continue

        try:
            if isinstance(page_data, dict):
                segmented_pages_map[page_index] = SegmentedPage.model_validate(
                    page_data
                )
            elif isinstance(page_data, str):
                segmented_pages_map[page_index] = SegmentedPage.model_validate_json(
                    page_data
                )
            elif isinstance(page_data, SegmentedPage):
                segmented_pages_map[page_index] = page_data
            else:
                _log.warning(
                    f"Unrecognized page_data format for doc {document_id}, page {page_index}: {type(page_data)}"
                )
                continue
        except Exception as e_page_val:
            _log.error(
                f"Error validating page data for doc {document_id}, page {page_index}: {e_page_val}"
            )
            traceback.print_exc()
            continue
    return segmented_pages_map if segmented_pages_map else None