Merge pull request #7 from dhdaines/more_boxes

More boxes

Author: dhdaines

README.md

@@ -44,7 +44,27 @@ page = pdf.pages[3]
 pi.show(page)
 ```
 
-You can of course draw boxes around those individual PDF objects, or
+Something quite interesting to do is, if your PDF contains a logical
+structure tree, to look at the bounding boxes of the contents of those
+structure elements (FIXME: This is not a very efficient way to do
+this, and it will be optimized in an upcoming PLAYA):
+
+```python
+pi.box(pdf.structure.find_all(lambda el: el.page is page))
+```
+
+![Structure Elements](./docs/page3-elements.png)
+
+Alternately, if you have annotations (such as links), you can look at
+those too:
+
+```python
+pi.box(page.annotations)
+```
+
+![Annotations](./docs/page2-annotations.png)
+
+You can of course draw boxes around individual PDF objects, or
 one particular sort of object, or filter them with a generator
 expression:
 
@@ -62,12 +82,28 @@ semi-transparent colour, which otherwise works the same way:
 pi.mark(page.images)
 ```
 
+![Annotations](./docs/page298-images.png)
+
 If you wish you can give each type of object a different colour:
 
 ```python
-pi.mark(page, color={"text": "red", "image": "blue", "path": green"})
+pi.mark(page, color={"text": "red", "image": "blue", "path": "green"})
 ```
 
+![Annotations](./docs/page298-colors.png)
+
+You can also add outlines and labels around the highlighting:
+
+```python
+pi.mark(page, outline=True, label=True,
+        color={"text": "red", "image": "blue", "path": "green"})
+```
+
+![Annotations](./docs/page298-outlines.png)
+
+There are even more options!  For now you will need to look at the
+source code, documentation is Coming Soon.
+
 ## Working in the PDF mine
 
 `pdfminer.six` is widely used for text extraction and layout analysis

docs/page2-annotations.png

@@ -3,16 +3,20 @@
 models and/or visualisation.`
 """
 
+import itertools
 import functools
 import subprocess
 import tempfile
 from os import PathLike
 from pathlib import Path
-from typing import Dict, Iterable, Iterator, Union, List, TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List, Union
+
 from PIL import Image, ImageDraw, ImageFont
 from playa.document import Document, PageList
-from playa.page import ContentObject, Page
-
+from playa.page import ContentObject, Page, Annotation
+from playa.structure import Element
+from playa.utils import Rect, get_transformed_bound, get_bound
+from playa import resolve
 
 if TYPE_CHECKING:
     import pypdfium2  # types: ignore
@@ -301,6 +305,109 @@ def show(page: Page, dpi: int = 72) -> Image.Image:
     return next(convert(page, dpi=dpi))
 
 
+LabelFunc = Callable[[ContentObject], str]
+BoxFunc = Callable[[ContentObject], Rect]
+
+
+@functools.singledispatch
+def get_box(obj: Union[Annotation, ContentObject, Element, Rect]) -> Rect:
+    """Default function to get the bounding box for an object."""
+    raise RuntimeError(f"Don't know how to get the box for {obj!r}")
+
+
+@get_box.register(tuple)
+def get_box_rect(obj: Rect) -> Rect:
+    """Get the bounding box of a ContentObject"""
+    return obj
+
+
+@get_box.register(ContentObject)
+def get_box_content(obj: ContentObject) -> Rect:
+    """Get the bounding box of a ContentObject"""
+    return obj.bbox
+
+
+@get_box.register(Annotation)
+def get_box_annotation(obj: Annotation) -> Rect:
+    """Get the bounding box of an Annotation"""
+    return get_transformed_bound(obj.page.ctm, obj.rect)
+
+
+@get_box.register(Element)
+def get_box_element(obj: Element) -> Rect:
+    """Get the bounding box for a structural Element"""
+    # It might just *have* a BBox already
+    page = obj.page
+    if page is None:
+        raise ValueError("Has no page, has no content! No box for you!")
+    if "BBox" in obj.props:
+        return get_transformed_bound(page.ctm, resolve(obj.props["BBox"]))
+    else:
+        return _get_marked_content_box(obj)
+
+
+def _get_marked_content_box(el: Element) -> Rect:
+    """Get the bounding box of an Element's marked content.
+
+    This will be superseded in PLAYA 0.3.x so do not use!
+    """
+    page = el.page
+    if page is None:
+        raise ValueError("Has no page, has no content! No box for you!")
+
+    def get_mcids(k):
+        k = resolve(k)
+        if isinstance(k, int):
+            yield k
+        elif isinstance(k, list):
+            for kk in k:
+                yield from get_mcids(kk)
+        elif isinstance(k, dict):
+            if "K" in k:
+                yield from get_mcids(k["K"])
+
+    mcids = set(get_mcids(el.props["K"]))
+    pts = itertools.chain.from_iterable(
+        ((x0, y0), (x1, y1))
+        for x0, y0, x1, y1 in (obj.bbox for obj in page if obj.mcid in mcids)
+    )
+    return get_bound(pts)
+
+
+@functools.singledispatch
+def get_label(obj: Union[Annotation, ContentObject, Element, Rect]) -> str:
+    """Default function to get the label text for an object."""
+    return str(obj)
+
+
+@get_label.register(ContentObject)
+def get_label_content(obj: ContentObject) -> str:
+    """Get the label text for a ContentObject."""
+    return obj.object_type
+
+
+@get_label.register(Annotation)
+def get_label_annotation(obj: Annotation) -> str:
+    """Get the default label text for an Annotation.
+
+    Note: This is just a default.
+        This is one of many possible options, so you may wish to
+        define your own custom LabelFunc.
+    """
+    return obj.subtype
+
+
+@get_label.register(Element)
+def get_label_element(obj: Element) -> str:
+    """Get the default label text for an Element.
+
+    Note: This is just a default.
+        This is one of many possible options, so you may wish to
+        define your own custom LabelFunc.
+    """
+    return obj.type
+
+
 def box(
     objs: Iterable[ContentObject],
     *,
@@ -310,30 +417,41 @@ def box(
     label_size: int = 9,
     label_margin: int = 1,
     image: Union[Image.Image, None] = None,
+    labelfunc: LabelFunc = get_label,
+    boxfunc: BoxFunc = get_box,
 ) -> Union[Image.Image, None]:
     """Draw boxes around things in a page of a PDF."""
     draw: ImageDraw.ImageDraw
     font = ImageFont.load_default(label_size)
     for obj in objs:
         if image is None:
             image = show(obj.page)
-        left, top, right, bottom = obj.bbox
+        try:
+            left, top, right, bottom = boxfunc(obj)
+        except ValueError:  # it has no content and no box
+            continue
         draw = ImageDraw.ImageDraw(image)
         obj_color = (
             color if isinstance(color, str) else color.get(obj.object_type, "red")
         )
         draw.rectangle((left, top, right, bottom), outline=obj_color)
         if label:
-            text = obj.object_type
+            text = labelfunc(obj)
             tl, tt, tr, tb = font.getbbox(text)
+            label_box = (
+                left,
+                top - tb - label_margin * 2,
+                left + tr + label_margin * 2,
+                top,
+            )
             draw.rectangle(
-                (left, top - tb - label_margin * 2, left + tr + label_margin * 2, top),
+                label_box,
                 outline=obj_color,
                 fill=obj_color,
             )
             draw.text(
                 xy=(left + label_margin, top - label_margin),
-                text=obj.object_type,
+                text=text,
                 font=font,
                 fill="white",
                 anchor="ld",
@@ -350,7 +468,10 @@ def mark(
     label_color: str = "white",
     label_size: int = 9,
     label_margin: int = 1,
+    outline: bool = False,
     image: Union[Image.Image, None] = None,
+    labelfunc: LabelFunc = get_label,
+    boxfunc: BoxFunc = get_box,
 ) -> Union[Image.Image, None]:
     """Highlight things in a page of a PDF."""
     overlay: Image.Image
@@ -363,34 +484,69 @@ def mark(
             image = show(obj.page)
             overlay = Image.new("RGB", image.size)
             mask = Image.new("L", image.size, 255)
-        left, top, right, bottom = obj.bbox
+        try:
+            left, top, right, bottom = boxfunc(obj)
+        except ValueError:  # it has no content and no box
+            continue
         draw = ImageDraw.ImageDraw(overlay)
         obj_color = (
             color if isinstance(color, str) else color.get(obj.object_type, "red")
         )
         draw.rectangle((left, top, right, bottom), fill=obj_color)
+        mask_draw = ImageDraw.ImageDraw(mask)
+        mask_draw.rectangle((left, top, right, bottom), fill=alpha)
+        if outline:
+            draw.rectangle((left, top, right, bottom), outline="black")
+            mask_draw.rectangle((left, top, right, bottom), outline=0)
         if label:
-            text = obj.object_type
+            text = labelfunc(obj)
             tl, tt, tr, tb = font.getbbox(text)
+            label_box = (
+                left,
+                top - tb - label_margin * 2,
+                left + tr + label_margin * 2,
+                top,
+            )
             draw.rectangle(
-                (left, top - tb - label_margin * 2, left + tr + label_margin * 2, top),
+                label_box,
                 outline=obj_color,
                 fill=obj_color,
             )
-            draw.text(
-                xy=(left + label_margin, top - label_margin),
-                text=obj.object_type,
-                font=font,
-                fill="white",
-                anchor="ld",
-            )
-        draw = ImageDraw.ImageDraw(mask)
-        draw.rectangle((left, top, right, bottom), fill=alpha)
-        if label:
-            draw.rectangle(
-                (left, top - tb - label_margin * 2, left + tr + label_margin * 2, top),
+            mask_draw.rectangle(
+                label_box,
                 fill=alpha,
             )
+            if outline:
+                draw.rectangle(
+                    label_box,
+                    outline="black",
+                )
+                mask_draw.rectangle(
+                    label_box,
+                    outline=0,
+                )
+                draw.text(
+                    xy=(left + label_margin, top - label_margin),
+                    text=text,
+                    font=font,
+                    fill="black",
+                    anchor="ld",
+                )
+                mask_draw.text(
+                    xy=(left + label_margin, top - label_margin),
+                    text=text,
+                    font=font,
+                    fill=0,
+                    anchor="ld",
+                )
+            else:
+                draw.text(
+                    xy=(left + label_margin, top - label_margin),
+                    text=text,
+                    font=font,
+                    fill="white",
+                    anchor="ld",
+                )
     if image is None:
         return None
     return Image.composite(image, overlay, mask)