Draw matched content stronger than expanded context in visualizations
This commit is contained in:
Yiorgis Gozadinos
parent
494f774473
commit
6d86237dd6
4 changed files with 97 additions and 11 deletions
|
|
@ -18,6 +18,7 @@
|
|||
- `update_document` accepts a `uri` argument to change a document's URI.
|
||||
- `SearchResult.chunk_ids` and `Citation.chunk_ids` carry the chunk ids merged into an expanded result.
|
||||
- `visualize_chunk` accepts multiple chunks and reproduces merged-result expansion; chat, inspector, and app visualizations pass all cited chunks.
|
||||
- Chunk visualizations draw matched content in a stronger highlight than expanded context.
|
||||
- docling-serve requests fail over to another instance on transport/5xx errors and skip instances whose circuit breaker is open; tune via `providers.docling_serve.max_attempts` and `providers.docling_serve.circuit_breaker`.
|
||||
|
||||
### Fixed
|
||||
|
|
|
|||
|
|
@ -246,7 +246,7 @@ Display visual grounding for a chunk - shows page images with highlighted boundi
|
|||
haiku-rag visualize <chunk_id>
|
||||
```
|
||||
|
||||
This renders the source document pages with the chunk's location highlighted. Useful for verifying chunk boundaries and understanding document structure.
|
||||
This renders the source document pages with the chunk's location highlighted. The chunk itself draws in a strong highlight, while surrounding context swept in by expansion draws fainter. Useful for verifying chunk boundaries and understanding document structure.
|
||||
|
||||
!!! note
|
||||
Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored.
|
||||
|
|
|
|||
|
|
@ -234,6 +234,10 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]")
|
|||
reproduces the merged expansion; chunks from a different document than
|
||||
the first are ignored.
|
||||
|
||||
The chunks' own items draw in a strong highlight; items swept in by
|
||||
expansion draw fainter, so the matched content stands out from its
|
||||
surrounding context.
|
||||
|
||||
Returns a list of PIL Image objects, one per page with bounding boxes.
|
||||
Empty list if no bounding boxes or page images available.
|
||||
"""
|
||||
|
|
@ -279,19 +283,29 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]")
|
|||
refs.extend(r for r in result.doc_item_refs if r not in refs)
|
||||
if not refs:
|
||||
refs = [r for sr in search_results for r in sr.doc_item_refs]
|
||||
meta = ChunkMetadata(doc_item_refs=refs)
|
||||
else:
|
||||
meta = chunks[0].get_chunk_metadata()
|
||||
bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
|
||||
if not bounding_boxes:
|
||||
refs = chunks[0].get_chunk_metadata().doc_item_refs
|
||||
|
||||
matched_refs = {r for sr in search_results for r in sr.doc_item_refs} or set(refs)
|
||||
swept_refs = [r for r in refs if r not in matched_refs]
|
||||
|
||||
matched_boxes = ChunkMetadata(
|
||||
doc_item_refs=sorted(matched_refs)
|
||||
).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
|
||||
# Group bounding boxes by page; swept boxes first so matched draw on top
|
||||
boxes_by_page: dict[int, list] = {}
|
||||
for bbox in bounding_boxes:
|
||||
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)
|
||||
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)
|
||||
|
|
@ -319,7 +333,7 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]")
|
|||
image = deepcopy(pil_image)
|
||||
draw = ImageDraw.Draw(image, "RGBA")
|
||||
|
||||
for bbox in boxes_by_page[page_no]:
|
||||
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
|
||||
|
|
@ -329,8 +343,12 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]")
|
|||
if y0 > y1:
|
||||
y0, y1 = y1, y0
|
||||
|
||||
fill_color = (255, 255, 0, 40) # Yellow with transparency
|
||||
outline_color = (255, 165, 0, 100) # Orange outline
|
||||
if is_matched:
|
||||
fill_color = (255, 255, 0, 40) # Yellow with transparency
|
||||
outline_color = (255, 165, 0, 100) # Orange outline
|
||||
else:
|
||||
fill_color = (255, 255, 0, 15)
|
||||
outline_color = (255, 165, 0, 50)
|
||||
|
||||
draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None)
|
||||
draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=1)
|
||||
|
|
|
|||
|
|
@ -1666,6 +1666,73 @@ async def test_client_visualize_chunk_merged_chunks_union_pages(temp_db_path):
|
|||
assert len(merged_images) == 2
|
||||
|
||||
|
||||
async def test_client_visualize_chunk_two_tone_highlights(temp_db_path):
|
||||
"""Matched content draws stronger than context swept in by expansion."""
|
||||
from docling_core.types.doc.base import BoundingBox, Size
|
||||
from docling_core.types.doc.document import (
|
||||
DoclingDocument,
|
||||
ImageRef,
|
||||
ProvenanceItem,
|
||||
)
|
||||
from docling_core.types.doc.labels import DocItemLabel
|
||||
from PIL import Image as PilImageModule
|
||||
|
||||
docling_doc = DoclingDocument(name="two-tone-test")
|
||||
page_size = Size(width=612.0, height=792.0)
|
||||
docling_doc.add_page(
|
||||
page_no=1,
|
||||
size=page_size,
|
||||
image=ImageRef.from_pil(
|
||||
PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
|
||||
),
|
||||
)
|
||||
|
||||
# Three small paragraphs; the chunk matches only the middle one, so
|
||||
# expansion sweeps in its neighbors.
|
||||
for i in range(3):
|
||||
docling_doc.add_text(
|
||||
label=DocItemLabel.PARAGRAPH,
|
||||
text=f"Paragraph {i}.",
|
||||
prov=ProvenanceItem(
|
||||
page_no=1,
|
||||
bbox=BoundingBox(l=50, t=700 - i * 100, r=550, b=650 - i * 100),
|
||||
charspan=(0, 12),
|
||||
),
|
||||
)
|
||||
|
||||
chunks = [
|
||||
Chunk(
|
||||
content="Paragraph 1.",
|
||||
metadata={
|
||||
"doc_item_refs": ["#/texts/1"],
|
||||
"page_numbers": [1],
|
||||
"labels": ["paragraph"],
|
||||
},
|
||||
order=0,
|
||||
embedding=[0.1] * 2560,
|
||||
)
|
||||
]
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.import_document(docling_doc, chunks, uri="test://two-tone")
|
||||
stored_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
assert len(stored_chunks) == 1
|
||||
|
||||
images = await client.visualize_chunk(stored_chunks[0])
|
||||
assert len(images) == 1
|
||||
image = images[0]
|
||||
|
||||
# Page dpi 72 == document coords, bottom-left origin flipped to
|
||||
# top-left: item i's box spans y = 92 + i * 100 .. 142 + i * 100.
|
||||
matched = image.getpixel((300, 217)) # inside #/texts/1
|
||||
swept = image.getpixel((300, 117)) # inside #/texts/0
|
||||
background = image.getpixel((300, 30)) # outside all boxes
|
||||
|
||||
assert matched != background
|
||||
assert swept != background
|
||||
assert matched != swept
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# convert() method tests
|
||||
# =============================================================================
|
||||
|
|
|
|||
Loading…
Reference in a new issue