From 2ee38e6cd5c924fa272d661eceb80e53ca9eb42e Mon Sep 17 00:00:00 2001 From: Benjamin Torres Date: Wed, 18 Oct 2023 21:02:27 -0500 Subject: [PATCH] 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. --- CHANGELOG.md | 1 + test_unstructured_inference/test_visualization.py | 6 +++--- unstructured_inference/inference/layout.py | 2 +- unstructured_inference/visualize.py | 14 +++++++------- 4 files changed, 12 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 33966f18..45d58139 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/test_unstructured_inference/test_visualization.py b/test_unstructured_inference/test_visualization.py index 0524451f..6855aca9 100644 --- a/test_unstructured_inference/test_visualization.py +++ b/test_unstructured_inference/test_visualization.py @@ -4,7 +4,7 @@ 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 @@ -12,8 +12,8 @@ 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]): diff --git a/unstructured_inference/inference/layout.py b/unstructured_inference/inference/layout.py index 9a86c2fc..5ab39aab 100644 --- a/unstructured_inference/inference/layout.py +++ b/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): diff --git a/unstructured_inference/visualize.py b/unstructured_inference/visualize.py index 8cba46ac..a01d6aba 100644 --- a/unstructured_inference/visualize.py +++ b/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