Spaces:
Paused
Paused
File size: 6,812 Bytes
f66643d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """Table models: table structure and cell recognition."""
from __future__ import annotations
import logging
from typing import Any
import numpy as np
from PIL import Image
from pdf2zh.parser.ai_models.base import BaseImageToTextModel
from pdf2zh.parser.utils.bbox import bbox_area, bbox_intersection
logger = logging.getLogger(__name__)
class SuryaTableModel(BaseImageToTextModel):
"""
Wraps Surya's TableRecPredictor.
Identifies row/column structure and cell bounding boxes within a cropped
table image. Text extraction is handled separately.
Models are loaded lazily upon first inference call.
"""
model_name = "SuryaTable"
def __init__(self) -> None:
"""Initialize empty state to defer model loading."""
super().__init__()
def load_model(self) -> None:
"""Load Surya model into VRAM."""
logger.info("Initializing %s...", self.model_name)
from surya.table_rec import TableRecPredictor
self.model = TableRecPredictor()
logger.info("Loaded TableRecPredictor")
def prepare(
self, images: list[Image.Image], *args: Any, **kwargs: Any
) -> list[Image.Image]:
"""Preprocess a batch of cropped table images."""
# Surya models accept raw PIL images directly
return images
def predict(
self,
prepared_inputs: list[Image.Image],
batch_size: int | None = None,
*args: Any,
**kwargs: Any,
) -> list[Any]:
"""
Recognize table structure for a batch of prepared table images.
"""
try:
# self.model is guaranteed to be loaded by the Base class
raw_results = self.model(
prepared_inputs,
batch_size=batch_size,
)
return raw_results
except Exception:
logger.exception(
"Table recognition failed for batch of %d crops — returning nulls.",
len(prepared_inputs),
)
return [None] * len(prepared_inputs)
def postprocess(
self, raw_results: list[Any], *args: Any, **kwargs: Any
) -> list[list[list[float]]]:
"""Convert objects into a simple list of bounding boxes."""
batch_boxes = []
for result in raw_results:
if result is None:
batch_boxes.append([])
continue
# Extract only bboxes and ensure float type
boxes = [
[float(x) for x in cell.bbox] for cell in getattr(result, "cells", [])
]
batch_boxes.append(boxes)
return batch_boxes
class PaddleCellTableModule(BaseImageToTextModel):
"""
Wraps Paddle's Table Cell Detection Module.
Models are loaded lazily upon first inference call.
"""
model_name = "PaddleCellTableModule"
def __init__(self) -> None:
"""Initialize empty state to defer model loading."""
super().__init__()
def load_model(self) -> None:
"""Load Paddle model into memory/VRAM."""
logger.info("Initializing %s...", self.model_name)
from paddleocr import TableCellsDetection
self.model = TableCellsDetection(model_name="RT-DETR-L_wireless_table_cell_det")
logger.info("Loaded TableCellsDetection")
def prepare(
self, images: list[Image.Image], *args: Any, **kwargs: Any
) -> list[np.ndarray]:
"""
Convert PIL images to numpy arrays to satisfy PaddleOCR requirements.
"""
return [np.array(img.convert("RGB")) for img in images]
def predict(
self,
prepared_inputs: list[np.ndarray],
batch_size: int | None = None,
threshold: float = 0.3,
*args: Any,
**kwargs: Any,
) -> list[Any]:
"""
Recognize cell detection for a batch of prepared table images.
"""
try:
raw_results = self.model.predict(
prepared_inputs,
threshold=threshold,
batch_size=batch_size,
)
return raw_results
except Exception:
logger.exception(
"Paddle table cell detection failed for batch of %d crops — returning nulls.",
len(prepared_inputs),
)
return [None] * len(prepared_inputs)
def postprocess(
self, raw_results: list[Any], *args: Any, **kwargs: Any
) -> list[list[list[float]]]:
"""Normalize Paddle output into simple bbox lists."""
batch_boxes = []
for result in raw_results:
if result is None:
batch_boxes.append([])
continue
# Check both 'boxes' and 'coordinate' attributes
raw_cells = result.get("boxes", [])
boxes = []
for cell in raw_cells:
coords = cell.get("coordinate")
if coords:
boxes.append([float(x) for x in coords])
batch_boxes.append(self._prune_nested_cell_boxes(boxes))
return batch_boxes
def _prune_nested_cell_boxes(
self,
boxes: list[list[float]],
containment_threshold: float = 0.8,
) -> list[list[float]]:
if len(boxes) < 2:
return boxes
kept_boxes: list[list[float]] = []
sorted_boxes = sorted(boxes, key=bbox_area)
for box in sorted_boxes:
box_area = max(1.0, bbox_area(box))
is_duplicate = False
for kept in kept_boxes:
intersection = bbox_intersection(box, kept)
if intersection is None:
continue
overlap_ratio = bbox_area(intersection) / box_area
if overlap_ratio >= containment_threshold:
is_duplicate = True
break
if not is_duplicate:
kept_boxes.append(box)
filtered_boxes: list[list[float]] = []
for box in kept_boxes:
box_area = max(1.0, bbox_area(box))
contains_smaller_box = False
for other in kept_boxes:
if other is box:
continue
other_area = bbox_area(other)
if other_area >= box_area:
continue
intersection = bbox_intersection(box, other)
if intersection is None:
continue
overlap_ratio = bbox_area(intersection) / max(1.0, other_area)
if overlap_ratio >= containment_threshold:
contains_smaller_box = True
break
if not contains_smaller_box:
filtered_boxes.append(box)
return filtered_boxes
|