fix: annotate image is fixed after refactor (#262)

Currently the code for annotate() don't add details when
add_details=True, for example:

```
from unstructured_inference.inference.layout import DocumentLayout
from unstructured_inference.models.base import get_model

file = "example-docs/layout-parser-paper-fast.pdf"
model = get_model("yolox_quantized")
doc = DocumentLayout.from_file(file,model)

for i,page in enumerate(doc.pages):
    page.annotate(add_details=True).save(f"{i}.png")
```

After this PR is merged, the resulting image should show origin for all
elements using the same code.

Author: benjats07

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 * fix: Reduce Chipper memory consumption on x86_64 cpus
 * fix: Skips ordering elements coming from Chipper
+* fix: After refactoring to introduce Chipper, annotate() weren't able to show text with extra info from elements, this is fixed now.
 
 ## 0.7.9
 

test_unstructured_inference/test_visualization.py

@@ -4,16 +4,16 @@
 import pytest
 from PIL import Image
 
-from unstructured_inference.inference.elements import Rectangle
+from unstructured_inference.inference.elements import TextRegion
 from unstructured_inference.visualize import draw_bbox, show_plot
 
 
 def test_draw_bbox():
     test_image_arr = np.ones((100, 100, 3), dtype="uint8")
     image = Image.fromarray(test_image_arr)
     x1, y1, x2, y2 = (1, 10, 7, 11)
-    rect = Rectangle(x1, y1, x2, y2)
-    annotated_image = draw_bbox(image=image, rect=rect)
+    rect = TextRegion.from_coords(x1, y1, x2, y2)
+    annotated_image = draw_bbox(image=image, element=rect, details=False)
     annotated_array = np.array(annotated_image)
     # Make sure the pixels on the edge of the box are red
     for i, expected in zip(range(3), [255, 0, 0]):

unstructured_inference/inference/layout.py

@@ -360,7 +360,7 @@ def annotate(
         if annotation_data is None:
             for el, color in zip(self.elements, colors):
                 if sources is None or el.source in sources:
-                    img = draw_bbox(img, el.bbox, color=color, details=add_details)
+                    img = draw_bbox(img, el, color=color, details=add_details)
         else:
             for attribute, style in annotation_data.items():
                 if hasattr(self, attribute) and getattr(self, attribute):

unstructured_inference/visualize.py

@@ -10,13 +10,13 @@
 from PIL.Image import Image
 from PIL.ImageDraw import ImageDraw
 
-from unstructured_inference.inference.elements import Rectangle
+from unstructured_inference.inference.elements import TextRegion
 
 
 @typing.no_type_check
 def draw_bbox(
     image: Image,
-    rect: Rectangle,
+    element: TextRegion,
     color: str = "red",
     width=1,
     details: bool = False,
@@ -25,17 +25,17 @@ def draw_bbox(
     try:
         img = image.copy()
         draw = ImageDraw(img)
-        topleft, _, bottomright, _ = rect.coordinates
-        c = getattr(rect, "color", color)
+        topleft, _, bottomright, _ = element.bbox.coordinates
+        c = getattr(element, "color", color)
         if details:
-            source = getattr(rect, "source", "Unknown")
-            type = getattr(rect, "type", "")
+            source = getattr(element, "source", "Unknown")
+            type = getattr(element, "type", "")
             kbd = ImageFont.truetype("Keyboard.ttf", 20)
             draw.text(topleft, text=f"{type} {source}", fill=c, font=kbd)
         draw.rectangle((topleft, bottomright), outline=c, width=width)
     except OSError:
         print("Failed to find font file. Skipping details.")
-        img = draw_bbox(image, rect, color, width)
+        img = draw_bbox(image, element, color, width)
     except Exception as e:
         print(f"Failed to draw bounding box: {e}")
     return img