Skip to content

Commit

Permalink
fix: annotate image is fixed after refactor (#262)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
benjats07 authored Oct 19, 2023
1 parent 7b99828 commit 2ee38e6
Show file tree
Hide file tree
Showing 4 changed files with 12 additions and 11 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
6 changes: 3 additions & 3 deletions test_unstructured_inference/test_visualization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]):
Expand Down
2 changes: 1 addition & 1 deletion unstructured_inference/inference/layout.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
14 changes: 7 additions & 7 deletions unstructured_inference/visualize.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down

0 comments on commit 2ee38e6

Please sign in to comment.