fdb5710491
Chunk 2 gave search a configured set to fan out over. ask and analyze
covered one database still: the RAG capability had no way to be told which
databases a question spanned, and the analysis sandbox mounted one
document tree.
The selection travels as sources on EvidenceState, beside the filter it
scopes with, so both capabilities read it the same way. clients_covering
is the one rule that turns a selection into clients, used by search, the
sandbox mount and the cite fallback, so a question scoped to some
databases cannot search, mount or cite another. Citations carry the
database they came from, and format_for_agent names it, so the model can
attribute evidence while it answers rather than only afterwards.
The sandbox keeps one flat /documents/{id}/ namespace and resolves each id
to the client holding it, which rests on ids being UUID4. A database
copied from another breaks that, so an id held twice is refused rather
than resolved to whichever arrived last.
On the CLI, search, ask and analyze cover the configured set and label
each result with its database. Every other command works on one, named
with --database NAME (a name reaches a database behind a URI, which --db
cannot) or --db PATH, and refuses a set it cannot choose from instead of
silently reading the default database. Cold databases open together, so a
first query costs the slowest open rather than their sum.
623 lines
23 KiB
Python
623 lines
23 KiB
Python
import asyncio
|
|
import base64
|
|
from collections.abc import Sequence
|
|
from typing import TYPE_CHECKING
|
|
|
|
from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType
|
|
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
|
|
|
|
if TYPE_CHECKING:
|
|
from PIL import Image as PILImage
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
|
|
async def search(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
limit: int | None = None,
|
|
search_type: SearchType | None = None,
|
|
filter: str | None = None,
|
|
include_images: bool = True,
|
|
) -> list[SearchResult]:
|
|
"""Search for relevant chunks with optional reranking.
|
|
|
|
Args:
|
|
client: The HaikuRAG client (provides config + chunk repository).
|
|
query: Text (``str``) or image (``bytes`` / ``PIL.Image.Image``).
|
|
Image queries require a multimodal embedder and run vector-only.
|
|
limit: Maximum number of results to return. Defaults to config.search.limit.
|
|
search_type: "vector", "fts", or "hybrid".
|
|
Applicable only for text queries, where the default is "hybrid".
|
|
filter: Optional SQL WHERE clause to filter documents before searching chunks.
|
|
include_images: When True, populate ``SearchResult.image_data`` with
|
|
base64 picture bytes for picture-labeled chunks.
|
|
|
|
Returns:
|
|
List of SearchResult objects ordered by relevance.
|
|
"""
|
|
if limit is None:
|
|
limit = client._config.search.limit
|
|
|
|
candidates = await _fetch(
|
|
client, query, _fetch_limit(client, query, limit), search_type, filter
|
|
)
|
|
chunk_results = await _rank(client, query, candidates, limit)
|
|
|
|
results = [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results]
|
|
results = _dedup_picture_chunks(results)
|
|
|
|
if include_images:
|
|
await _populate_image_data(client, results)
|
|
|
|
return results
|
|
|
|
|
|
async def search_sources(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
limit: int | None = None,
|
|
search_type: SearchType | None = None,
|
|
filter: str | None = None,
|
|
include_images: bool = True,
|
|
sources: list[str] | None = None,
|
|
) -> list[SearchResult]:
|
|
"""Search several databases and fuse their results into one ranked list.
|
|
|
|
Fetch, fuse, truncate, then enrich: enrichment runs on the survivors through
|
|
the database each came from, so it costs the same as a single-database search
|
|
rather than multiplying by the number searched.
|
|
"""
|
|
if limit is None:
|
|
limit = client._config.search.limit
|
|
|
|
names = list(client._federated) if sources is None else list(sources)
|
|
if not names:
|
|
return []
|
|
selected = await client.clients_for(names)
|
|
|
|
# One over-fetch decision, and one reranker, for the whole set.
|
|
fetch_limit = _fetch_limit(client, query, limit)
|
|
per_source = await asyncio.gather(
|
|
*(_fetch(c, query, fetch_limit, search_type, filter) for c in selected)
|
|
)
|
|
|
|
ranked = await _fuse(client, selected, query, per_source, limit)
|
|
|
|
results: list[SearchResult] = []
|
|
for owner, chunk, score in ranked:
|
|
result = SearchResult.from_chunk(chunk, score)
|
|
result.source = owner._source
|
|
results.append(result)
|
|
results = _dedup_picture_chunks(results)
|
|
|
|
if include_images:
|
|
by_owner: dict[str, list[SearchResult]] = {}
|
|
for result in results:
|
|
if result.source:
|
|
by_owner.setdefault(result.source, []).append(result)
|
|
await asyncio.gather(
|
|
*(
|
|
_populate_image_data(client._clients[name], owned)
|
|
for name, owned in by_owner.items()
|
|
)
|
|
)
|
|
|
|
return results
|
|
|
|
|
|
async def _fuse(
|
|
federator: "HaikuRAG",
|
|
clients: list["HaikuRAG"],
|
|
query: "str | bytes | PILImage.Image",
|
|
per_source: list[list[tuple[Chunk, float]]],
|
|
limit: int,
|
|
) -> list[tuple["HaikuRAG", Chunk, float]]:
|
|
"""One ranked list from several, keeping each candidate's owner.
|
|
|
|
A configured reranker scores the union directly, which is what makes ranking
|
|
across databases tractable: it compares query against document and does not
|
|
care where a candidate came from. Without one, reciprocal rank fusion over the
|
|
per-database rankings, since scores from separate indexes are not comparable.
|
|
"""
|
|
owned = [
|
|
(client, chunk, score)
|
|
for client, candidates in zip(clients, per_source, strict=True)
|
|
for chunk, score in candidates
|
|
]
|
|
if not owned:
|
|
return []
|
|
|
|
# An image query has no text for a reranker to score against, and the check
|
|
# precedes `reranker`, which builds the reranker on first access.
|
|
if isinstance(query, str):
|
|
reranker = federator.reranker
|
|
if reranker is not None:
|
|
chunks = [chunk for _, chunk, _ in owned]
|
|
if federator._config.reranking.multimodal:
|
|
await asyncio.gather(
|
|
*(
|
|
_attach_picture_data(
|
|
c, [chunk for owner, chunk, _ in owned if owner is c]
|
|
)
|
|
for c in clients
|
|
)
|
|
)
|
|
reranked = await reranker.rerank(query, chunks, top_n=limit)
|
|
owner_of = {id(chunk): client for client, chunk, _ in owned}
|
|
return [(owner_of[id(chunk)], chunk, score) for chunk, score in reranked]
|
|
|
|
scored: list[tuple[float, HaikuRAG, Chunk]] = []
|
|
for client, candidates in zip(clients, per_source, strict=True):
|
|
for rank, (chunk, _) in enumerate(candidates):
|
|
scored.append((1.0 / (_RRF_K + rank + 1), client, chunk))
|
|
scored.sort(key=lambda item: item[0], reverse=True)
|
|
return [(client, chunk, score) for score, client, chunk in scored[:limit]]
|
|
|
|
|
|
# Reciprocal rank fusion's smoothing constant, the value the literature uses.
|
|
_RRF_K = 60
|
|
|
|
|
|
# Candidates per requested result when a reranker will re-order them.
|
|
_RERANK_OVERFETCH = 10
|
|
|
|
|
|
def _fetch_limit(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
limit: int,
|
|
) -> int:
|
|
"""How many candidates to fetch per database.
|
|
|
|
Only a text query with a reranker over-fetches: an image query keeps its
|
|
vector ranking, and the type is checked before `reranker`, which loads model
|
|
weights for a local one on first access.
|
|
"""
|
|
if not isinstance(query, str):
|
|
return limit
|
|
return limit * _RERANK_OVERFETCH if client.reranker else limit
|
|
|
|
|
|
async def _fetch(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
limit: int,
|
|
search_type: SearchType | None,
|
|
filter: str | None,
|
|
) -> list[tuple[Chunk, float]]:
|
|
"""Candidates from one database, ranked by that database.
|
|
|
|
`limit` is how many to fetch, already including any over-fetch the caller
|
|
wants. Deciding that here would have each database consult its own reranker,
|
|
and a local reranker loads model weights per instance.
|
|
"""
|
|
if isinstance(query, str):
|
|
if search_type is None:
|
|
search_type = "hybrid"
|
|
return await client.chunk_repository.search(query, limit, search_type, filter)
|
|
|
|
embedder = client.embedder
|
|
if not embedder.supports_images:
|
|
raise ValueError(
|
|
"Image queries require a multimodal embedder. Set "
|
|
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere "
|
|
"model."
|
|
)
|
|
query_vector = await embedder.embed_image(query)
|
|
return await client.chunk_repository.search(
|
|
query="",
|
|
limit=limit,
|
|
filter=filter,
|
|
query_vector=query_vector,
|
|
)
|
|
|
|
|
|
async def _rank(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
candidates: list[tuple[Chunk, float]],
|
|
limit: int,
|
|
) -> list[tuple[Chunk, float]]:
|
|
"""Order candidates and cut them to `limit`.
|
|
|
|
An image query carries no text for a reranker to score against, so its
|
|
candidates keep the vector ranking. Its type is checked before
|
|
`client.reranker`, which builds the reranker on first access and loads model
|
|
weights for a local one.
|
|
"""
|
|
if not isinstance(query, str):
|
|
return candidates[:limit]
|
|
|
|
reranker = client.reranker
|
|
if reranker is None:
|
|
return candidates[:limit]
|
|
|
|
chunks = [chunk for chunk, _ in candidates]
|
|
if client._config.reranking.multimodal:
|
|
await _attach_picture_data(client, chunks)
|
|
return await reranker.rerank(query, chunks, top_n=limit)
|
|
|
|
|
|
async def _attach_picture_data(client: "HaikuRAG", chunks: list[Chunk]) -> None:
|
|
"""Attach picture bytes to synthetic picture chunks in-place, so a
|
|
multimodal reranker can score the pixels instead of just the chunk's
|
|
description text.
|
|
|
|
One query however many documents the candidates span, which matters here
|
|
more than anywhere: reranking fetches `limit * 10` candidates.
|
|
"""
|
|
by_doc: dict[str, list[tuple[Chunk, str]]] = {}
|
|
for chunk in chunks:
|
|
if chunk.document_id is None:
|
|
continue
|
|
refs = chunk.get_chunk_metadata().doc_item_refs
|
|
if len(refs) == 1 and refs[0].startswith(PICTURE_REF_PREFIX):
|
|
by_doc.setdefault(chunk.document_id, []).append((chunk, refs[0]))
|
|
|
|
bytes_by_document, _ = await client.document_item_repository.get_pictures_grouped(
|
|
{doc_id: [ref for _, ref in pairs] for doc_id, pairs in by_doc.items()}
|
|
)
|
|
for doc_id, doc_chunks in by_doc.items():
|
|
bytes_by_ref = bytes_by_document.get(doc_id, {})
|
|
for chunk, ref in doc_chunks:
|
|
data = bytes_by_ref.get(ref)
|
|
if data:
|
|
chunk._picture_data = data
|
|
|
|
|
|
def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]:
|
|
"""Collapse duplicate picture-only chunks to one result per ``self_ref``.
|
|
|
|
Keyed by database as well, since a database copied from another holds the same
|
|
document id: collapsing across them would drop one of two real results.
|
|
|
|
A single picture can produce two chunks for the same self_ref: one whose
|
|
vector is the text embedding of the picture's description, and one whose
|
|
vector is the image embedding of the picture's bytes. Both can rank for
|
|
the same query. When two results share a single picture self_ref as
|
|
their only ref, keep the higher-scoring one. Wider chunks that span the
|
|
picture plus surrounding items pass through untouched.
|
|
"""
|
|
seen: dict[tuple[str | None, str | None, str], int] = {}
|
|
keep: list[bool] = [True] * len(results)
|
|
for i, r in enumerate(results):
|
|
if len(r.doc_item_refs) == 1 and r.doc_item_refs[0].startswith(
|
|
PICTURE_REF_PREFIX
|
|
):
|
|
key = (r.source, r.document_id, r.doc_item_refs[0])
|
|
prior = seen.get(key)
|
|
if prior is None:
|
|
seen[key] = i
|
|
elif r.score > results[prior].score:
|
|
keep[prior] = False
|
|
seen[key] = i
|
|
else:
|
|
keep[i] = False
|
|
return [r for r, k in zip(results, keep) if k]
|
|
|
|
|
|
async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) -> None:
|
|
"""Attach base64 picture bytes to ``SearchResult.image_data`` in-place.
|
|
|
|
A result carries a picture when its refs include the picture directly, or
|
|
when they include the picture's caption — the common case where a prose
|
|
chunk carrying a figure's caption ranks while the picture is its own chunk.
|
|
Costs a fixed number of reads however many documents the result set spans.
|
|
"""
|
|
repo = client.document_item_repository
|
|
by_doc: dict[str, list[SearchResult]] = {}
|
|
for r in results:
|
|
if r.document_id and r.doc_item_refs:
|
|
by_doc.setdefault(r.document_id, []).append(r)
|
|
if not by_doc:
|
|
return
|
|
|
|
refs_by_document = {
|
|
doc_id: list({ref for r in doc_results for ref in r.doc_item_refs})
|
|
for doc_id, doc_results in by_doc.items()
|
|
}
|
|
captions_to_pictures = await repo.get_caption_picture_refs_grouped(refs_by_document)
|
|
|
|
# Which pictures each result wants, and which to fetch per document.
|
|
result_pictures: list[tuple[SearchResult, list[str]]] = []
|
|
wanted: dict[str, list[str]] = {}
|
|
for doc_id, doc_results in by_doc.items():
|
|
caption_to_picture = captions_to_pictures.get(doc_id, {})
|
|
seen: set[str] = set()
|
|
for r in doc_results:
|
|
pictures: list[str] = []
|
|
for ref in r.doc_item_refs:
|
|
picture = (
|
|
ref
|
|
if ref.startswith(PICTURE_REF_PREFIX)
|
|
else caption_to_picture.get(ref)
|
|
)
|
|
if picture and picture not in pictures:
|
|
pictures.append(picture)
|
|
if pictures:
|
|
result_pictures.append((r, pictures))
|
|
for picture in pictures:
|
|
if picture not in seen:
|
|
wanted.setdefault(doc_id, []).append(picture)
|
|
seen.add(picture)
|
|
if not wanted:
|
|
return
|
|
|
|
bytes_by_document, captions_by_document = await repo.get_pictures_grouped(
|
|
wanted, with_text=True
|
|
)
|
|
if not bytes_by_document:
|
|
return
|
|
|
|
for r, pictures in result_pictures:
|
|
bytes_by_ref = bytes_by_document.get(r.document_id or "", {})
|
|
captions_by_ref = captions_by_document.get(r.document_id or "", {})
|
|
attached: dict[str, str] = {}
|
|
captions: dict[str, str] = {}
|
|
for ref in pictures:
|
|
blob = bytes_by_ref.get(ref)
|
|
if blob:
|
|
attached[ref] = base64.b64encode(blob).decode("ascii")
|
|
caption = captions_by_ref.get(ref)
|
|
if caption:
|
|
captions[ref] = caption
|
|
if attached:
|
|
r.image_data = attached
|
|
if captions:
|
|
r.picture_captions = captions
|
|
|
|
|
|
async def expand_context(
|
|
client: "HaikuRAG",
|
|
search_results: list[SearchResult],
|
|
) -> list[SearchResult]:
|
|
"""Expand search results with surrounding content from the document.
|
|
|
|
Uses the document_items table for section-bounded expansion.
|
|
See haiku.rag.context for the algorithm description.
|
|
|
|
Results without doc_item_refs pass through unexpanded. This happens when
|
|
chunks were created without docling metadata (e.g., custom chunks passed
|
|
to import_document).
|
|
"""
|
|
# A federating client has no repositories of its own, so each result expands
|
|
# through the database it came from.
|
|
if client._federated:
|
|
by_source: dict[str, list[SearchResult]] = {}
|
|
unsourced: list[SearchResult] = []
|
|
for result in search_results:
|
|
if result.source:
|
|
by_source.setdefault(result.source, []).append(result)
|
|
else:
|
|
unsourced.append(result)
|
|
owners = await client.clients_for(list(by_source))
|
|
expanded_groups = await asyncio.gather(
|
|
*(
|
|
expand_context(owner, by_source[owner._source])
|
|
for owner in owners
|
|
if owner._source
|
|
)
|
|
)
|
|
merged = unsourced + [r for group in expanded_groups for r in group]
|
|
# Grouping by database must not become the tiebreak: fused scores tie
|
|
# often, so equal scores keep the order they were fused in.
|
|
arrival = {
|
|
result.chunk_id: rank
|
|
for rank, result in enumerate(search_results)
|
|
if result.chunk_id
|
|
}
|
|
|
|
def fused_rank(result: SearchResult) -> int:
|
|
return min(
|
|
(
|
|
arrival[cid]
|
|
for cid in (result.chunk_id, *result.chunk_ids)
|
|
if cid in arrival
|
|
),
|
|
default=len(arrival),
|
|
)
|
|
|
|
merged.sort(key=lambda r: (-r.score, fused_rank(r)))
|
|
return merged
|
|
|
|
from haiku.rag.context import expand_with_items, window_for
|
|
|
|
max_chars = client._config.search.max_context_chars
|
|
|
|
# Group by document_id for efficient processing
|
|
document_groups: dict[str | None, list[SearchResult]] = {}
|
|
for result in search_results:
|
|
doc_id = result.document_id
|
|
if doc_id not in document_groups:
|
|
document_groups[doc_id] = []
|
|
document_groups[doc_id].append(result)
|
|
|
|
expanded_results = []
|
|
expandable = {
|
|
doc_id: doc_results
|
|
for doc_id, doc_results in document_groups.items()
|
|
if doc_id is not None and any(r.doc_item_refs for r in doc_results)
|
|
}
|
|
repo = client.document_item_repository
|
|
positions_by_document = await repo.resolve_refs_grouped(
|
|
{
|
|
doc_id: [ref for r in doc_results for ref in r.doc_item_refs]
|
|
for doc_id, doc_results in expandable.items()
|
|
}
|
|
)
|
|
windows = {
|
|
doc_id: window_for(positions)
|
|
for doc_id, positions in positions_by_document.items()
|
|
if positions
|
|
}
|
|
items_by_document = await repo.get_items_in_ranges(windows)
|
|
|
|
# In document_groups order: the score sort below is stable, so assembling
|
|
# expandable and passthrough documents in separate passes would reorder
|
|
# equal-scored results.
|
|
for doc_id, doc_results in document_groups.items():
|
|
if doc_id not in expandable:
|
|
expanded_results.extend(doc_results)
|
|
continue
|
|
expanded_results.extend(
|
|
expand_with_items(
|
|
doc_results,
|
|
max_chars,
|
|
positions_by_document.get(doc_id, {}),
|
|
items_by_document.get(doc_id, []),
|
|
)
|
|
)
|
|
|
|
expanded_results.sort(key=lambda r: r.score, reverse=True)
|
|
# image_data and picture_captions are preserved through expansion by
|
|
# expand_with_items — we deliberately do not re-attach bytes for refs
|
|
# introduced by section expansion, so the multimodal payload stays
|
|
# bounded by what was originally retrieved.
|
|
return expanded_results
|
|
|
|
|
|
async def visualize_chunk(
|
|
client: "HaikuRAG",
|
|
chunk: "Chunk | Sequence[Chunk]",
|
|
refs: list[str] | None = None,
|
|
expand: bool = True,
|
|
) -> list:
|
|
"""Render page images with bounding box highlights for one or more chunks.
|
|
|
|
When ``refs`` is given (the ``doc_item_refs`` of the citation, i.e. the
|
|
exact items the model saw), bounding boxes are resolved from them directly
|
|
so the visualization matches the cited context precisely. Otherwise, with
|
|
``expand=True`` (default) the chunks' context is re-expanded to recover the
|
|
surrounding section; with ``expand=False`` only the chunks' own items are
|
|
drawn, so the visualization shows just the retrieved chunk with no context.
|
|
|
|
The chunks' own items draw in a strong highlight; the remaining items draw
|
|
fainter, so the matched content stands out from its surrounding context.
|
|
Chunks from a different document than the first are ignored.
|
|
|
|
Returns a list of PIL Image objects, one per page with bounding boxes.
|
|
Empty list if no bounding boxes or page images available.
|
|
"""
|
|
from copy import deepcopy
|
|
|
|
from PIL import ImageDraw
|
|
|
|
from haiku.rag.store.models.chunk import ChunkMetadata
|
|
|
|
chunks = [chunk] if isinstance(chunk, Chunk) else list(chunk)
|
|
if not chunks:
|
|
return []
|
|
document_id = chunks[0].document_id
|
|
if not document_id:
|
|
return []
|
|
chunks = [c for c in chunks if c.document_id == document_id]
|
|
|
|
doc = await client.document_repository.get_docling_data(document_id)
|
|
if not doc:
|
|
return []
|
|
|
|
docling_doc = doc.get_docling_document()
|
|
if not docling_doc:
|
|
return []
|
|
|
|
matched_refs = {r for c in chunks for r in c.get_chunk_metadata().doc_item_refs}
|
|
|
|
if refs is not None:
|
|
all_refs = list(refs)
|
|
elif not expand:
|
|
# Chunk-only: draw just the retrieved chunks' own items, no context.
|
|
all_refs = list(matched_refs)
|
|
else:
|
|
# No stored context: re-expand the chunks to recover their section.
|
|
search_results = [
|
|
SearchResult(
|
|
content=c.content,
|
|
score=1.0,
|
|
chunk_id=c.id,
|
|
document_id=c.document_id,
|
|
doc_item_refs=meta.doc_item_refs,
|
|
page_numbers=meta.page_numbers,
|
|
)
|
|
for c in chunks
|
|
if (meta := c.get_chunk_metadata()).doc_item_refs
|
|
]
|
|
if search_results:
|
|
expanded = await expand_context(client, search_results)
|
|
all_refs = []
|
|
for result in expanded:
|
|
all_refs.extend(r for r in result.doc_item_refs if r not in all_refs)
|
|
if not all_refs:
|
|
all_refs = [r for sr in search_results for r in sr.doc_item_refs]
|
|
else:
|
|
all_refs = list(chunks[0].get_chunk_metadata().doc_item_refs)
|
|
|
|
matched_draw = [r for r in all_refs if r in matched_refs]
|
|
swept_refs = [r for r in all_refs if r not in matched_refs]
|
|
|
|
matched_boxes = ChunkMetadata(doc_item_refs=matched_draw).resolve_bounding_boxes(
|
|
docling_doc
|
|
)
|
|
swept_boxes = ChunkMetadata(doc_item_refs=swept_refs).resolve_bounding_boxes(
|
|
docling_doc
|
|
)
|
|
if not matched_boxes and not swept_boxes:
|
|
return []
|
|
|
|
# Group bounding boxes by page; swept boxes first so matched draw on top
|
|
boxes_by_page: dict[int, list] = {}
|
|
for bbox, is_matched in [(b, False) for b in swept_boxes] + [
|
|
(b, True) for b in matched_boxes
|
|
]:
|
|
if bbox.page_no not in boxes_by_page:
|
|
boxes_by_page[bbox.page_no] = []
|
|
boxes_by_page[bbox.page_no].append((bbox, is_matched))
|
|
|
|
# Load only the needed page images
|
|
pages_doc = await client.document_repository.get_pages_data(document_id)
|
|
if not pages_doc:
|
|
return []
|
|
page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))
|
|
|
|
images = []
|
|
for page_no in sorted(boxes_by_page.keys()):
|
|
if page_no not in page_images:
|
|
continue
|
|
|
|
page = page_images[page_no]
|
|
if page.image is None or page.image.pil_image is None:
|
|
continue
|
|
|
|
pil_image = page.image.pil_image
|
|
page_height = page.size.height
|
|
|
|
# Scale factor: image pixels vs document coordinates
|
|
scale_x = pil_image.width / page.size.width
|
|
scale_y = pil_image.height / page.size.height
|
|
|
|
image = deepcopy(pil_image)
|
|
draw = ImageDraw.Draw(image, "RGBA")
|
|
|
|
for bbox, is_matched in boxes_by_page[page_no]:
|
|
# Document coords are bottom-left origin; PIL uses top-left
|
|
x0 = bbox.left * scale_x
|
|
y0 = (page_height - bbox.top) * scale_y
|
|
x1 = bbox.right * scale_x
|
|
y1 = (page_height - bbox.bottom) * scale_y
|
|
|
|
if y0 > y1:
|
|
y0, y1 = y1, y0
|
|
|
|
if is_matched:
|
|
fill_color = (255, 150, 0, 55) # Orange, matched content
|
|
outline_color = (240, 130, 0, 150) # Orange outline
|
|
else:
|
|
fill_color = (255, 255, 0, 40) # Yellow, surrounding context
|
|
outline_color = (255, 165, 0, 100)
|
|
|
|
draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None)
|
|
draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=1)
|
|
|
|
images.append(image)
|
|
|
|
return images
|