Coverage for src / quber / core / extractors / dual / vision.py: 100%

1"""Visual flow: identify every table by looking at the rendered page. 

2 

3Renders each page and runs the grid locator on it; the located full-anatomy 

4boxes ARE the table count and identity for the document. No Camelot, no 

5detector, no recovery. Pages are processed concurrently under a small 

6semaphore so the per-page vision calls overlap without flooding the API. 

7""" 

8 

9from __future__ import annotations 

10 

11import asyncio 

12import tempfile 

13from pathlib import Path 

14from typing import List 

15 

16from quber.agents.grid_locator import GridLocator 

17from quber.core.extractors.camelot.acquire import render_pages 

18from quber.core.extractors.dual.models import VisualTable 

19 

20 

21async def run_vision_flow( 

22 source: Path, 

23 locator: GridLocator, 

24 dpi: int = 200, 

25 concurrency: int = 4, 

26) -> List[VisualTable]: 

27 """Locate tables on every page of `source` and return them in page order.""" 

28 with tempfile.TemporaryDirectory(prefix="quber-dual-vision-") as tmp: 

29 page_images = await asyncio.to_thread(render_pages, source, dpi, Path(tmp)) 

30 semaphore = asyncio.Semaphore(concurrency) 

31 

32 async def locate_page(page: int, image: Path) -> List[VisualTable]: 

33 async with semaphore: 

34 located = await locator.locate(image, source, page) 

35 return [ 

36 VisualTable( 

37 page=page, 

38 ordinal=t.ordinal, 

39 title=t.title, 

40 region=t.region, 

41 tightened=t.tightened, 

42 ) 

43 for t in located 

44 ] 

45 

46 per_page = await asyncio.gather( 

47 *(locate_page(i, image) for i, image in enumerate(page_images, start=1)) 

48 ) 

49 return [table for page_tables in per_page for table in page_tables]