documentExtractionag051/ExtractDocument

1import json2import os3import uuid4import re5from typing import List6 7class Rect:8    # A shared ID counter for disambiguation9    _id = 010    11    def __init__(self, x1=0, y1=0, x2=0, y2=0, block_type=None, confidence=0.0, uuid_value=None):12        self.id = Rect._id13        Rect._id += 114        15        self.x1 = x116        self.y1 = y117        self.x2 = x218        self.y2 = y219        self.block_type = block_type20        self.confidence = confidence21        self.uuid = str(uuid_value or uuid.uuid4()).upper()22    23    @property24    def xcenter(self):25        return (self.x2 + self.x1) / 226    27    @property28    def ycenter(self):29        return (self.y2 + self.y1) / 230    31    @property32    def width(self):33        return self.x2 - self.x134    35    @property36    def height(self):37        return self.y2 - self.y138    39    def get_geometry(self):40        return {41            'x1': int(self.x1),42            'y1': int(self.y1),43            'x2': int(self.x2),44            'y2': int(self.y2)45        }46    47    def as_block(self):48        return {49            "id": str(self.uuid),50            "blockType": self.block_type,51            "geometry": self.get_geometry(),52            "confidence": self.confidence,53        }54    55    def intersection_pct(self, boxB):56        """57        Given two bounding boxes, it returns the percentage of the boxB within this box.58        59        Returns:60            (int) Percent IOU score61        """62        # Determine the (x, y)-coordinates of the intersection rectangle63        xA = max(self.x1, boxB.x1)64        yA = max(self.y1, boxB.y1)65        xB = min(self.x2, boxB.x2)66        yB = min(self.y2, boxB.y2)67        68        # Compute the area of the intersection rectangle69        interArea = round(abs(max((xB - xA, 0)) * max((yB - yA, 0))), 2)70        if interArea == 0:71            return 072        73        # Compute the intersection over the area of boxB74        boxB_area = boxB.width * boxB.height75        if boxB_area == 0:76            return 077        78        iou = round((interArea / float(boxB_area)) * 100, 2)79        return iou80    81    def vertical_alignment_pct(self, boxB):82        """83        Measures the percentage of vertical (y-axis) overlap between two boxes,84        relative to the height of boxB.85        86        Returns:87            (float) percentage of vertical alignment88        """89        # Determine the vertical overlap90        yA = max(self.y1, boxB.y1)91        yB = min(self.y2, boxB.y2)92        93        interHeight = max(yB - yA, 0)94        if interHeight == 0:95            return 0.096        97        boxB_height = boxB.height98        if boxB_height == 0:99            return 0.0100        101        val = round((interHeight / float(boxB_height)) * 100, 2)102        return val103    104    def horizontal_alignment_pct(self, boxB):105        """106        Measures the percentage of horizontal (x-axis) overlap between two boxes,107        relative to the width of boxB.108        109        Returns:110            (float) percentage of horizontal alignment111        """112        # Determine the horizontal overlap113        xA = max(self.x1, boxB.x1)114        xB = min(self.x2, boxB.x2)115        116        interWidth = max(xB - xA, 0)117        if interWidth == 0:118            return 0.0119        120        boxB_width = boxB.width121        if boxB_width == 0:122            return 0.0123        124        val = round((interWidth / float(boxB_width)) * 100, 2)125        return val126    127    def iou_pct(self, boxB):128        """129        Given two bounding boxes, it returns the percentage of the intersection over union.130        131        Returns:132            (int) Percent IOU score133        """134        # Determine the (x, y)-coordinates of the intersection rectangle135        xA = max(self.x1, boxB.x1)136        yA = max(self.y1, boxB.y1)137        xB = min(self.x2, boxB.x2)138        yB = min(self.y2, boxB.y2)139        140        # Compute the area of the intersection rectangle141        interArea = round(abs(max((xB - xA, 0)) * max((yB - yA, 0))), 2)142        if interArea == 0:143            return 0144        145        # Compute the union area146        xA = min(self.x1, boxB.x1)147        yA = min(self.y1, boxB.y1)148        xB = max(self.x2, boxB.x2)149        yB = max(self.y2, boxB.y2)150        union_area = round(max((xB - xA, 0)) * max((yB - yA, 0)), 2)151        if union_area == 0:152            return 0153        154        iou = round((interArea / float(union_area)) * 100, 2)155        return iou156 157 158class TextRect:159    def __init__(self, x1, y1, x2, y2, confidence, block_type, text, page_num):160        """Initialize a TextRect object.161        162        Args:163            x1 (int): The x-coordinate of the top-left corner.164            y1 (int): The y-coordinate of the top-left corner.165            x2 (int): The x-coordinate of the bottom-right corner.166            y2 (int): The y-coordinate of the bottom-right corner.167            confidence (float): The confidence score of the element.168            block_type (str): The type of block (e.g., 'text', 'image').169            text (str): The text content of the element.170            page_num (int): The page number where the element is located.171        """172        self.x1 = x1173        self.y1 = y1174        self.x2 = x2175        self.y2 = y2176        self.confidence = confidence177        self.block_type = block_type178        self.text = text179        self.page_num = page_num180    181    @property182    def width(self):183        return self.x2 - self.x1184    185    @property186    def height(self):187        return self.y2 - self.y1188    189    @property190    def midpoint(self):191        return ((self.x1 + self.x2) / 2, (self.y1 + self.y2) / 2)192    193    def intersection_pct(self, boxB):194        """Calculate the percentage of boxB within the current box.195        196        Args:197            boxB (TextRect): Another bounding box.198        199        Returns:200            float: The percentage of boxB within the current box.201        """202        xA = max(self.x1, boxB.x1)203        yA = max(self.y1, boxB.y1)204        xB = min(self.x2, boxB.x2)205        yB = min(self.y2, boxB.y2)206        207        interArea = round(abs(max((xB - xA, 0)) * max((yB - yA, 0))), 2)208        if interArea == 0:209            return 0210        211        boxB_area = boxB.width * boxB.height212        if boxB_area == 0:213            return 0214        215        iou = round((interArea / float(boxB_area)) * 100, 2)216        return iou217    218    def vertical_alignment_pct(self, boxB):219        """220        Measures the percentage of vertical (y-axis) overlap between two boxes,221        relative to the height of boxB.222        223        Returns:224            (float) percentage of vertical alignment225        """226        # Determine the vertical overlap227        yA = max(self.y1, boxB.y1)228        yB = min(self.y2, boxB.y2)229        230        interHeight = max(yB - yA, 0)231        if interHeight == 0:232            return 0.0233        234        boxB_height = boxB.height235        if boxB_height == 0:236            return 0.0237        238        val = round((interHeight / float(boxB_height)) * 100, 2)239        return val240    241    def horizontal_alignment_pct(self, boxB):242        """243        Measures the percentage of horizontal (x-axis) overlap between two boxes,244        relative to the width of boxB.245        246        Returns:247            (float) percentage of horizontal alignment248        """249        # Determine the horizontal overlap250        xA = max(self.x1, boxB.x1)251        xB = min(self.x2, boxB.x2)252        253        interWidth = max(xB - xA, 0)254        if interWidth == 0:255            return 0.0256        257        boxB_width = boxB.width258        if boxB_width == 0:259            return 0.0260        261        val = round((interWidth / float(boxB_width)) * 100, 2)262        return val263    264    def __repr__(self):265        return f"TextRect(x1={self.x1}, y1={self.y1}, x2={self.x2}, y2={self.y2}, confidence={self.confidence}, block_type={self.block_type}, text={self.text}, page_num={self.page_num})"266    267    def _find_words_in_direction(self, words, threshold, direction):268        """Common helper for finding words in any direction.269        270        Args:271            words: List of TextRect objects to search through272            threshold: Search threshold (0, positive, or negative)273            direction: One of 'left', 'right', 'above', 'below'274        275        Returns:276            List of TextRect objects in the specified direction277        """278        result_words = []279        280        for word in words:281            if word.page_num != self.page_num:282                continue283            284            # Check positional constraint based on direction285            if direction == 'left' and word.x2 > self.x1:286                continue287            elif direction == 'right' and word.x1 < self.x2:288                continue289            elif direction == 'above' and word.y2 > self.y1:290                continue291            elif direction == 'below' and word.y1 < self.y2:292                continue293            294            # For negative threshold, collect all qualifying words295            if threshold < 0:296                result_words.append(word)297                continue298            299            # Check alignment based on direction300            if direction in ['left', 'right']:301                # Horizontal directions - check y-axis alignment302                if threshold == 0:303                    if word.y1 < self.y2 and word.y2 > self.y1:304                        result_words.append(word)305                else:306                    expanded_y1 = self.y1 - threshold307                    expanded_y2 = self.y2 + threshold308                    if word.y1 < expanded_y2 and word.y2 > expanded_y1:309                        result_words.append(word)310            else:  # 'above' or 'below'311                # Vertical directions - check x-axis alignment312                if threshold == 0:313                    if word.x1 < self.x2 and word.x2 > self.x1:314                        result_words.append(word)315                else:316                    expanded_x1 = self.x1 - threshold317                    expanded_x2 = self.x2 + threshold318                    if word.x1 < expanded_x2 and word.x2 > expanded_x1:319                        result_words.append(word)320        321        # Return only nearest word if threshold is negative322        if threshold < 0 and result_words:323            nearest = self.find_nearest_word(result_words)324            return [nearest] if nearest is not None else []325        326        return result_words327    328    def find_words_to_the_left_of_text(self, words, threshold=0):329        """Find words to the left of the current text."""330        return self._find_words_in_direction(words, threshold, 'left')331    332    def find_words_to_the_right_of_text(self, words, threshold=0):333        """Find words to the right of the current text."""334        return self._find_words_in_direction(words, threshold, 'right')335    336    def find_words_above_the_text(self, words, threshold=0):337        """Find words above the current text."""338        return self._find_words_in_direction(words, threshold, 'above')339    340    def find_words_below_the_text(self, words, threshold=0):341        """Find words below the current text."""342        return self._find_words_in_direction(words, threshold, 'below')343    344    def get_coordinates(self):345        return (self.x1, self.y1, self.x2, self.y2)346    347    def find_nearest_word(self, words):348        nearest_word = None349        nearest_distance = float('inf')350        351        for word in words:352            if word.page_num != self.page_num:353                continue354            355            distance = ((word.x1 - self.x1) ** 2 + (word.y1 - self.y1) ** 2) ** 0.5356            if distance < nearest_distance:357                nearest_distance = distance358                nearest_word = word359        360        return nearest_word361 362 363class DocumentWords:364    365    @staticmethod366    def _convert_lower(word):367        return word.strip().lower()368    369    @staticmethod370    def find_words_in_area(words, x1, y1, x2, y2):371        area_words = []372        for word in words:373            if word.x1 >= x1 and word.x2 <= x2 and word.y1 >= y1 and word.y2 <= y2:374                area_words.append(word)375        return area_words376    377    @staticmethod378    def find_words_on_document(self, word_texts, words, only_first=False)->List[TextRect]:379        if isinstance(word_texts, str):380            word_texts = [word_texts]381        all_matches = []382        for word in words:383            for word_text in word_texts:384                if self._convert_lower(word.text) == self._convert_lower(word_text):385                    all_matches.append(word)386                    if only_first:387                        return all_matches388        return all_matches389    390    @staticmethod391    def find_words_in_column(x1, x2, words):392        column_words = []393        for word in words:394            if word.x1 >= x1 and word.x2 <= x2:395                column_words.append(word)396        return column_words397    398    @staticmethod399    def find_words_in_row(y1, y2, words):400        row_words = []401        for word in words:402            if word.y1 >= y1 and word.y2 <= y2:403                row_words.append(word)404        return row_words405    406    @staticmethod407    def find_in_box_words(x1, y1, x2, y2, words):408        box_words = []409        for word in words:410            if word.x1 >= x1 and word.x2 <= x2 and word.y1 >= y1 and word.y2 <= y2:411                box_words.append(word)412        return box_words413    414    def find_value_for_key(self, key_word: TextRect, value_candidates)->TextRect|None:415        closest_value = None416        closest_distance = float('inf')417        418        for value_word in value_candidates:419            if value_word.page_num != key_word.page_num:420                continue421            422            # Check if value is to the right of the key423            if value_word.x1 >= key_word.x2:424                distance = value_word.x1 - key_word.x2425                if distance < closest_distance:426                    closest_distance = distance427                    closest_value = value_word428        429        return closest_value430    431    @staticmethod432    def find_words_by_regex(pattern: str, words)->List[TextRect]:433        regex = re.compile(pattern)434        matched_words = []435        for word in words:436            if regex.match(word.text):437                matched_words.append(word)438        return matched_words439 440 441def compute_global_bounds(obj, page_heights):442    """443    Compute global bounds for an object whose y-coordinates are page-relative.444    445    Parameters446    ----------447    obj : Any object with attributes:448        - page_num449        - x1, y1450        - x2, y2451    page_heights : list[int|float]452        A list of heights of each document page in order.453    454    Returns455    -------456    (x, y, w, h) : tuple[float]457        Global coordinates with y adjusted using cumulative height.458    """459    # Precompute cumulative heights only once per document460    # (You can move this out if doing many calls)461    cumulative = [0]462    for h in page_heights[:-1]:463        cumulative.append(cumulative[-1] + h)464    465    pg_idx = max(obj.page_num - 1, 0)  # 1-indexed → 0-indexed466    y_offset = cumulative[pg_idx]467    468    x = obj.x1469    y = obj.y1 + y_offset470    w = obj.x2 - obj.x1471    h = obj.y2 - obj.y1472    473    return x, y, w, h474