mirror of
https://github.com/Unstructured-IO/unstructured.git
synced 2025-07-04 23:52:23 +00:00
f673ea4c40
This PR was initially created to close GitHub Issue #1604 (Synchronizing the default layout model), but since it was already resolved in PR [#1607](https://github.com/Unstructured-IO/unstructured/pull/1607), this PR now only adds the visualization script used to investigate the issue. ### Summary - add python script to annotate elements PDF: [references.pdf](https://github.com/Unstructured-IO/unstructured/files/12778270/references.pdf) ### Evaluation ``` PYTHONPATH=. python examples/layout-analysis/visualization.py references.pdf hi_res ```
83 lines
2.5 KiB
Python
83 lines
2.5 KiB
Python
import os
|
|
import pathlib
|
|
import sys
|
|
|
|
import pdf2image
|
|
from unstructured_inference.inference.elements import Rectangle
|
|
from unstructured_inference.visualize import draw_bbox
|
|
|
|
from unstructured.documents.elements import PageBreak
|
|
from unstructured.partition.pdf import partition_pdf
|
|
|
|
CUR_DIR = pathlib.Path(__file__).parent.resolve()
|
|
|
|
|
|
def extract_element_coordinates(elements):
|
|
elements_coordinates = []
|
|
page_elements_coordinates = []
|
|
|
|
for el in elements:
|
|
if isinstance(el, PageBreak):
|
|
if page_elements_coordinates:
|
|
elements_coordinates.append(page_elements_coordinates)
|
|
page_elements_coordinates = []
|
|
else:
|
|
page_elements_coordinates.append(el.metadata.coordinates)
|
|
|
|
if page_elements_coordinates:
|
|
elements_coordinates.append(page_elements_coordinates)
|
|
|
|
return elements_coordinates
|
|
|
|
|
|
def run_partition_pdf(f_path, strategy, images, output_dir):
|
|
|
|
elements = partition_pdf(
|
|
f_path,
|
|
strategy=strategy,
|
|
include_page_breaks=True,
|
|
)
|
|
|
|
elements_coordinates = extract_element_coordinates(elements)
|
|
assert len(images) == len(elements_coordinates)
|
|
|
|
for idx, (img, coords_per_page) in enumerate(zip(images, elements_coordinates)):
|
|
for coordinate in coords_per_page:
|
|
points = coordinate.points
|
|
x1, y1 = points[0]
|
|
x2, y2 = points[2]
|
|
rect = Rectangle(x1, y1, x2, y2)
|
|
img = draw_bbox(img, rect, color="red")
|
|
|
|
output_image_path = os.path.join(output_dir, f"{strategy}-{idx + 1}.jpg")
|
|
print(f"output_image_path: {output_image_path}")
|
|
|
|
img.save(output_image_path)
|
|
|
|
|
|
def run(f_path, strategy):
|
|
f_basename = os.path.splitext(os.path.basename(f_path))[0]
|
|
output_dir_path = os.path.join(output_basedir_path, f_basename)
|
|
os.makedirs(output_dir_path, exist_ok=True)
|
|
|
|
images = pdf2image.convert_from_path(f_path)
|
|
run_partition_pdf(f_path, strategy, images, output_dir_path)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if len(sys.argv) < 3:
|
|
print(
|
|
"Please provide the path to the file name as the first argument and the strategy as the "
|
|
"second argument.",
|
|
)
|
|
sys.exit(1)
|
|
|
|
if sys.argv[2] not in ["auto", "hi_res", "ocr_only", "fast"]:
|
|
print("Invalid strategy")
|
|
sys.exit(1)
|
|
|
|
output_basedir_path = os.path.join(CUR_DIR, "output")
|
|
os.makedirs(output_basedir_path, exist_ok=True)
|
|
|
|
run(f_path=sys.argv[1], strategy=sys.argv[2])
|