Introduce visualize_chunk(), refactor inspector to use it
This commit is contained in:
Yiorgis Gozadinos
parent
bfdaa4652a
commit
910d382748
5 changed files with 236 additions and 140 deletions
|
|
@ -870,6 +870,94 @@ class HaikuRAG:
|
|||
qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt)
|
||||
return await qa_agent.answer(question)
|
||||
|
||||
async def visualize_chunk(self, chunk: Chunk) -> list:
|
||||
"""Render page images with bounding box highlights for a chunk.
|
||||
|
||||
Gets the DoclingDocument from the chunk's document, resolves bounding boxes
|
||||
from chunk metadata, and renders all pages that contain bounding boxes with
|
||||
yellow/orange highlight overlays.
|
||||
|
||||
Args:
|
||||
chunk: The chunk to visualize.
|
||||
|
||||
Returns:
|
||||
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
|
||||
|
||||
# Get the document
|
||||
if not chunk.document_id:
|
||||
return []
|
||||
|
||||
doc = await self.document_repository.get_by_id(chunk.document_id)
|
||||
if not doc:
|
||||
return []
|
||||
|
||||
# Get DoclingDocument
|
||||
docling_doc = doc.get_docling_document()
|
||||
if not docling_doc:
|
||||
return []
|
||||
|
||||
# Resolve bounding boxes from chunk metadata
|
||||
chunk_meta = chunk.get_chunk_metadata()
|
||||
bounding_boxes = chunk_meta.resolve_bounding_boxes(docling_doc)
|
||||
if not bounding_boxes:
|
||||
return []
|
||||
|
||||
# Group bounding boxes by page
|
||||
boxes_by_page: dict[int, list] = {}
|
||||
for bbox in bounding_boxes:
|
||||
if bbox.page_no not in boxes_by_page:
|
||||
boxes_by_page[bbox.page_no] = []
|
||||
boxes_by_page[bbox.page_no].append(bbox)
|
||||
|
||||
# Render each page with its bounding boxes
|
||||
images = []
|
||||
for page_no in sorted(boxes_by_page.keys()):
|
||||
if page_no not in docling_doc.pages:
|
||||
continue
|
||||
|
||||
page = docling_doc.pages[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
|
||||
|
||||
# Calculate scale factor (image pixels vs document coordinates)
|
||||
scale_x = pil_image.width / page.size.width
|
||||
scale_y = pil_image.height / page.size.height
|
||||
|
||||
# Draw bounding boxes
|
||||
image = deepcopy(pil_image)
|
||||
draw = ImageDraw.Draw(image, "RGBA")
|
||||
|
||||
for bbox in boxes_by_page[page_no]:
|
||||
# Convert from document coordinates to image coordinates
|
||||
# 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
|
||||
|
||||
# Ensure proper ordering (y0 should be less than y1 for PIL)
|
||||
if y0 > y1:
|
||||
y0, y1 = y1, y0
|
||||
|
||||
# Draw filled rectangle with transparency
|
||||
fill_color = (255, 255, 0, 80) # Yellow with transparency
|
||||
outline_color = (255, 165, 0, 255) # Orange outline
|
||||
|
||||
draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None)
|
||||
draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=3)
|
||||
|
||||
images.append(image)
|
||||
|
||||
return images
|
||||
|
||||
async def rebuild_database(
|
||||
self, mode: RebuildMode = RebuildMode.FULL
|
||||
) -> AsyncGenerator[str, None]:
|
||||
|
|
|
|||
|
|
@ -187,31 +187,10 @@ class InspectorApp(App): # type: ignore[misc]
|
|||
return
|
||||
|
||||
chunk = chunk_list.chunks[idx]
|
||||
if not chunk.document_id:
|
||||
return
|
||||
|
||||
document = await self.client.get_document_by_id(chunk.document_id)
|
||||
if not document:
|
||||
return
|
||||
|
||||
docling_doc = document.get_docling_document()
|
||||
if not docling_doc:
|
||||
self.notify("No DoclingDocument available", severity="warning")
|
||||
return
|
||||
|
||||
meta = chunk.get_chunk_metadata()
|
||||
bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
|
||||
|
||||
from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
|
||||
|
||||
await self.push_screen(
|
||||
VisualGroundingModal(
|
||||
docling_document=docling_doc,
|
||||
bounding_boxes=bounding_boxes,
|
||||
page_numbers=meta.page_numbers,
|
||||
document_uri=chunk.document_uri,
|
||||
)
|
||||
)
|
||||
await self.push_screen(VisualGroundingModal(chunk=chunk, client=self.client))
|
||||
|
||||
|
||||
def run_inspector(db_path: Path | None = None) -> None:
|
||||
|
|
|
|||
|
|
@ -166,39 +166,17 @@ class SearchModal(Screen):
|
|||
"""Show visual grounding for the current chunk."""
|
||||
list_view = self.query_one("#search-results", ListView)
|
||||
idx = list_view.index
|
||||
if idx is None or not self.search_results:
|
||||
if idx is None or not self.chunks:
|
||||
return
|
||||
|
||||
search_result = self.search_results[idx]
|
||||
if not search_result.document_id or not search_result.chunk_id:
|
||||
return
|
||||
|
||||
status_label = self.query_one("#status-label", Static)
|
||||
|
||||
document = await self.client.get_document_by_id(search_result.document_id)
|
||||
if not document:
|
||||
status_label.update("[yellow]Document not found[/yellow]")
|
||||
return
|
||||
|
||||
docling_doc = document.get_docling_document()
|
||||
if not docling_doc:
|
||||
status_label.update(
|
||||
"[yellow]No DoclingDocument available for visual[/yellow]"
|
||||
)
|
||||
return
|
||||
|
||||
chunk = await self.client.chunk_repository.get_by_id(search_result.chunk_id)
|
||||
bounding_boxes = []
|
||||
if chunk:
|
||||
meta = chunk.get_chunk_metadata()
|
||||
bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
|
||||
chunk = self.chunks[idx]
|
||||
|
||||
from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
|
||||
|
||||
modal = VisualGroundingModal(
|
||||
docling_document=docling_doc,
|
||||
bounding_boxes=bounding_boxes,
|
||||
page_numbers=search_result.page_numbers,
|
||||
document_uri=search_result.document_uri,
|
||||
await self.app.push_screen(
|
||||
VisualGroundingModal(
|
||||
chunk=chunk,
|
||||
client=self.client,
|
||||
document_uri=chunk.document_uri,
|
||||
)
|
||||
)
|
||||
await self.app.push_screen(modal)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from textual.app import ComposeResult
|
||||
|
|
@ -9,12 +8,12 @@ from textual.widget import Widget
|
|||
from textual.widgets import Static
|
||||
from textual_image.widget import Image as TextualImage
|
||||
|
||||
from haiku.rag.store.models import BoundingBox
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
from PIL.Image import Image as PILImage
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.store.models import Chunk
|
||||
|
||||
|
||||
class VisualGroundingModal(Screen):
|
||||
"""Modal screen for displaying visual grounding with bounding boxes."""
|
||||
|
|
@ -57,22 +56,21 @@ class VisualGroundingModal(Screen):
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
docling_document: "DoclingDocument",
|
||||
bounding_boxes: list[BoundingBox],
|
||||
page_numbers: list[int],
|
||||
chunk: "Chunk",
|
||||
client: "HaikuRAG",
|
||||
document_uri: str | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.docling_document = docling_document
|
||||
self.document_uri = document_uri
|
||||
self.bounding_boxes = bounding_boxes
|
||||
self.page_numbers = sorted(set(page_numbers)) if page_numbers else []
|
||||
self.chunk = chunk
|
||||
self.client = client
|
||||
self.document_uri = document_uri or chunk.document_uri
|
||||
self.images: list[PILImage] = []
|
||||
self.current_page_idx = 0
|
||||
self._image_widget: Widget = Static("Loading...", id="image-display")
|
||||
self._page_info = Static("", id="page-info")
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
uri_display = self.document_uri or self.docling_document.name or "Document"
|
||||
uri_display = self.document_uri or "Document"
|
||||
with Vertical(id="visual-header"):
|
||||
yield Static(f"[bold]Visual Grounding[/bold] - {uri_display}")
|
||||
with Horizontal(id="visual-content"):
|
||||
|
|
@ -81,96 +79,36 @@ class VisualGroundingModal(Screen):
|
|||
yield self._page_info
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
"""Load and display the first page."""
|
||||
"""Load images and display the first page."""
|
||||
self.images = await self.client.visualize_chunk(self.chunk)
|
||||
await self._render_current_page()
|
||||
|
||||
async def _render_current_page(self) -> None:
|
||||
"""Render the current page with bounding boxes."""
|
||||
if not self.page_numbers:
|
||||
"""Render the current page."""
|
||||
if not self.images:
|
||||
if isinstance(self._image_widget, Static):
|
||||
self._image_widget.update("[red]No page information available[/red]")
|
||||
self._image_widget.update(
|
||||
"[yellow]No page images available[/yellow]\n"
|
||||
"This document was converted without page images."
|
||||
)
|
||||
self._page_info.update("")
|
||||
return
|
||||
|
||||
current_page = self.page_numbers[self.current_page_idx]
|
||||
self._page_info.update(
|
||||
f"Page {current_page} "
|
||||
f"({self.current_page_idx + 1}/{len(self.page_numbers)}) "
|
||||
f"- Use ←/→ to navigate"
|
||||
f"Page {self.current_page_idx + 1}/{len(self.images)} - Use ←/→ to navigate"
|
||||
)
|
||||
|
||||
try:
|
||||
image = self._render_page_with_boxes(current_page)
|
||||
if image:
|
||||
new_widget = TextualImage(image, id="rendered-image")
|
||||
await self._image_widget.remove()
|
||||
content = self.query_one("#visual-content", Horizontal)
|
||||
await content.mount(new_widget)
|
||||
self._image_widget = new_widget
|
||||
elif isinstance(self._image_widget, Static):
|
||||
self._image_widget.update(
|
||||
"[yellow]No page image available[/yellow]\n"
|
||||
"This document was converted without page images."
|
||||
)
|
||||
image = self.images[self.current_page_idx]
|
||||
new_widget = TextualImage(image, id="rendered-image")
|
||||
await self._image_widget.remove()
|
||||
content = self.query_one("#visual-content", Horizontal)
|
||||
await content.mount(new_widget)
|
||||
self._image_widget = new_widget
|
||||
except Exception as e:
|
||||
if isinstance(self._image_widget, Static):
|
||||
self._image_widget.update(f"[red]Error: {e}[/red]")
|
||||
|
||||
def _render_page_with_boxes(self, page_no: int) -> "PILImage | None":
|
||||
"""Render a page from DoclingDocument with bounding boxes."""
|
||||
from PIL import ImageDraw
|
||||
|
||||
# Get the page from DoclingDocument
|
||||
if page_no not in self.docling_document.pages:
|
||||
return None
|
||||
|
||||
page = self.docling_document.pages[page_no]
|
||||
if page.image is None:
|
||||
return None
|
||||
|
||||
pil_image = page.image.pil_image
|
||||
if pil_image is None:
|
||||
return None
|
||||
|
||||
# Get page dimensions
|
||||
page_height = page.size.height
|
||||
|
||||
# Calculate scale factor (image pixels vs document coordinates)
|
||||
scale_x = pil_image.width / page.size.width
|
||||
scale_y = pil_image.height / page.size.height
|
||||
|
||||
# Get bounding boxes for this page
|
||||
page_boxes = [bb for bb in self.bounding_boxes if bb.page_no == page_no]
|
||||
|
||||
if page_boxes:
|
||||
# Draw bounding boxes
|
||||
image = deepcopy(pil_image)
|
||||
draw = ImageDraw.Draw(image, "RGBA")
|
||||
|
||||
for bbox in page_boxes:
|
||||
# Convert from document coordinates to image coordinates
|
||||
# Document coords are typically bottom-left origin
|
||||
# PIL uses top-left origin
|
||||
x0 = bbox.left * scale_x
|
||||
y0 = (page_height - bbox.top) * scale_y # Flip Y
|
||||
x1 = bbox.right * scale_x
|
||||
y1 = (page_height - bbox.bottom) * scale_y # Flip Y
|
||||
|
||||
# Ensure proper ordering (y0 should be less than y1 for PIL)
|
||||
if y0 > y1:
|
||||
y0, y1 = y1, y0
|
||||
|
||||
# Draw filled rectangle with transparency
|
||||
fill_color = (255, 255, 0, 80) # Yellow with transparency
|
||||
outline_color = (255, 165, 0, 255) # Orange outline
|
||||
|
||||
draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None)
|
||||
draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=3)
|
||||
|
||||
return image
|
||||
|
||||
return pil_image
|
||||
|
||||
async def action_dismiss(self, result=None) -> None:
|
||||
self.app.pop_screen()
|
||||
|
||||
|
|
@ -182,6 +120,6 @@ class VisualGroundingModal(Screen):
|
|||
|
||||
async def action_next_page(self) -> None:
|
||||
"""Navigate to the next page."""
|
||||
if self.current_page_idx < len(self.page_numbers) - 1:
|
||||
if self.current_page_idx < len(self.images) - 1:
|
||||
self.current_page_idx += 1
|
||||
await self._render_current_page()
|
||||
|
|
|
|||
|
|
@ -1282,3 +1282,116 @@ async def test_client_file_update_stores_docling_json(temp_db_path):
|
|||
assert doc2.docling_document_json is not None
|
||||
assert doc2.docling_document_json != original_json
|
||||
assert doc2.docling_version == original_version # Version stays same
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_visualize_chunk_no_document(temp_db_path):
|
||||
"""Test visualize_chunk returns empty list when chunk has no document_id."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
chunk = Chunk(content="Orphan chunk", order=0)
|
||||
images = await client.visualize_chunk(chunk)
|
||||
assert images == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_visualize_chunk_no_docling_document(temp_db_path):
|
||||
"""Test visualize_chunk returns empty list when document has no DoclingDocument."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
# Create document with custom chunks (no DoclingDocument)
|
||||
custom_chunks = [Chunk(content="Custom chunk", order=0)]
|
||||
doc = await client.create_document(content="Test content", chunks=custom_chunks)
|
||||
|
||||
assert doc.id is not None
|
||||
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
assert len(chunks) == 1
|
||||
|
||||
images = await client.visualize_chunk(chunks[0])
|
||||
assert images == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_visualize_chunk_no_bounding_boxes(temp_db_path):
|
||||
"""Test visualize_chunk returns empty list when chunk has no bounding boxes."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
# Create document from text (will have DoclingDocument but no page images)
|
||||
doc = await client.create_document(
|
||||
content="Simple text content without structure",
|
||||
uri="test://simple",
|
||||
)
|
||||
|
||||
assert doc.id is not None
|
||||
assert doc.docling_document_json is not None
|
||||
|
||||
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
assert len(chunks) >= 1
|
||||
|
||||
# Text documents converted via markdown won't have page images
|
||||
# so visualize_chunk should return empty list
|
||||
images = await client.visualize_chunk(chunks[0])
|
||||
assert images == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_visualize_chunk_returns_list(temp_db_path):
|
||||
"""Test visualize_chunk returns a list (empty or with images)."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
# Create a structured document
|
||||
markdown_content = """# Chapter 1
|
||||
|
||||
This is paragraph one about topic A.
|
||||
|
||||
This is paragraph two about topic A continued.
|
||||
|
||||
# Chapter 2
|
||||
|
||||
This is paragraph four about topic C.
|
||||
"""
|
||||
doc = await client.create_document(
|
||||
content=markdown_content,
|
||||
uri="test://structured",
|
||||
)
|
||||
|
||||
assert doc.id is not None
|
||||
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
|
||||
# Find a chunk with doc_item_refs
|
||||
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
|
||||
|
||||
if chunks_with_refs:
|
||||
# visualize_chunk should return a list (possibly empty if no page images)
|
||||
images = await client.visualize_chunk(chunks_with_refs[0])
|
||||
assert isinstance(images, list)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_visualize_chunk_with_pdf(temp_db_path):
|
||||
"""Test visualize_chunk returns images with bounding boxes for PDF documents."""
|
||||
from PIL.Image import Image as PILImage
|
||||
|
||||
pdf_path = Path("tests/data/doclaynet.pdf")
|
||||
if not pdf_path.exists():
|
||||
pytest.skip("doclaynet.pdf not found")
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document_from_source(pdf_path)
|
||||
assert isinstance(doc, Document)
|
||||
assert doc.id is not None
|
||||
assert doc.docling_document_json is not None
|
||||
|
||||
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
assert len(chunks) > 0
|
||||
|
||||
# Find a chunk with doc_item_refs (bounding box info)
|
||||
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
|
||||
assert len(chunks_with_refs) > 0, "PDF should have chunks with doc_item_refs"
|
||||
|
||||
# Visualize a chunk - should return images with bounding boxes drawn
|
||||
images = await client.visualize_chunk(chunks_with_refs[0])
|
||||
|
||||
assert isinstance(images, list)
|
||||
assert len(images) > 0, "PDF with page images should return visualizations"
|
||||
|
||||
# Verify returned objects are PIL Images
|
||||
for img in images:
|
||||
assert isinstance(img, PILImage)
|
||||
|
|
|
|||
Loading…
Reference in a new issue