Spaces:
Running on Zero
Running on Zero
| """PP-DocLayoutV3 the way the OCR SDK runs it. | |
| The picture is resized to 800x800 in BGR and scaled to 0-1. Scores, boxes, and | |
| reading order come from the raw model outputs. Same-class boxes are suppressed, | |
| and a box that sits inside a larger one is dropped. | |
| """ | |
| from __future__ import annotations | |
| from typing import List | |
| import numpy as np | |
| import torch | |
| LABELS = [ | |
| "abstract", | |
| "algorithm", | |
| "aside_text", | |
| "chart", | |
| "content", | |
| "display_formula", | |
| "doc_title", | |
| "figure_title", | |
| "footer", | |
| "footer_image", | |
| "footnote", | |
| "formula_number", | |
| "header", | |
| "header_image", | |
| "image", | |
| "inline_formula", | |
| "number", | |
| "paragraph_title", | |
| "reference", | |
| "reference_content", | |
| "seal", | |
| "table", | |
| "text", | |
| "vertical_text", | |
| "vision_footnote", | |
| ] | |
| # Keep the large region when one of these classes contains another box. | |
| _KEEP_LARGE = {3, 5, 6, 15, 17} | |
| def _sigmoid(values: np.ndarray) -> np.ndarray: | |
| values = np.clip(values, -50, 50) | |
| return 1.0 / (1.0 + np.exp(-values)) | |
| def _reading_order(order_logits: np.ndarray) -> np.ndarray: | |
| scores = _sigmoid(order_logits.astype(np.float32)) | |
| upper = np.triu(scores, 1) | |
| lower = np.tril(np.float32(1.0) - scores.T, -1) | |
| votes = upper.sum(axis=0) + lower.sum(axis=0) | |
| sequence = np.zeros(order_logits.shape[0], dtype=np.int64) | |
| sequence[np.argsort(votes)] = np.arange(order_logits.shape[0]) | |
| return sequence | |
| def _boxes_7col(logits, pred_boxes, order_logits, width: int, height: int) -> np.ndarray: | |
| num_classes = logits.shape[1] | |
| flat = _sigmoid(logits).reshape(-1) | |
| top = min(300, flat.size) | |
| chosen = np.argpartition(-flat, top - 1)[:top] | |
| chosen = chosen[np.argsort(-flat[chosen])] | |
| labels = (chosen % num_classes).astype(np.float32) | |
| query = chosen // num_classes | |
| centers, sizes = pred_boxes[:, :2], pred_boxes[:, 2:] | |
| corners = np.concatenate([centers - 0.5 * sizes, centers + 0.5 * sizes], axis=-1) | |
| scaled = corners * np.array([width, height, width, height], dtype=np.float32) | |
| order = _reading_order(order_logits)[query].astype(np.float32) | |
| picked = scaled[query] | |
| return np.column_stack([labels, flat[chosen], picked, order]) | |
| def _iou(left: np.ndarray, right: np.ndarray) -> float: | |
| x1 = max(left[2], right[2]) | |
| y1 = max(left[3], right[3]) | |
| x2 = min(left[4], right[4]) | |
| y2 = min(left[5], right[5]) | |
| intersection_w = x2 - x1 + 1 | |
| intersection_h = y2 - y1 + 1 | |
| if intersection_w <= 0 or intersection_h <= 0: | |
| return 0.0 | |
| intersection = intersection_w * intersection_h | |
| area_left = (left[4] - left[2] + 1) * (left[5] - left[3] + 1) | |
| area_right = (right[4] - right[2] + 1) * (right[5] - right[3] + 1) | |
| return float(intersection / (area_left + area_right - intersection)) | |
| def _nms(boxes: np.ndarray) -> List[int]: | |
| remaining = list(np.argsort(-boxes[:, 1])) | |
| kept: List[int] = [] | |
| while remaining: | |
| current = remaining.pop(0) | |
| kept.append(current) | |
| current_box = boxes[current] | |
| still = [] | |
| for other in remaining: | |
| other_box = boxes[other] | |
| limit = 0.6 if current_box[0] == other_box[0] else 0.98 | |
| if _iou(current_box, other_box) < limit: | |
| still.append(other) | |
| remaining = still | |
| return kept | |
| def _contained(inner: np.ndarray, outer: np.ndarray) -> bool: | |
| x1 = max(inner[2], outer[2]) | |
| y1 = max(inner[3], outer[3]) | |
| x2 = min(inner[4], outer[4]) | |
| y2 = min(inner[5], outer[5]) | |
| area = max(0.0, x2 - x1) * max(0.0, y2 - y1) | |
| inner_area = (inner[4] - inner[2]) * (inner[5] - inner[3]) | |
| return bool(inner_area > 0 and area / inner_area >= 0.9) | |
| def _drop_contained(boxes: np.ndarray) -> np.ndarray: | |
| contained = np.zeros(len(boxes), dtype=bool) | |
| for class_id in _KEEP_LARGE: | |
| for index, box in enumerate(boxes): | |
| for other_index, other in enumerate(boxes): | |
| if other_index == index or other[0] != class_id: | |
| continue | |
| if _contained(box, other): | |
| contained[index] = True | |
| break | |
| return boxes[~contained] | |
| def _drop_overlapping(regions: List[dict]) -> List[dict]: | |
| regions = [region for region in regions if region["label"] != "reference"] | |
| if not regions: | |
| return regions | |
| coords = np.array([region["bbox"] for region in regions], dtype=np.float64) | |
| widths = coords[:, 2] - coords[:, 0] | |
| heights = coords[:, 3] - coords[:, 1] | |
| areas = np.maximum(widths, 0) * np.maximum(heights, 0) | |
| dropped = set() | |
| for index, region in enumerate(regions): | |
| if widths[index] < 6 or heights[index] < 6: | |
| dropped.add(index) | |
| for left in range(len(regions)): | |
| for right in range(left + 1, len(regions)): | |
| if left in dropped or right in dropped: | |
| continue | |
| x1 = max(coords[left, 0], coords[right, 0]) | |
| y1 = max(coords[left, 1], coords[right, 1]) | |
| x2 = min(coords[left, 2], coords[right, 2]) | |
| y2 = min(coords[left, 3], coords[right, 3]) | |
| intersection = max(0.0, x2 - x1) * max(0.0, y2 - y1) | |
| smaller = min(areas[left], areas[right]) | |
| ratio = intersection / smaller if smaller > 0 else 0.0 | |
| labels = {regions[left]["label"], regions[right]["label"]} | |
| if "inline_formula" in labels and ratio > 0.5: | |
| if regions[left]["label"] == "inline_formula": | |
| dropped.add(left) | |
| if regions[right]["label"] == "inline_formula": | |
| dropped.add(right) | |
| continue | |
| if ratio > 0.7: | |
| dropped.add(right if areas[left] >= areas[right] else left) | |
| return [region for index, region in enumerate(regions) if index not in dropped] | |
| def detect(model, image) -> List[tuple]: | |
| import cv2 | |
| rgb = np.asarray(image.convert("RGB")) | |
| height, width = rgb.shape[:2] | |
| bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) | |
| resized = cv2.resize(bgr, (800, 800), interpolation=cv2.INTER_CUBIC) | |
| channels = [(channel.astype(np.float32) / 255.0) for channel in cv2.split(resized)] | |
| array = cv2.merge(channels).transpose(2, 0, 1) | |
| pixel_values = torch.from_numpy(array).unsqueeze(0).to("cuda") | |
| with torch.inference_mode(): | |
| outputs = model(pixel_values=pixel_values) | |
| logits = outputs.logits.detach().float().cpu().numpy()[0] | |
| pred_boxes = outputs.pred_boxes.detach().float().cpu().numpy()[0] | |
| order_logits = outputs.order_logits.detach().float().cpu().numpy()[0] | |
| boxes = _boxes_7col(logits, pred_boxes, order_logits, width, height) | |
| boxes[:, 2:6] = np.round(boxes[:, 2:6]) | |
| boxes = boxes[(boxes[:, 1] > 0.3) & (boxes[:, 0] > -1)] | |
| if len(boxes) == 0: | |
| return [] | |
| boxes = boxes[_nms(boxes[:, :6])] | |
| if len(boxes) > 1: | |
| image_id = LABELS.index("image") | |
| area_limit = 0.82 if width > height else 0.93 | |
| page_area = width * height | |
| kept = [] | |
| for box in boxes: | |
| if box[0] != image_id: | |
| kept.append(box) | |
| continue | |
| x1, y1 = max(0, int(box[2])), max(0, int(box[3])) | |
| x2, y2 = min(width, int(box[4])), min(height, int(box[5])) | |
| if (x2 - x1) * (y2 - y1) <= area_limit * page_area: | |
| kept.append(box) | |
| boxes = np.array(kept, dtype=boxes.dtype) if kept else boxes[:0] | |
| if len(boxes) >= 2: | |
| boxes = _drop_contained(boxes) | |
| if len(boxes) == 0: | |
| return [] | |
| boxes = boxes[np.argsort(boxes[:, 6])] | |
| regions = [] | |
| for box in boxes: | |
| x1 = int(max(0, box[2])) | |
| y1 = int(max(0, box[3])) | |
| x2 = int(min(width, box[4])) | |
| y2 = int(min(height, box[5])) | |
| if x2 <= x1 or y2 <= y1: | |
| continue | |
| class_id = int(box[0]) | |
| regions.append({"label": LABELS[class_id], "bbox": [x1, y1, x2, y2]}) | |
| return [(region["label"], region["bbox"]) for region in _drop_overlapping(regions)] | |