documentExtractionag051/ExtractDocument

1from enum import Enum2 3from utils.geometry_utils import TextRect4 5 6class SortMethod(Enum):7    PAGE = 'page'8    HORIZONTAL = 'horizontal'9    VERTICAL = 'vertical'10 11 12class SortRects:13    14    @staticmethod15    def sort(**kwargs):16        rects = kwargs.get('rects', [])17        method = kwargs.get('method', SortMethod.PAGE.value)18        reverse = kwargs.get('order', 'asc') == 'desc'19        20        if method == SortMethod.PAGE.value:21            return SortRects.sort_by_page(rects, reverse=reverse)22        elif method == SortMethod.HORIZONTAL.value:23            return SortRects.sort_horizontally(rects, reverse=reverse)24        elif method == SortMethod.VERTICAL.value:25            return SortRects.sort_vertically(rects, reverse=reverse)26        else:27            raise ValueError(f"Unknown sorting method: {method}")28    29    @staticmethod30    def sort_by_page(rects: list[TextRect], reverse=False):31        return sorted(rects, key=lambda r: r.page_num, reverse=reverse)32    33    @staticmethod34    def sort_horizontally(rects: list[TextRect], reverse=False):35        # sort by page first and then on x136        return sorted(rects, key=lambda r: (r.page_num, r.x1), reverse=reverse)37    38    @staticmethod39    def sort_vertically(rects: list[TextRect], reverse=False):40        return sorted(rects, key=lambda r: (r.page_num, r.y1), reverse=reverse)41