"""Define classes for layout visualization.""" from typing import Literal, Optional, Union from PIL import Image as PILImage from PIL import ImageDraw, ImageFont from PIL.Image import Image from PIL.ImageFont import FreeTypeFont from pydantic import BaseModel from typing_extensions import override from docling_core.transforms.visualizer.base import BaseVisualizer from docling_core.types.doc import ( BoundingRectangle, ContentLayer, CoordOrigin, DocItem, DocItemLabel, DoclingDocument, TextCell, ) class _TLBoundingRectangle(BoundingRectangle): coord_origin: Literal[CoordOrigin.TOPLEFT] = CoordOrigin.TOPLEFT class _TLTextCell(TextCell): rect: _TLBoundingRectangle class _TLCluster(BaseModel): id: int label: DocItemLabel brec: _TLBoundingRectangle confidence: float = 1.0 cells: list[_TLTextCell] = [] children: list["_TLCluster"] = [] # Add child cluster support class LayoutVisualizer(BaseVisualizer): """Layout visualizer.""" class Params(BaseModel): """Layout visualization parameters.""" show_label: bool = True content_layers: set[ContentLayer] = set(ContentLayer) base_visualizer: BaseVisualizer | None = None params: Params = Params() def _draw_clusters(self, image: Image, clusters: list[_TLCluster], scale_x: float, scale_y: float) -> None: """Draw clusters on an image.""" # Create transparent overlay for proper alpha compositing overlay = PILImage.new("RGBA", image.size, (255, 255, 255, 0)) overlay_draw = ImageDraw.Draw(overlay) main_draw = ImageDraw.Draw(image) # Create a smaller font for the labels font: ImageFont.ImageFont | FreeTypeFont try: font = ImageFont.truetype("arial.ttf", 12) except OSError: # Fallback to default font if arial is not available font = ImageFont.load_default() for c_tl in clusters: all_clusters = [c_tl, *c_tl.children] for c in all_clusters: # Draw cells first (underneath) cell_color = (0, 0, 0, 40) # Transparent black for cells for tc in c.cells: cx0, cy0, cx1, cy1 = tc.rect.to_bounding_box().as_tuple() cx0 *= scale_x cx1 *= scale_x cy0 *= scale_y cy1 *= scale_y # Draw fill on overlay (for transparency) overlay_draw.rectangle( [(cx0, cy0), (cx1, cy1)], outline=None, fill=cell_color, ) # Draw cluster rectangle x0, y0, x1, y1 = c.brec.to_bounding_box().as_tuple() x0 *= scale_x x1 *= scale_x y0 *= scale_y y1 *= scale_y cluster_fill_color = (*list(DocItemLabel.get_color(c.label)), 70) cluster_outline_color = ( *list(DocItemLabel.get_color(c.label)), 255, ) # Draw fill on overlay (for transparency) overlay_draw.rectangle( [(x0, y0), (x1, y1)], outline=None, fill=cluster_fill_color, ) # Draw outline on main image main_draw.rectangle( [(x0, y0), (x1, y1)], outline=cluster_outline_color, fill=None, ) if self.params.show_label: # Add label name and confidence label_text = f"{c.label.name} ({c.confidence:.2f})" # Create semi-transparent background for text text_bbox = overlay_draw.textbbox((x0, y0), label_text, font=font) text_bg_padding = 2 # Draw label background on overlay (for transparency) overlay_draw.rectangle( [ ( text_bbox[0] - text_bg_padding, text_bbox[1] - text_bg_padding, ), ( text_bbox[2] + text_bg_padding, text_bbox[3] + text_bg_padding, ), ], fill=(255, 255, 255, 180), # Semi-transparent white ) # Draw text on overlay overlay_draw.text( (x0, y0), label_text, fill=(0, 0, 0, 255), # Solid black font=font, ) # Alpha composite the overlay onto the image composited = PILImage.alpha_composite(image.convert("RGBA"), overlay) image.paste(composited.convert(image.mode)) def _draw_doc_layout( self, doc: DoclingDocument, images: dict[int | None, Image] | None = None, included_content_layers: set[ContentLayer] | None = None, ): """Draw the document clusters and optionally the reading order.""" clusters = [] my_images: dict[int | None, Image] = {} if images is not None: my_images = images if included_content_layers is None: included_content_layers = set(ContentLayer) # Initialise `my_images` beforehand: sometimes, you have the # page-images but no DocItems! for page_nr, page in doc.pages.items(): page_image = doc.pages[page_nr].image if page_image is None or (pil_img := page_image.pil_image) is None: raise RuntimeError("Cannot visualize document without images") elif page_nr not in my_images: image = pil_img.copy() my_images[page_nr] = image prev_image = None prev_page_nr = None for idx, (elem, _) in enumerate( doc.iterate_items(included_content_layers=included_content_layers, traverse_pictures=True) ): if not isinstance(elem, DocItem): continue if len(elem.prov) == 0: continue # Skip elements without provenances for prov in elem.prov: page_nr = prov.page_no if page_nr in my_images: image = my_images[page_nr] else: raise RuntimeError(f"Cannot visualize page-image for {page_nr}") if prev_page_nr is None or page_nr != prev_page_nr: # changing page # dump previous drawing if prev_page_nr is not None and prev_image and clusters: self._draw_clusters( image=prev_image, clusters=clusters, scale_x=prev_image.width / doc.pages[prev_page_nr].size.width, scale_y=prev_image.height / doc.pages[prev_page_nr].size.height, ) clusters = [] tlo_bbox = prov.bbox.to_top_left_origin(page_height=doc.pages[prov.page_no].size.height) cluster = _TLCluster( id=idx, label=elem.label, brec=_TLBoundingRectangle.from_bounding_box(bbox=tlo_bbox), cells=[], ) clusters.append(cluster) prev_page_nr = page_nr prev_image = image # complete last drawing if prev_page_nr is not None and prev_image and clusters: self._draw_clusters( image=prev_image, clusters=clusters, scale_x=prev_image.width / doc.pages[prev_page_nr].size.width, scale_y=prev_image.height / doc.pages[prev_page_nr].size.height, ) return my_images @override def get_visualization( self, *, doc: DoclingDocument, **kwargs, ) -> dict[int | None, Image]: """Get visualization of the document as images by page.""" base_images = self.base_visualizer.get_visualization(doc=doc, **kwargs) if self.base_visualizer else None return self._draw_doc_layout( doc=doc, images=base_images, )