unstructured-io--unstructured/unstructured/partition/utils/sorting.py

from __future__ import annotations

import os
from typing import TYPE_CHECKING, Any

import numpy as np

from unstructured.documents.elements import CoordinatesMetadata, Element
from unstructured.logger import trace_logger
from unstructured.partition.utils.constants import SORT_MODE_BASIC, SORT_MODE_XY_CUT
from unstructured.partition.utils.xycut import recursive_xy_cut, recursive_xy_cut_swapped

if TYPE_CHECKING:
    from unstructured_inference.inference.elements import TextRegions


def coordinates_to_bbox(coordinates: CoordinatesMetadata) -> tuple[int, int, int, int]:
    """
    Convert coordinates to a bounding box representation.

    Parameters:
        coordinates (CoordinatesMetadata): Metadata containing points to represent the bounding box.

    Returns:
        tuple[int, int, int, int]: A tuple representing the bounding box in the format
        (left, top, right, bottom).
    """

    points = coordinates.points
    left, top = points[0]
    right, bottom = points[2]
    return int(left), int(top), int(right), int(bottom)


def shrink_bbox(bbox: tuple[int, int, int, int], shrink_factor) -> tuple[int, int, int, int]:
    """
    Shrink a bounding box by a given shrink factor while maintaining its top and left.

    Parameters:
        bbox (tuple[int, int, int, int]): The original bounding box represented by
        (left, top, right, bottom).
        shrink_factor (float): The factor by which to shrink the bounding box (0.0 to 1.0).

    Returns:
        tuple[int, int, int, int]: The shrunken bounding box represented by
        (left, top, right, bottom).
    """

    left, top, right, bottom = bbox
    width = right - left
    height = bottom - top
    new_width = width * shrink_factor
    new_height = height * shrink_factor
    dw = width - new_width
    dh = height - new_height

    new_right = right - dw
    new_bottom = bottom - dh
    return int(left), int(top), int(new_right), int(new_bottom)


def coord_has_valid_points(coordinates: CoordinatesMetadata) -> bool:
    """
    Verifies all 4 points in a coordinate exist and are positive.
    """
    if not coordinates:
        return False
    if len(coordinates.points) != 4:
        return False
    for point in coordinates.points:
        if len(point) != 2:
            return False
        try:
            if point[0] < 0 or point[1] < 0:
                return False
        except TypeError:
            return False
    return True


def bbox_is_valid(bbox: Any) -> bool:
    """
    Verifies all 4 values in a bounding box exist and are positive.
    """

    if not bbox:
        return False
    if len(bbox) != 4:
        return False
    for v in bbox:
        try:
            if v < 0:
                return False
        except TypeError:
            return False
    return True


def sort_page_elements(
    page_elements: list[Element],
    sort_mode: str = SORT_MODE_XY_CUT,
    shrink_factor: float = 0.9,
    xy_cut_primary_direction: str = "x",
) -> list[Element]:
    """
    Sorts a list of page elements based on the specified sorting mode.

    Parameters:
    - page_elements (list[Element]): A list of elements representing parts of a page. Each element
     should have metadata containing coordinates.
    - sort_mode (str, optional): The mode by which the elements will be sorted. Default is
     SORT_MODE_XY_CUT.
        - SORT_MODE_XY_CUT: Sorts elements based on XY-cut sorting approach. Requires the
         recursive_xy_cut function and coordinates_to_bbox function to be defined. And requires all
         elements to have valid cooridnates
        - SORT_MODE_BASIC: Sorts elements based on their coordinates. Elements without coordinates
         will be pushed to the end.
        - If an unrecognized sort_mode is provided, the function returns the elements as-is.

    Returns:
    - list[Element]: A list of sorted page elements.
    """

    shrink_factor = float(
        os.environ.get("UNSTRUCTURED_XY_CUT_BBOX_SHRINK_FACTOR", shrink_factor),
    )

    xy_cut_primary_direction = os.environ.get(
        "UNSTRUCTURED_XY_CUT_PRIMARY_DIRECTION",
        xy_cut_primary_direction,
    )

    if not page_elements:
        return []

    coordinates_list = [el.metadata.coordinates for el in page_elements]

    def _coords_ok(strict_points: bool):
        warned = False

        for coord in coordinates_list:
            if coord is None or not coord.points:
                trace_logger.detail(  # type: ignore
                    "some or all elements are missing coordinates, skipping sort",
                )
                return False
            elif not coord_has_valid_points(coord):
                if not warned:
                    trace_logger.detail(f"coord {coord} does not have valid points")  # type: ignore
                    warned = True
                if strict_points:
                    return False
        return True

    if sort_mode == SORT_MODE_XY_CUT:
        if not _coords_ok(strict_points=True):
            return page_elements
        shrunken_bboxes = []
        for coords in coordinates_list:
            bbox = coordinates_to_bbox(coords)
            shrunken_bbox = shrink_bbox(bbox, shrink_factor)
            shrunken_bboxes.append(shrunken_bbox)

        res: list[int] = []
        xy_cut_sorting_func = (
            recursive_xy_cut_swapped if xy_cut_primary_direction == "x" else recursive_xy_cut
        )
        xy_cut_sorting_func(
            np.asarray(shrunken_bboxes).astype(int),
            np.arange(len(shrunken_bboxes)),
            res,
        )
        sorted_page_elements = [page_elements[i] for i in res]
    elif sort_mode == SORT_MODE_BASIC:
        if not _coords_ok(strict_points=False):
            return page_elements
        sorted_page_elements = sorted(
            page_elements,
            key=lambda el: (
                el.metadata.coordinates.points[0][1] if el.metadata.coordinates else float("inf"),
                el.metadata.coordinates.points[0][0] if el.metadata.coordinates else float("inf"),
            ),
        )
    else:
        sorted_page_elements = page_elements

    return sorted_page_elements


def sort_bboxes_by_xy_cut(
    bboxes,
    shrink_factor: float = 0.9,
    xy_cut_primary_direction: str = "x",
):
    """Sort bounding boxes using XY-cut algorithm."""

    shrunken_bboxes = []
    for bbox in bboxes:
        shrunken_bbox = shrink_bbox(bbox, shrink_factor)
        shrunken_bboxes.append(shrunken_bbox)

    res: list[int] = []
    xy_cut_sorting_func = (
        recursive_xy_cut_swapped if xy_cut_primary_direction == "x" else recursive_xy_cut
    )
    xy_cut_sorting_func(
        np.asarray(shrunken_bboxes).astype(int),
        np.arange(len(shrunken_bboxes)),
        res,
    )
    return res


def sort_text_regions(
    elements: TextRegions,
    sort_mode: str = SORT_MODE_XY_CUT,
    shrink_factor: float = 0.9,
    xy_cut_primary_direction: str = "x",
) -> TextRegions:
    """Sort a list of TextRegion elements based on the specified sorting mode."""

    if not elements:
        return elements

    bboxes = elements.element_coords

    def _bboxes_ok(strict_points: bool):
        if np.isnan(bboxes).any():
            trace_logger.detail(  # type: ignore
                "some or all elements are missing bboxes, skipping sort",
            )
            return False

        if bboxes.shape[1] != 4 or np.where(bboxes < 0)[0].size:
            trace_logger.detail("at least one bbox contains invalid values")  # type: ignore
            if strict_points:
                return False
        return True

    if sort_mode == SORT_MODE_XY_CUT:
        if not _bboxes_ok(strict_points=True):
            return elements

        shrink_factor = float(
            os.environ.get("UNSTRUCTURED_XY_CUT_BBOX_SHRINK_FACTOR", shrink_factor),
        )

        xy_cut_primary_direction = os.environ.get(
            "UNSTRUCTURED_XY_CUT_PRIMARY_DIRECTION",
            xy_cut_primary_direction,
        )

        res = sort_bboxes_by_xy_cut(
            bboxes=bboxes,
            shrink_factor=shrink_factor,
            xy_cut_primary_direction=xy_cut_primary_direction,
        )
        sorted_elements = elements.slice(res)
    elif sort_mode == SORT_MODE_BASIC:
        # NOTE (yao): lexsort order is revese from the input sequence; so below is first sort by y1,
        # then x1, then y2, lastly x2
        sorted_elements = elements.slice(
            np.lexsort((elements.x2, elements.y2, elements.x1, elements.y1))
        )
    else:
        sorted_elements = elements

    return sorted_elements