File size: 4,642 Bytes
7b2177e | 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 | # GENERATED by hub/build_hub_package.py from src/bodhan_genai/ocr/engine/layout.py -- do not edit.
# Vendored so this repo is self-contained: `pip install transformers torch pillow` is the
# whole install. See indic_doc_parser.py for usage.
"""Layout backends: page image -> cleaned, reading-ordered blocks.
The two stages hand off a plain JSON layout, so stage 2 does not care where the layout came
from. :class:`LayoutBackend` makes that an interface rather than a claim.
"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING, Protocol, runtime_checkable
from idp_blocks import clamp_to_page, clean_layout
from idp_types import Block, DedupConfig, LayoutConfig, PageResult
from idp_contract import map_label
if TYPE_CHECKING: # pragma: no cover
from PIL.Image import Image
@runtime_checkable
class LayoutBackend(Protocol):
"""``detect`` must return blocks already cleaned and densely ordered: ``order`` a gap-free
0-based rank. Stage 2 matches transcriptions back by ``order``, so gaps mis-assign text."""
def detect(self, image: Image) -> list[Block]: ...
def close(self) -> None: ...
def _densify(blocks: list[Block]) -> list[Block]:
"""Sort by the detector's reading order and renumber to a gap-free 0-based rank."""
ordered = sorted(blocks, key=lambda b: b.order)
for rank, block in enumerate(ordered):
block.order = rank
return ordered
class IndicDocLayoutBackend:
"""Our finetuned PP-DocLayoutV3 with an integrated reading-order head. Torch only --
constructing this does not load vLLM, which is what lets stage 1 run alone."""
def __init__(
self,
ckpt: str | None = None,
config: LayoutConfig | None = None,
dedup: DedupConfig | None = None,
) -> None:
from idp_model_infer import get_model
self.config = config or LayoutConfig()
self.dedup = dedup or DedupConfig()
if ckpt is None:
raise ValueError(
"no layout weights given -- "
"IndicDocParser.from_pretrained(snapshot_download(REPO))"
)
self.ckpt = ckpt
self.model = get_model(self.ckpt, device=self.config.device)
def detect(self, image: Image) -> list[Block]:
from idp_model_infer import infer
width, height = image.size
detections = infer(
self.model,
image,
conf=self.config.conf,
img_size=self.config.img_size,
device=self.config.device,
)
# The model emits [y0, x0, y1, x1] normalised to 0-1000; the pipeline works in pixel
# [x0, y0, x1, y1]. Axis swap and rescale happen here, once.
blocks = []
for det in detections:
y0, x0, y1, x1 = det["bbox"]
bbox = [x0 / 1000 * width, y0 / 1000 * height, x1 / 1000 * width, y1 / 1000 * height]
label = str(det["label"])
blocks.append(
Block(
order=det["reading_order"],
label=label,
type=map_label(label),
bbox_xyxy=[round(v, 1) for v in clamp_to_page(bbox, width, height)],
conf=round(float(det.get("score", 1.0)), 3),
)
)
return _densify(clean_layout(blocks, self.dedup))
def close(self) -> None:
self.model = None
class JsonLayoutBackend:
"""Replay a layout produced elsewhere -- by stage 1, by hand, or by another detector.
Assumed already clean, so no cleanup runs; blocks are only renumbered, which makes a
hand-edited file usable without fixing ranks. Needs no torch.
"""
def __init__(self, layout: str | dict | PageResult, *, strict: bool = True) -> None:
if isinstance(layout, PageResult):
self.page = layout
else:
if isinstance(layout, str):
with open(layout, encoding="utf-8") as fh:
layout = json.load(fh)
# Validates by default: this is the door a foreign layout comes through, and an
# unrecognised label would otherwise become Text without a word. `strict=False`
# replays a file written before validation existed.
self.page = PageResult.from_record(layout, strict=strict)
def detect(self, image: Image) -> list[Block]:
"""Copies, so renumbering cannot write back into the stored layout. ``image`` is
accepted for interface parity and not read."""
return _densify([b.copy() for b in self.page.blocks])
def close(self) -> None:
return None
|