Implements a document builder Args: resolve_lines: whether words should be automatically grouped into lines resolve_blocks: whether lines should be automatically grouped into blocks paragraph_break: relative length of the minimum space separating paragraphs expor
| 17 | |
| 18 | |
| 19 | class DocumentBuilder(NestedObject): |
| 20 | """Implements a document builder |
| 21 | |
| 22 | Args: |
| 23 | resolve_lines: whether words should be automatically grouped into lines |
| 24 | resolve_blocks: whether lines should be automatically grouped into blocks |
| 25 | paragraph_break: relative length of the minimum space separating paragraphs |
| 26 | export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle |
| 27 | box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is. |
| 28 | """ |
| 29 | |
| 30 | def __init__( |
| 31 | self, |
| 32 | resolve_lines: bool = True, |
| 33 | resolve_blocks: bool = False, |
| 34 | paragraph_break: float = 0.035, |
| 35 | export_as_straight_boxes: bool = False, |
| 36 | ) -> None: |
| 37 | self.resolve_lines = resolve_lines |
| 38 | self.resolve_blocks = resolve_blocks |
| 39 | self.paragraph_break = paragraph_break |
| 40 | self.export_as_straight_boxes = export_as_straight_boxes |
| 41 | |
| 42 | @staticmethod |
| 43 | def _sort_boxes(boxes: np.ndarray) -> tuple[np.ndarray, np.ndarray]: |
| 44 | """Sort bounding boxes from top to bottom, left to right |
| 45 | |
| 46 | Args: |
| 47 | boxes: bounding boxes of shape (N, 4) or (N, 4, 2) (in case of rotated bbox) |
| 48 | |
| 49 | Returns: |
| 50 | tuple: indices of ordered boxes of shape (N,), boxes |
| 51 | If straight boxes are passed tpo the function, boxes are unchanged |
| 52 | else: boxes returned are straight boxes fitted to the straightened rotated boxes |
| 53 | so that we fit the lines afterwards to the straigthened page |
| 54 | """ |
| 55 | if boxes.ndim == 3: |
| 56 | boxes = rotate_boxes( |
| 57 | loc_preds=boxes, |
| 58 | angle=-estimate_page_angle(boxes), |
| 59 | orig_shape=(1024, 1024), |
| 60 | min_angle=5.0, |
| 61 | ) |
| 62 | boxes = np.concatenate((boxes.min(1), boxes.max(1)), -1) |
| 63 | return (boxes[:, 0] + 2 * boxes[:, 3] / np.median(boxes[:, 3] - boxes[:, 1])).argsort(), boxes |
| 64 | |
| 65 | def _resolve_sub_lines(self, boxes: np.ndarray, word_idcs: list[int]) -> list[list[int]]: |
| 66 | """Split a line in sub_lines |
| 67 | |
| 68 | Args: |
| 69 | boxes: bounding boxes of shape (N, 4) |
| 70 | word_idcs: list of indexes for the words of the line |
| 71 | |
| 72 | Returns: |
| 73 | A list of (sub-)lines computed from the original line (words) |
| 74 | """ |
| 75 | lines = [] |
| 76 | # Sort words horizontally |